fluidattacks_blends_domain/stack/resolve/
stitcher.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3
4use crate::stack::resolve::{PartialPath, PartialPathDatabase, PartialSymbolStack};
5use crate::stack::StackNode;
6use crate::NodeId;
7
8/// A walk the stitcher built, and the file of the node it ends at.
9#[derive(Clone, PartialEq, Eq, Debug)]
10#[non_exhaustive]
11pub struct StitchedPath<'a> {
12    /// The file of the node the walk ends at.
13    pub file: &'a str,
14    /// The walk.
15    pub path: PartialPath<'a>,
16}
17
18/// The stored paths a stitcher tries from one reference when no other limit is given.
19pub const DEFAULT_MAX_STITCH_WORK: usize = 1_000;
20
21type State<'a> = (NodeId, &'a str, NodeId, PartialSymbolStack<'a>);
22
23/// Extends walks from the references of one file, phase by phase, with the stored partial paths.
24#[derive(Debug)]
25pub struct ForwardStitcher<'d, 'a> {
26    database: &'d PartialPathDatabase<'a>,
27    queue: Vec<StitchedPath<'a>>,
28    seen: BTreeSet<State<'a>>,
29    max_work: usize,
30    work: BTreeMap<NodeId, usize>,
31    exhausted: BTreeSet<NodeId>,
32}
33
34impl<'d, 'a> ForwardStitcher<'d, 'a> {
35    /// A stitcher whose walks start at those `references` of `file` that are references, and
36    /// that tries at most `max_work` stored paths from each of them.
37    pub fn from_references(
38        database: &'d PartialPathDatabase<'a>,
39        file: &'a str,
40        references: impl IntoIterator<Item = NodeId>,
41        max_work: usize,
42    ) -> Self {
43        let mut seen = BTreeSet::new();
44        let queue = database.graph(file).map_or_else(Vec::new, |graph| {
45            references
46                .into_iter()
47                .filter(|node| graph.nodes.get(node).is_some_and(StackNode::is_reference))
48                .flat_map(|node| PartialPath::seed(graph, node))
49                .map(|path| StitchedPath { file, path })
50                .filter(|seed| seen.insert(state(seed)))
51                .collect()
52        });
53        Self {
54            database,
55            queue,
56            seen,
57            max_work,
58            work: BTreeMap::new(),
59            exhausted: BTreeSet::new(),
60        }
61    }
62
63    /// Extends every pending walk by one stored path, and returns the walks that now reach a
64    /// definition with no symbol left to find. A walk that reaches an import also goes on to
65    /// what the import names. A walk that needs symbols on the stack before the reference is
66    /// dropped, since a reference starts from an empty stack: at a root, no stored path that needs
67    /// symbols below the reference is tried. A reference that runs out of work with stored paths
68    /// left to try stops there and is reported by [`Self::exhausted`].
69    pub fn process_next_phase(&mut self) -> Vec<StitchedPath<'a>> {
70        let pending = core::mem::take(&mut self.queue);
71        let extensions: Vec<_> = pending.iter().flat_map(|walk| self.extend(walk)).collect();
72        let mut complete = Vec::new();
73        for walk in extensions {
74            if !self.seen.insert(state(&walk)) {
75                continue;
76            }
77            let graph = self.database.graph(walk.file);
78            let end = graph.and_then(|graph| graph.nodes.get(&walk.path.end));
79            let answers =
80                walk.path.post.symbols.is_empty() && end.is_some_and(StackNode::is_definition);
81            if answers {
82                complete.push(walk.clone());
83            }
84            if !answers || end.is_some_and(|node| node.import_binding().is_some()) {
85                self.queue.push(walk);
86            }
87        }
88        complete
89    }
90
91    /// Whether no walk is left to extend.
92    #[must_use]
93    pub fn is_complete(&self) -> bool {
94        self.queue.is_empty()
95    }
96
97    /// The references that stopped with stored paths left to try, in id order. A reference
98    /// listed here may still have answers that were never found.
99    pub fn exhausted(&self) -> impl Iterator<Item = NodeId> + '_ {
100        self.exhausted.iter().copied()
101    }
102
103    /// The stored paths tried so far, over every reference.
104    #[must_use]
105    pub fn work(&self) -> usize {
106        self.work.values().sum()
107    }
108
109    fn extend(&mut self, walk: &StitchedPath<'a>) -> Vec<StitchedPath<'a>> {
110        let Some(graph) = self.database.graph(walk.file) else {
111            return Vec::new();
112        };
113        let candidates: Vec<_> = if matches!(graph.nodes.get(&walk.path.end), Some(StackNode::Root))
114        {
115            let below_reference = walk.path.post.variable == walk.path.pre.variable;
116            let left = PartialSymbolStack {
117                symbols: walk.path.post.symbols.clone(),
118                variable: walk.path.post.variable.filter(|_| !below_reference),
119            };
120            let roots = self.database.starting_at_root(&left);
121            roots
122                .into_iter()
123                .map(|(id, path)| (id.file, path))
124                .collect()
125        } else {
126            let here = self.database.starting_at(walk.file, walk.path.end);
127            here.map(|(_, path)| (walk.file, path)).collect()
128        };
129        let mut extended = Vec::new();
130        let spent = self.work.entry(walk.path.start).or_default();
131        for (file, candidate) in candidates {
132            if *spent >= self.max_work {
133                self.exhausted.insert(walk.path.start);
134                break;
135            }
136            *spent = spent.saturating_add(1);
137            let at_this_end = PartialPath {
138                start: walk.path.end,
139                ..candidate.clone()
140            };
141            match walk.path.concatenate(graph, &at_this_end) {
142                Ok(path) if path.pre.symbols.is_empty() => {
143                    extended.push(StitchedPath { file, path });
144                }
145                _ => {}
146            }
147        }
148        extended
149    }
150}
151
152fn state<'a>(walk: &StitchedPath<'a>) -> State<'a> {
153    (
154        walk.path.start,
155        walk.file,
156        walk.path.end,
157        walk.path.post.clone(),
158    )
159}
160
161#[cfg(test)]
162mod tests {
163    use super::{ForwardStitcher, DEFAULT_MAX_STITCH_WORK};
164    use crate::stack::resolve::{PartialPathDatabase, DEFAULT_MAX_WORK};
165    use crate::stack::{
166        Definition, ImportBinding, StackGraph, StackNode, StackScope, SymbolPop, SymbolPush,
167    };
168    use crate::NodeId;
169    use alloc::borrow::ToOwned;
170    use alloc::vec::Vec;
171
172    const ROOT: NodeId = NodeId(1);
173    const REFERENCE: NodeId = NodeId(2);
174    const SCOPE: NodeId = NodeId(3);
175    const DEFINITION: NodeId = NodeId(4);
176    const RELAY: NodeId = NodeId(5);
177    const MODULE: NodeId = NodeId(6);
178    const MODULE_DEFINITION: NodeId = NodeId(7);
179    const IMPORT: NodeId = NodeId(8);
180    const ORIGIN: NodeId = NodeId(9);
181    const MODULE_REFERENCE: NodeId = NodeId(10);
182    const LOCAL_REFERENCE: NodeId = NodeId(11);
183    const INNER: NodeId = NodeId(12);
184    const ALL_PHASES: usize = 16;
185    const LIBRARY: (&str, &[(NodeId, NodeId)]) = ("m.py", &[(ROOT, MODULE), (MODULE, DEFINITION)]);
186
187    type Edges = &'static [(NodeId, NodeId)];
188
189    fn push(symbol: &str, is_reference: bool) -> StackNode {
190        StackNode::PushSymbol(SymbolPush {
191            symbol: symbol.to_owned(),
192            precedence: 0,
193            is_reference,
194        })
195    }
196
197    fn pop(symbol: &str, definition: bool, import: bool) -> StackNode {
198        StackNode::PopSymbol(SymbolPop {
199            symbol: symbol.to_owned(),
200            precedence: 0,
201            definition: definition.then(Definition::default),
202            import: import.then(|| ImportBinding {
203                module_stem: "m".to_owned(),
204                source_symbol: None,
205                relative_level: 0,
206                binds_module: false,
207            }),
208        })
209    }
210
211    fn node(id: NodeId) -> StackNode {
212        match id {
213            ROOT => StackNode::Root,
214            REFERENCE | LOCAL_REFERENCE => push("x", true),
215            MODULE_REFERENCE => push("m", true),
216            ORIGIN => push("x", false),
217            RELAY => push("m", false),
218            DEFINITION => pop("x", true, false),
219            IMPORT => pop("x", true, true),
220            MODULE => pop("m", false, false),
221            MODULE_DEFINITION => pop("m", true, false),
222            INNER => pop("x", false, false),
223            _ => StackNode::Scope(StackScope::default()),
224        }
225    }
226
227    fn graph(edges: Edges) -> StackGraph {
228        let mut graph = StackGraph::new();
229        for (source, sink) in edges {
230            graph.add_node(*source, node(*source));
231            graph.add_node(*sink, node(*sink));
232            graph.add_edge(*source, *sink, 0);
233        }
234        graph
235    }
236
237    #[derive(PartialEq, Debug)]
238    struct Run {
239        answers: Vec<(&'static str, NodeId)>,
240        complete: bool,
241        exhausted: Vec<NodeId>,
242        work: usize,
243    }
244
245    fn stitch(
246        files: &[(&'static str, Edges)],
247        reference: NodeId,
248        phases: usize,
249    ) -> (Vec<(&'static str, NodeId)>, bool) {
250        let run = stitch_within(files, &[reference], phases, DEFAULT_MAX_STITCH_WORK);
251        (run.answers, run.complete)
252    }
253
254    fn stitch_within(
255        files: &[(&'static str, Edges)],
256        references: &[NodeId],
257        phases: usize,
258        max_work: usize,
259    ) -> Run {
260        let graphs: Vec<StackGraph> = files.iter().map(|(_, edges)| graph(edges)).collect();
261        let mut database = PartialPathDatabase::default();
262        for ((file, _), graph) in files.iter().zip(&graphs) {
263            database.add_file(file, graph, DEFAULT_MAX_WORK);
264        }
265        let mut stitcher = ForwardStitcher::from_references(
266            &database,
267            "a.py",
268            references.iter().copied(),
269            max_work,
270        );
271        let answers = (0..phases)
272            .flat_map(|_| stitcher.process_next_phase())
273            .filter_map(|walk| {
274                Some((
275                    files.iter().find(|(file, _)| *file == walk.file)?.0,
276                    walk.path.end,
277                ))
278            })
279            .collect();
280        Run {
281            answers,
282            complete: stitcher.is_complete(),
283            exhausted: stitcher.exhausted().collect(),
284            work: stitcher.work(),
285        }
286    }
287
288    #[test]
289    fn a_reference_reaches_a_definition_of_its_own_file() {
290        let own = ("a.py", &[(REFERENCE, SCOPE), (SCOPE, DEFINITION)][..]);
291        assert_eq!(stitch(&[own], REFERENCE, 1).0, [("a.py", DEFINITION)]);
292    }
293
294    #[test]
295    fn a_node_that_is_not_a_reference_starts_no_walk() {
296        let own = ("a.py", &[(REFERENCE, SCOPE), (SCOPE, DEFINITION)][..]);
297        assert_eq!(stitch(&[own], SCOPE, 0), (Vec::new(), true));
298    }
299
300    #[test]
301    fn a_reference_reaches_a_definition_of_another_file_through_the_root() {
302        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
303        let answers = stitch(&[importer, LIBRARY], REFERENCE, ALL_PHASES).0;
304        assert_eq!(answers, [("m.py", DEFINITION)]);
305    }
306
307    #[test]
308    fn a_reference_that_reaches_an_import_also_reaches_what_it_imports() {
309        let edges = &[
310            (REFERENCE, IMPORT),
311            (IMPORT, ORIGIN),
312            (ORIGIN, RELAY),
313            (RELAY, ROOT),
314        ];
315        let answers = stitch(&[("a.py", edges), LIBRARY], REFERENCE, ALL_PHASES).0;
316        assert_eq!(answers, [("a.py", IMPORT), ("m.py", DEFINITION)]);
317    }
318
319    #[test]
320    fn a_definition_reached_with_symbols_left_to_find_is_walked_past() {
321        let edges = &[
322            (REFERENCE, RELAY),
323            (RELAY, MODULE_DEFINITION),
324            (MODULE_DEFINITION, DEFINITION),
325        ];
326        let answers = stitch(&[("a.py", edges)], REFERENCE, ALL_PHASES).0;
327        assert_eq!(answers, [("a.py", DEFINITION)]);
328    }
329
330    #[test]
331    fn a_walk_that_needs_symbols_the_reference_did_not_push_is_dropped() {
332        let importer = ("a.py", &[(MODULE_REFERENCE, ROOT)][..]);
333        assert_eq!(
334            stitch(&[importer, LIBRARY], MODULE_REFERENCE, 2),
335            (Vec::new(), true)
336        );
337    }
338
339    #[test]
340    fn a_walk_back_to_a_state_already_seen_is_not_extended_again() {
341        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
342        let looping = (
343            "m.py",
344            &[(ROOT, MODULE), (MODULE, RELAY), (RELAY, ROOT)][..],
345        );
346        assert!(stitch(&[importer, looping], REFERENCE, ALL_PHASES).1);
347    }
348
349    #[test]
350    fn a_reference_with_no_work_left_is_reported_as_exhausted() {
351        let own = ("a.py", &[(REFERENCE, SCOPE), (SCOPE, DEFINITION)][..]);
352        let run = stitch_within(&[own], &[REFERENCE], ALL_PHASES, 0);
353        assert_eq!(
354            (run.answers, run.exhausted),
355            (Vec::new(), Vec::from([REFERENCE]))
356        );
357    }
358
359    #[test]
360    fn a_reference_with_work_for_every_stored_path_is_not_exhausted() {
361        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
362        let run = stitch_within(&[importer, LIBRARY], &[REFERENCE], ALL_PHASES, 2);
363        assert_eq!(
364            (run.answers, run.exhausted),
365            (Vec::from([("m.py", DEFINITION)]), Vec::new())
366        );
367    }
368
369    #[test]
370    fn the_work_counts_every_stored_path_tried() {
371        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
372        let run = stitch_within(&[importer, LIBRARY], &[REFERENCE], ALL_PHASES, 2);
373        assert_eq!(run.work, 2);
374    }
375
376    #[test]
377    fn a_reference_that_runs_out_of_work_leaves_the_answers_of_the_others() {
378        let edges = &[
379            (REFERENCE, RELAY),
380            (RELAY, ROOT),
381            (LOCAL_REFERENCE, SCOPE),
382            (SCOPE, DEFINITION),
383        ];
384        let references = [REFERENCE, LOCAL_REFERENCE];
385        let run = stitch_within(&[("a.py", edges), LIBRARY], &references, ALL_PHASES, 1);
386        assert_eq!(run.answers, [("a.py", DEFINITION)]);
387    }
388
389    #[test]
390    fn a_walk_whose_stack_grows_without_end_stops_at_the_work_limit() {
391        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
392        let growing = ("m.py", &[(ROOT, RELAY), (RELAY, ROOT)][..]);
393        let run = stitch_within(&[importer, growing], &[REFERENCE], ALL_PHASES, 8);
394        assert_eq!(
395            (run.complete, run.exhausted),
396            (true, Vec::from([REFERENCE]))
397        );
398    }
399
400    #[test]
401    fn a_walk_at_the_root_tries_no_path_that_needs_symbols_below_the_reference() {
402        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
403        let deeper = (
404            "n.py",
405            &[(ROOT, MODULE), (MODULE, INNER), (INNER, MODULE_DEFINITION)][..],
406        );
407        let run = stitch_within(
408            &[importer, LIBRARY, deeper],
409            &[REFERENCE],
410            ALL_PHASES,
411            DEFAULT_MAX_STITCH_WORK,
412        );
413        assert_eq!(run.work, 2);
414    }
415}
fluidattacks_blends_domain/stack/resolve/stitcher.rs

fluidattacks_blends_domain/stack/resolve/
stitcher.rs
