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#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
14#[non_exhaustive]
15pub struct PartialPathId<'a> {
16 pub file: &'a str,
18 pub index: usize,
20}
21
22#[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 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 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 #[must_use]
76 pub fn graph(&self, file: &str) -> Option<&'a StackGraph> {
77 self.files.get(file).map(|stored| stored.graph)
78 }
79
80 #[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 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 #[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 #[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}