dfir_rs/scheduled/handoff/
handoff_list.rs

1//! Module for variadic handoff port lists, [`PortList`].
2
3use ref_cast::RefCast;
4use sealed::sealed;
5use variadics::{Variadic, variadic_trait};
6
7use super::Handoff;
8use crate::scheduled::graph::HandoffData;
9use crate::scheduled::port::{Polarity, Port, PortCtx};
10use crate::scheduled::{HandoffId, HandoffTag, SubgraphId};
11use crate::util::slot_vec::SlotVec;
12
13/// Sealed trait for variadic lists of ports.
14///
15/// See the [`variadics`] crate for the strategy we use to implement variadics in Rust.
16#[sealed]
17pub trait PortList<S>: Variadic
18where
19    S: Polarity,
20{
21    /// Iteratively/recursively set the graph metadata for each port in this list.
22    ///
23    /// Specifically sets:
24    /// - `HandoffData::preds` and `HandoffData::succs` in the `handoffs` slice for the
25    ///   handoffs in this [`PortList`] (using `pred` and/or `succ`).
26    /// - `out_handoff_ids` will be extended with all the handoff IDs in this [`PortList`].
27    ///
28    /// `handoffs_are_preds`:
29    /// - `true`: Handoffs are predecessors (inputs) to subgraph `sg_id`.
30    /// - `false`: Handoffs are successors (outputs) from subgraph `sg_id`.
31    fn set_graph_meta(
32        &self,
33        handoffs: &mut SlotVec<HandoffTag, HandoffData>,
34        out_handoff_ids: &mut Vec<HandoffId>,
35        sg_id: SubgraphId,
36        handoffs_are_preds: bool,
37    );
38
39    /// The [`Variadic`] return type of [`Self::make_ctx`].
40    type Ctx<'a>: Variadic;
41    /// Iteratively/recursively construct a `Ctx` variadic list.
42    ///
43    /// (Note that unlike [`Self::set_graph_meta`], this does not mess with pred/succ handoffs for
44    /// teeing).
45    ///
46    /// # Safety
47    /// The handoffs in this port list (`self`) must come from the `handoffs` [`SlotVec`].
48    /// This ensure the types will match.
49    ///
50    /// Use [`Self::assert_is_from`] to check this.
51    unsafe fn make_ctx<'a>(&self, handoffs: &'a SlotVec<HandoffTag, HandoffData>) -> Self::Ctx<'a>;
52
53    /// Asserts that `self` is a valid port list from `handoffs`. Panics if not.
54    fn assert_is_from(&self, handoffs: &SlotVec<HandoffTag, HandoffData>);
55}
56#[sealed]
57impl<S, Rest, H> PortList<S> for (Port<S, H>, Rest)
58where
59    S: Polarity,
60    H: Handoff,
61    Rest: PortList<S>,
62{
63    fn set_graph_meta(
64        &self,
65        handoffs: &mut SlotVec<HandoffTag, HandoffData>,
66        out_handoff_ids: &mut Vec<HandoffId>,
67        sg_id: SubgraphId,
68        handoffs_are_preds: bool,
69    ) {
70        let (this, rest) = self;
71        let this_handoff = &mut handoffs[this.handoff_id];
72
73        // Set subgraph's info (`out_handoff_ids`) about neighbor handoffs.
74        // Use the "representative" handoff (pred or succ) for teeing handoffs, for the subgraph metadata.
75        // For regular Vec handoffs, `pred_handoffs` and `succ_handoffs` will just be the handoff itself.
76        out_handoff_ids.extend(if handoffs_are_preds {
77            this_handoff.pred_handoffs.iter().copied()
78        } else {
79            this_handoff.succ_handoffs.iter().copied()
80        });
81
82        // Set handoff's info (`preds`/`succs`) about neighbor subgraph (`sg_id`).
83        if handoffs_are_preds {
84            for succ_hoff in this_handoff.succ_handoffs.clone() {
85                handoffs[succ_hoff].succs.push(sg_id);
86            }
87        } else {
88            for pred_hoff in this_handoff.pred_handoffs.clone() {
89                handoffs[pred_hoff].preds.push(sg_id);
90            }
91        }
92        rest.set_graph_meta(handoffs, out_handoff_ids, sg_id, handoffs_are_preds);
93    }
94
95    type Ctx<'a> = (&'a PortCtx<S, H>, Rest::Ctx<'a>);
96    unsafe fn make_ctx<'a>(&self, handoffs: &'a SlotVec<HandoffTag, HandoffData>) -> Self::Ctx<'a> {
97        let (this, rest) = self;
98        let hoff_any = handoffs.get(this.handoff_id).unwrap().handoff.any_ref();
99        debug_assert!(hoff_any.is::<H>());
100
101        let handoff = unsafe {
102            // SAFETY: Caller must ensure `self` is from `handoffs`.
103            // TODO(shadaj): replace with `downcast_ref_unchecked` when it's stabilized
104            &*(hoff_any as *const dyn std::any::Any as *const H)
105        };
106
107        let ctx = RefCast::ref_cast(handoff);
108        let ctx_rest = unsafe {
109            // SAFETY: Same invariants hold, as we recurse through the list.
110            rest.make_ctx(handoffs)
111        };
112        (ctx, ctx_rest)
113    }
114
115    fn assert_is_from(&self, handoffs: &SlotVec<HandoffTag, HandoffData>) {
116        let (this, rest) = self;
117        let Some(hoff_data) = handoffs.get(this.handoff_id) else {
118            panic!("Handoff ID {} not found in `handoffs`.", this.handoff_id);
119        };
120        let hoff_any = hoff_data.handoff.any_ref();
121        assert!(
122            hoff_any.is::<H>(),
123            "Handoff ID {} is not of type {} in `handoffs`.",
124            this.handoff_id,
125            std::any::type_name::<H>(),
126        );
127        rest.assert_is_from(handoffs);
128    }
129}
130#[sealed]
131impl<S> PortList<S> for ()
132where
133    S: Polarity,
134{
135    fn set_graph_meta(
136        &self,
137        _handoffs: &mut SlotVec<HandoffTag, HandoffData>,
138        _out_handoff_ids: &mut Vec<HandoffId>,
139        _sg_id: SubgraphId,
140        _handoffs_are_preds: bool,
141    ) {
142    }
143
144    type Ctx<'a> = ();
145    unsafe fn make_ctx<'a>(
146        &self,
147        _handoffs: &'a SlotVec<HandoffTag, HandoffData>,
148    ) -> Self::Ctx<'a> {
149    }
150
151    fn assert_is_from(&self, _handoffs: &SlotVec<HandoffTag, HandoffData>) {}
152}
153
154/// Trait for splitting a list of ports into two.
155#[sealed]
156pub trait PortListSplit<S, A>: PortList<S>
157where
158    S: Polarity,
159    A: PortList<S>,
160{
161    /// The suffix, second half of the split.
162    type Suffix: PortList<S>;
163
164    /// Split the port list, returning the prefix and [`Self::Suffix`] as the two halves.
165    fn split_ctx(ctx: Self::Ctx<'_>) -> (A::Ctx<'_>, <Self::Suffix as PortList<S>>::Ctx<'_>);
166}
167#[sealed]
168impl<S, H, T, U> PortListSplit<S, (Port<S, H>, U)> for (Port<S, H>, T)
169where
170    S: Polarity,
171    H: Handoff,
172    T: PortListSplit<S, U>,
173    U: PortList<S>,
174{
175    type Suffix = T::Suffix;
176
177    fn split_ctx(
178        ctx: Self::Ctx<'_>,
179    ) -> (
180        <(Port<S, H>, U) as PortList<S>>::Ctx<'_>,
181        <Self::Suffix as PortList<S>>::Ctx<'_>,
182    ) {
183        let (x, t) = ctx;
184        let (u, v) = T::split_ctx(t);
185        ((x, u), v)
186    }
187}
188#[sealed]
189impl<S, T> PortListSplit<S, ()> for T
190where
191    S: Polarity,
192    T: PortList<S>,
193{
194    type Suffix = T;
195
196    fn split_ctx(ctx: Self::Ctx<'_>) -> ((), T::Ctx<'_>) {
197        ((), ctx)
198    }
199}
200
201variadic_trait! {
202    /// A variadic list of Handoff types, represented using a lisp-style tuple structure.
203    ///
204    /// This trait is sealed and not meant to be implemented or used directly. Instead tuple lists (which already implement this trait) should be used, for example:
205    /// ```ignore
206    /// type MyHandoffList = (VecHandoff<usize>, (VecHandoff<String>, (TeeingHandoff<u32>, ())));
207    /// ```
208    /// The [`var_expr!`](crate::var) macro simplifies usage of this kind:
209    /// ```ignore
210    /// type MyHandoffList = var_expr!(VecHandoff<usize>, VecHandoff<String>, TeeingHandoff<u32>);
211    /// ```
212    #[sealed]
213    pub variadic<T> HandoffList where T: 'static + Handoff {}
214}