Skip to main content

fluidattacks_blends_domain/stack/resolve/
database.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::vec::Vec;
3use core::ops::Bound;
4
5use crate::stack::resolve::{minimal_partial_paths, PartialPath, PartialSymbolStack};
6use crate::stack::{StackGraph, StackNode};
7use crate::NodeId;
8
9/// Where a stored partial path lives: its file and its place among the file's paths.
10///
11/// An id lives no longer than the borrow of the database that gave it, so that database does
12/// not change while the id lives. It names a path only in that database.
13#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
14#[non_exhaustive]
15pub struct PartialPathId<'a> {
16    /// The file whose graph the path walks.
17    pub file: &'a str,
18    /// The path's place among the file's paths.
19    pub index: usize,
20}
21
22/// The minimal partial paths of each file, kept with the graph they walk.
23#[derive(Clone, Debug, Default)]
24pub struct PartialPathDatabase<'a> {
25    files: BTreeMap<&'a str, FilePaths<'a>>,
26    roots: BTreeMap<Vec<&'a str>, BTreeSet<PartialPathId<'a>>>,
27}
28
29#[derive(Clone, Debug)]
30struct FilePaths<'a> {
31    graph: &'a StackGraph,
32    paths: Vec<PartialPath<'a>>,
33    exhausted: Vec<NodeId>,
34}
35
36impl<'a> PartialPathDatabase<'a> {
37    /// Stores the minimal partial paths of `graph` as `file`'s, in place of any it had, trying
38    /// at most `max_work` edges from each endpoint.
39    pub fn add_file(&mut self, file: &'a str, graph: &'a StackGraph, max_work: usize) {
40        self.drop_file(file);
41        let found = minimal_partial_paths(graph, max_work);
42        let mut paths = found.paths;
43        paths.sort_by_key(|path| path.start);
44        let stored = FilePaths {
45            graph,
46            paths,
47            exhausted: found.exhausted,
48        };
49        for (index, needed) in stored.root_paths() {
50            self.roots
51                .entry(needed)
52                .or_default()
53                .insert(PartialPathId { file, index });
54        }
55        self.files.insert(file, stored);
56    }
57
58    /// Forgets `file` and its paths.
59    pub fn drop_file(&mut self, file: &str) {
60        let Some((file, stored)) = self.files.remove_entry(file) else {
61            return;
62        };
63        for (index, needed) in stored.root_paths() {
64            let Some(ids) = self.roots.get_mut(&needed) else {
65                continue;
66            };
67            ids.remove(&PartialPathId { file, index });
68            if ids.is_empty() {
69                self.roots.remove(&needed);
70            }
71        }
72    }
73
74    /// The graph `file`'s paths walk.
75    #[must_use]
76    pub fn graph(&self, file: &str) -> Option<&'a StackGraph> {
77        self.files.get(file).map(|stored| stored.graph)
78    }
79
80    /// The path stored under `id`, an id this database gave.
81    #[must_use]
82    pub fn path(&self, id: PartialPathId<'_>) -> Option<&PartialPath<'a>> {
83        self.files.get(id.file)?.paths.get(id.index)
84    }
85
86    /// The paths of `file` that start at `node`.
87    pub fn starting_at(
88        &self,
89        file: &str,
90        node: NodeId,
91    ) -> impl Iterator<Item = (PartialPathId<'_>, &PartialPath<'a>)> + '_ {
92        self.files
93            .get_key_value(file)
94            .into_iter()
95            .flat_map(move |(&file, stored)| {
96                let first = stored.paths.partition_point(|path| path.start < node);
97                stored
98                    .paths
99                    .iter()
100                    .enumerate()
101                    .skip(first)
102                    .take_while(move |(_, path)| path.start == node)
103                    .map(move |(index, path)| (PartialPathId { file, index }, path))
104            })
105    }
106
107    /// The paths of any file that start at its root and whose needed stack may follow `left`,
108    /// a stack left at a root: one is a prefix of the other, and the longer one is only
109    /// reachable through the variable of the shorter.
110    #[must_use]
111    pub fn starting_at_root(
112        &self,
113        left: &PartialSymbolStack<'a>,
114    ) -> Vec<(PartialPathId<'a>, &PartialPath<'a>)> {
115        let symbols: Vec<&'a str> = left.symbols.iter().copied().collect();
116        let equal = self.roots.get(symbols.as_slice()).into_iter().flatten();
117        let shorter = (0..symbols.len())
118            .rev()
119            .filter_map(|len| symbols.get(..len))
120            .filter_map(|prefix| self.roots.get(prefix))
121            .flatten();
122        let longer = self
123            .roots
124            .range::<[&'a str], _>((Bound::Excluded(symbols.as_slice()), Bound::Unbounded))
125            .take_while(|(needed, _)| left.variable.is_some() && needed.starts_with(&symbols))
126            .flat_map(|(_, ids)| ids);
127        let with_rest = |id: &PartialPathId<'a>| {
128            self.path(*id)
129                .is_some_and(|path| path.pre.variable.is_some())
130        };
131        equal
132            .chain(shorter.filter(|id| with_rest(id)))
133            .chain(longer)
134            .filter_map(|id| Some((*id, self.path(*id)?)))
135            .collect()
136    }
137
138    /// The endpoints of `file` whose search stopped with edges left to try.
139    #[must_use]
140    pub fn exhausted(&self, file: &str) -> &[NodeId] {
141        self.files
142            .get(file)
143            .map_or(&[], |stored| stored.exhausted.as_slice())
144    }
145}
146
147impl<'a> FilePaths<'a> {
148    fn root_paths(&self) -> impl Iterator<Item = (usize, Vec<&'a str>)> + '_ {
149        self.paths
150            .iter()
151            .enumerate()
152            .filter(|(_, path)| matches!(self.graph.nodes.get(&path.start), Some(StackNode::Root)))
153            .map(|(index, path)| (index, path.pre.symbols.iter().copied().collect()))
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::{PartialPathDatabase, PartialPathId};
160    use crate::stack::resolve::{PartialSymbolStack, SymbolStackVariable, DEFAULT_MAX_WORK};
161    use crate::stack::{Definition, StackGraph, StackNode, StackScope, SymbolPop, SymbolPush};
162    use crate::NodeId;
163    use alloc::borrow::ToOwned;
164    use alloc::vec::Vec;
165
166    const ROOT: NodeId = NodeId(1);
167    const REFERENCE: NodeId = NodeId(2);
168    const SCOPE: NodeId = NodeId(3);
169    const DEFINITION: NodeId = NodeId(4);
170    const RELAY: NodeId = NodeId(5);
171    const MODULE: NodeId = NodeId(6);
172
173    fn push(symbol: &str, is_reference: bool) -> StackNode {
174        StackNode::PushSymbol(SymbolPush {
175            symbol: symbol.to_owned(),
176            precedence: 0,
177            is_reference,
178        })
179    }
180
181    fn pop(symbol: &str, definition: Option<Definition>) -> StackNode {
182        StackNode::PopSymbol(SymbolPop {
183            symbol: symbol.to_owned(),
184            precedence: 0,
185            definition,
186            import: None,
187        })
188    }
189
190    fn node(id: NodeId) -> StackNode {
191        match id {
192            ROOT => StackNode::Root,
193            REFERENCE => push("x", true),
194            RELAY => push("m", false),
195            DEFINITION => pop("x", Some(Definition::default())),
196            MODULE => pop("m", None),
197            _ => StackNode::Scope(StackScope::default()),
198        }
199    }
200
201    fn graph(edges: &[(NodeId, NodeId)]) -> StackGraph {
202        let mut graph = StackGraph::new();
203        for (source, sink) in edges {
204            graph.add_node(*source, node(*source));
205            graph.add_node(*sink, node(*sink));
206            graph.add_edge(*source, *sink, 0);
207        }
208        graph
209    }
210
211    fn ends(database: &PartialPathDatabase<'_>, file: &str, start: NodeId) -> Vec<NodeId> {
212        database
213            .starting_at(file, start)
214            .map(|(_, path)| path.end)
215            .collect()
216    }
217
218    fn reached<'a>(
219        database: &PartialPathDatabase<'a>,
220        symbols: &[&'a str],
221        with_variable: bool,
222    ) -> Vec<(&'a str, NodeId)> {
223        let variable = with_variable.then_some(SymbolStackVariable::INITIAL);
224        let left = PartialSymbolStack::new(symbols.iter().copied().collect(), variable);
225        database
226            .starting_at_root(&left)
227            .into_iter()
228            .map(|(id, path)| (id.file, path.end))
229            .collect()
230    }
231
232    fn module_graph() -> StackGraph {
233        graph(&[(ROOT, MODULE), (MODULE, DEFINITION)])
234    }
235
236    #[test]
237    fn a_stack_with_the_symbols_a_root_path_needs_reaches_it() {
238        let graph = module_graph();
239        let mut database = PartialPathDatabase::default();
240        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
241
242        assert_eq!(
243            reached(&database, &["m", "x"], false),
244            [("a.py", DEFINITION)]
245        );
246    }
247
248    #[test]
249    fn a_stack_with_more_symbols_reaches_a_root_path_that_leaves_the_rest() {
250        let graph = module_graph();
251        let mut database = PartialPathDatabase::default();
252        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
253
254        assert_eq!(
255            reached(&database, &["m", "x", "y"], false),
256            [("a.py", DEFINITION)]
257        );
258    }
259
260    #[test]
261    fn a_stack_with_a_variable_reaches_a_root_path_that_needs_more_symbols() {
262        let graph = module_graph();
263        let mut database = PartialPathDatabase::default();
264        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
265
266        assert_eq!(reached(&database, &["m"], true), [("a.py", DEFINITION)]);
267    }
268
269    #[test]
270    fn a_stack_without_a_variable_does_not_reach_a_root_path_that_needs_more_symbols() {
271        let graph = module_graph();
272        let mut database = PartialPathDatabase::default();
273        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
274
275        assert_eq!(reached(&database, &["m"], false), []);
276    }
277
278    #[test]
279    fn a_stack_with_other_symbols_reaches_no_root_path() {
280        let graph = module_graph();
281        let mut database = PartialPathDatabase::default();
282        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
283
284        assert_eq!(reached(&database, &["n", "x"], true), []);
285    }
286
287    #[test]
288    fn a_stack_left_at_a_root_reaches_the_root_paths_of_every_file() {
289        let graph = module_graph();
290        let mut database = PartialPathDatabase::default();
291        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
292        database.add_file("b.py", &graph, DEFAULT_MAX_WORK);
293
294        assert_eq!(
295            reached(&database, &["m", "x"], false),
296            [("a.py", DEFINITION), ("b.py", DEFINITION)]
297        );
298    }
299
300    #[test]
301    fn a_path_that_starts_elsewhere_is_not_reached_from_the_root() {
302        let graph = graph(&[(ROOT, MODULE), (MODULE, DEFINITION), (REFERENCE, ROOT)]);
303        let mut database = PartialPathDatabase::default();
304        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
305
306        assert_eq!(reached(&database, &[], true), [("a.py", DEFINITION)]);
307    }
308
309    #[test]
310    fn a_dropped_file_is_not_reached_from_the_root() {
311        let graph = module_graph();
312        let mut database = PartialPathDatabase::default();
313        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
314        database.drop_file("a.py");
315
316        assert_eq!(reached(&database, &["m", "x"], false), []);
317    }
318
319    #[test]
320    fn a_file_stored_again_is_reached_from_the_root_once() {
321        let graph = module_graph();
322        let mut database = PartialPathDatabase::default();
323        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
324        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
325
326        assert_eq!(
327            reached(&database, &["m", "x"], false),
328            [("a.py", DEFINITION)]
329        );
330    }
331
332    #[test]
333    fn the_paths_starting_at_a_node_leave_out_those_starting_elsewhere() {
334        let graph = graph(&[(ROOT, DEFINITION), (REFERENCE, SCOPE), (SCOPE, ROOT)]);
335        let mut database = PartialPathDatabase::default();
336        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
337
338        assert_eq!(ends(&database, "a.py", REFERENCE), [ROOT]);
339    }
340
341    #[test]
342    fn a_node_id_of_one_file_does_not_reach_the_paths_of_another() {
343        let local = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
344        let other = graph(&[(REFERENCE, SCOPE), (SCOPE, ROOT)]);
345        let mut database = PartialPathDatabase::default();
346        database.add_file("a.py", &local, DEFAULT_MAX_WORK);
347        database.add_file("b.py", &other, DEFAULT_MAX_WORK);
348
349        assert_eq!(ends(&database, "a.py", REFERENCE), [DEFINITION]);
350    }
351
352    #[test]
353    fn a_file_never_stored_has_no_paths() {
354        let database = PartialPathDatabase::default();
355
356        assert_eq!(ends(&database, "a.py", REFERENCE), []);
357    }
358
359    #[test]
360    fn the_id_of_a_path_gives_back_that_path() {
361        let graph = graph(&[(ROOT, DEFINITION), (REFERENCE, SCOPE), (SCOPE, ROOT)]);
362        let mut database = PartialPathDatabase::default();
363        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
364        let found = database.starting_at("a.py", REFERENCE);
365
366        assert_eq!(
367            found
368                .map(|(id, path)| database.path(id) == Some(path))
369                .collect::<Vec<_>>(),
370            [true]
371        );
372    }
373
374    #[test]
375    fn storing_a_file_again_replaces_its_paths() {
376        let first = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
377        let second = graph(&[(REFERENCE, SCOPE), (SCOPE, ROOT)]);
378        let mut database = PartialPathDatabase::default();
379        database.add_file("a.py", &first, DEFAULT_MAX_WORK);
380        database.add_file("a.py", &second, DEFAULT_MAX_WORK);
381
382        assert_eq!(ends(&database, "a.py", REFERENCE), [ROOT]);
383    }
384
385    #[test]
386    fn the_id_of_a_dropped_file_gives_no_path() {
387        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
388        let mut database = PartialPathDatabase::default();
389        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
390        let id = PartialPathId {
391            file: "a.py",
392            index: 0,
393        };
394        database.drop_file("a.py");
395
396        assert_eq!(database.path(id), None);
397    }
398
399    #[test]
400    fn a_file_keeps_the_graph_its_paths_walk() {
401        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
402        let mut database = PartialPathDatabase::default();
403        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
404
405        assert_eq!(database.graph("a.py"), Some(&graph));
406    }
407
408    #[test]
409    fn a_file_keeps_the_endpoints_whose_search_ran_out_of_work() {
410        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, RELAY), (RELAY, SCOPE)]);
411        let mut database = PartialPathDatabase::default();
412        database.add_file("a.py", &graph, 100);
413
414        assert_eq!(database.exhausted("a.py"), [REFERENCE]);
415    }
416}