Skip to main content

meta_ast/graph/
resolver.rs

1//! Reference resolution via FlattenedScopeCache.
2//!
3//! Pre-computes the visible scope per file by DFS-ing the import graph
4//! once, then resolves references with O(1) lookups instead of
5//! per-reference BFS.
6
7use std::collections::{HashMap, HashSet, VecDeque};
8use std::path::PathBuf;
9
10use rayon::prelude::*;
11
12use crate::error::{Diagnostic, Severity};
13use crate::language::LangId;
14use crate::model::{FileExtraction, FileId, SymbolId, Visibility};
15
16type ScopeMap = HashMap<String, Vec<(SymbolId, f32)>>;
17pub(crate) type SymbolIndexEntry = (SymbolId, String, LangId, Option<Visibility>);
18pub(crate) type SymbolIndex = HashMap<FileId, Vec<SymbolIndexEntry>>;
19
20/// Bundles the data needed for scope resolution across files.
21pub struct ResolutionContext {
22    pub symbol_index: SymbolIndex,
23    pub import_adjacency: HashMap<FileId, Vec<FileId>>,
24    pub file_languages: HashMap<FileId, LangId>,
25    pub file_paths: HashMap<FileId, PathBuf>,
26}
27
28impl ResolutionContext {
29    /// Build a ResolutionContext from extraction results and graph data.
30    pub fn from_extractions(
31        extractions: &[FileExtraction],
32        path_to_file_id: &HashMap<PathBuf, FileId>,
33        import_adjacency: HashMap<FileId, Vec<FileId>>,
34    ) -> Self {
35        let symbol_index = build_symbol_index(extractions, path_to_file_id);
36        let file_languages: HashMap<_, _> = extractions
37            .iter()
38            .filter_map(|f| Some((path_to_file_id.get(&f.path)?.to_owned(), f.lang)))
39            .collect();
40        let file_paths: HashMap<_, _> = path_to_file_id
41            .iter()
42            .map(|(path, &fid)| (fid, path.clone()))
43            .collect();
44
45        Self {
46            symbol_index,
47            import_adjacency,
48            file_languages,
49            file_paths,
50        }
51    }
52}
53
54/// Pre-computed visible scope for each file.
55///
56/// Scope = own symbols + public symbols from imported files transitively.
57/// Local symbols take priority over imported (shadowing).
58pub struct FlattenedScopeCache {
59    scopes: HashMap<FileId, ScopeMap>,
60}
61
62impl FlattenedScopeCache {
63    /// Build the scope cache from the file->symbols index and import adjacency.
64    ///
65    /// For each file, DFS over import edges, collecting public symbols from
66    /// reachable files. Confidence decays with distance:
67    /// - 1.0: own file or direct import, same language
68    /// - 0.8: transitive import, same language
69    /// - 0.6: cross-language imports
70    pub fn build(ctx: &ResolutionContext, diagnostics: &mut Vec<Diagnostic>) -> Self {
71        let results: Vec<(FileId, ScopeMap, Vec<Diagnostic>)> = ctx
72            .symbol_index
73            .par_iter()
74            .map(|(&file_id, _)| {
75                let (scope, diags) = Self::compute_scope(file_id, ctx);
76                (file_id, scope, diags)
77            })
78            .collect();
79
80        let mut scopes = HashMap::with_capacity(results.len());
81        for (file_id, scope, diags) in results {
82            scopes.insert(file_id, scope);
83            diagnostics.extend(diags);
84        }
85
86        Self { scopes }
87    }
88
89    fn compute_scope(file_id: FileId, ctx: &ResolutionContext) -> (ScopeMap, Vec<Diagnostic>) {
90        let mut diagnostics = Vec::new();
91        let source_lang = ctx.file_languages.get(&file_id).copied();
92        let mut scope: ScopeMap = HashMap::new();
93        let mut visited: HashSet<FileId> = HashSet::new();
94        let mut queue: VecDeque<(FileId, usize)> = VecDeque::new();
95
96        queue.push_back((file_id, 0));
97
98        while let Some((current, distance)) = queue.pop_front() {
99            if !visited.insert(current) {
100                continue;
101            }
102
103            if let Some(symbols) = ctx.symbol_index.get(&current) {
104                let default_vis = ctx
105                    .file_languages
106                    .get(&current)
107                    .map(|lang| lang.spec().default_visibility)
108                    .unwrap_or(crate::language::DefaultVisibility::PublicByDefault);
109
110                for (sym_id, name, sym_lang, visibility) in symbols {
111                    let is_public = match visibility {
112                        Some(Visibility::Public) => true,
113                        Some(Visibility::Private) => current == file_id,
114                        None => {
115                            matches!(
116                                default_vis,
117                                crate::language::DefaultVisibility::PublicByDefault
118                            ) || current == file_id
119                        }
120                    };
121
122                    if !is_public {
123                        continue;
124                    }
125
126                    let same_lang = source_lang.is_some() && source_lang == Some(*sym_lang);
127                    let diff_lang = source_lang.is_some() && source_lang != Some(*sym_lang);
128
129                    let confidence = if distance == 0 || (distance == 1 && same_lang) {
130                        1.0
131                    } else if diff_lang {
132                        0.6
133                    } else {
134                        0.8
135                    };
136
137                    if let Some(entries) = scope.get_mut(name) {
138                        entries.push((*sym_id, confidence));
139                    } else {
140                        scope.insert(name.clone(), vec![(*sym_id, confidence)]);
141                    }
142                }
143            }
144
145            if let Some(neighbors) = ctx.import_adjacency.get(&current) {
146                for &neighbor in neighbors {
147                    if !visited.contains(&neighbor) {
148                        queue.push_back((neighbor, distance + 1));
149                    } else if neighbor == file_id {
150                        let path = ctx
151                            .file_paths
152                            .get(&current)
153                            .cloned()
154                            .unwrap_or_else(|| PathBuf::from("<unknown>"));
155                        let root_path = ctx
156                            .file_paths
157                            .get(&file_id)
158                            .map(|p| p.display().to_string())
159                            .unwrap_or_else(|| "<unknown>".to_string());
160                        diagnostics.push(Diagnostic {
161                            path,
162                            severity: Severity::Warning,
163                            message: format!("circular import: {} -> {}", current.0, root_path),
164                            source_range: None,
165                        });
166                    }
167                }
168            }
169        }
170
171        // Sort each entry: higher confidence first, then by symbol_id (stable)
172        for entries in scope.values_mut() {
173            entries.sort_by(|a, b| {
174                b.1.partial_cmp(&a.1)
175                    .unwrap_or(std::cmp::Ordering::Equal)
176                    .then(a.0.0.cmp(&b.0.0))
177            });
178        }
179
180        (scope, diagnostics)
181    }
182
183    /// Look up a name in a file's flattened scope.
184    ///
185    /// Returns matching symbols with confidence scores, or None if not found.
186    pub fn resolve(&self, file_id: FileId, name: &str) -> Option<&[(SymbolId, f32)]> {
187        self.scopes
188            .get(&file_id)
189            .and_then(|s| s.get(name).map(|v| v.as_slice()))
190    }
191
192    /// Returns the number of scopes in the cache.
193    pub fn len(&self) -> usize {
194        self.scopes.len()
195    }
196
197    /// Returns true if the cache is empty.
198    pub fn is_empty(&self) -> bool {
199        self.scopes.is_empty()
200    }
201}
202
203/// Resolve all references across extracted files.
204///
205/// Returns a list of (source_symbol_id, target_symbol_id, confidence) triples
206/// representing ReferenceEdges to add. Confidence is threaded from the
207/// FlattenedScopeCache (1.0 local/direct, 0.8 transitive, 0.6 cross-language).
208/// Warnings for unresolved references are appended to `diagnostics`.
209pub fn resolve_all_references(
210    extractions: &[FileExtraction],
211    path_to_file_id: &HashMap<PathBuf, FileId>,
212    scope_cache: &FlattenedScopeCache,
213    diagnostics: &mut Vec<Diagnostic>,
214) -> Vec<(SymbolId, SymbolId, f32)> {
215    #[allow(clippy::type_complexity)]
216    let results: Vec<(Vec<(SymbolId, SymbolId, f32)>, Vec<Diagnostic>)> = extractions
217        .par_iter()
218        .map(|file_ext| {
219            let mut local_edges = Vec::new();
220            let mut local_diags = Vec::new();
221
222            let file_id = match path_to_file_id.get(&file_ext.path) {
223                Some(&id) => id,
224                None => return (local_edges, local_diags),
225            };
226
227            let file_path = &file_ext.path;
228            for ref_ in &file_ext.references {
229                if let Some(matches) = scope_cache.resolve(file_id, &ref_.name) {
230                    // Find the source symbol that contains this reference range
231                    // We use rfind to get the innermost symbol (e.g. method inside class)
232                    let source_sym = file_ext.symbols.iter().rfind(|s| {
233                        s.source_range.byte_start <= ref_.range.byte_start
234                            && s.source_range.byte_end >= ref_.range.byte_end
235                    });
236
237                    if let Some(source) = source_sym {
238                        for &(target_id, confidence) in matches {
239                            // Don't add self-references (symbol to itself)
240                            if source.id != target_id {
241                                local_edges.push((source.id, target_id, confidence));
242                            }
243                        }
244                    }
245                } else {
246                    local_diags.push(Diagnostic {
247                        path: file_path.clone(),
248                        severity: Severity::Warning,
249                        message: format!("unresolved reference: '{}'", ref_.name),
250                        source_range: Some(ref_.range.clone()),
251                    });
252                }
253            }
254            (local_edges, local_diags)
255        })
256        .collect();
257
258    let mut edges = Vec::new();
259    for (local_edges, local_diags) in results {
260        edges.extend(local_edges);
261        diagnostics.extend(local_diags);
262    }
263
264    // Deduplicate: max-merge confidence for same (src, dst) pairs
265    let mut seen: HashMap<(SymbolId, SymbolId), f32> = HashMap::with_capacity(edges.len());
266    for (src, dst, conf) in edges {
267        seen.entry((src, dst))
268            .and_modify(|e| *e = e.max(conf))
269            .or_insert(conf);
270    }
271    let mut deduped: Vec<_> = seen
272        .into_iter()
273        .map(|((src, dst), conf)| (src, dst, conf))
274        .collect();
275    deduped.sort_by_key(|(a, b, _)| (a.0, b.0));
276    deduped
277}
278
279/// Build a symbol index from extracted files and a path-to-FileId mapping.
280///
281/// Returns: SymbolIndex
282pub fn build_symbol_index(
283    extractions: &[FileExtraction],
284    path_to_file_id: &HashMap<PathBuf, FileId>,
285) -> SymbolIndex {
286    let mut index: SymbolIndex = HashMap::new();
287
288    for file_ext in extractions {
289        if let Some(&file_id) = path_to_file_id.get(&file_ext.path) {
290            let entries: Vec<_> = file_ext
291                .symbols
292                .iter()
293                .map(|s| (s.id, s.name.clone(), s.language, s.visibility))
294                .collect();
295            index.entry(file_id).or_default().extend(entries);
296        }
297    }
298
299    index
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use std::path::PathBuf;
306
307    #[test]
308    fn empty_cache() {
309        let cache = FlattenedScopeCache {
310            scopes: HashMap::new(),
311        };
312        assert!(cache.is_empty());
313        assert_eq!(cache.len(), 0);
314        assert!(cache.resolve(FileId(0), "foo").is_none());
315    }
316
317    #[test]
318    fn scope_cache_resolve_own_file() {
319        let mut symbol_index: SymbolIndex = HashMap::new();
320        symbol_index.insert(
321            FileId(0),
322            vec![(SymbolId(10), "main".into(), LangId::Python, None)],
323        );
324
325        let ctx = ResolutionContext {
326            symbol_index,
327            import_adjacency: HashMap::new(),
328            file_languages: HashMap::from([(FileId(0), LangId::Python)]),
329            file_paths: HashMap::new(),
330        };
331
332        let cache = FlattenedScopeCache::build(&ctx, &mut Vec::new());
333        let result = cache.resolve(FileId(0), "main");
334        assert!(result.is_some());
335        let matches = result.unwrap();
336        assert_eq!(matches.len(), 1);
337        assert_eq!(matches[0].0, SymbolId(10));
338        assert_eq!(matches[0].1, 1.0);
339    }
340
341    #[test]
342    fn scope_cache_resolve_imported_symbol() {
343        let mut symbol_index = HashMap::new();
344        symbol_index.insert(FileId(0), vec![]);
345        symbol_index.insert(
346            FileId(1),
347            vec![(
348                SymbolId(20),
349                "helper".into(),
350                LangId::Python,
351                Some(Visibility::Public),
352            )],
353        );
354
355        let ctx = ResolutionContext {
356            symbol_index,
357            import_adjacency: HashMap::from([(FileId(0), vec![FileId(1)])]),
358            file_languages: HashMap::from([
359                (FileId(0), LangId::Python),
360                (FileId(1), LangId::Python),
361            ]),
362            file_paths: HashMap::new(),
363        };
364
365        let cache = FlattenedScopeCache::build(&ctx, &mut Vec::new());
366        let result = cache.resolve(FileId(0), "helper");
367        assert!(result.is_some());
368        let matches = result.unwrap();
369        assert_eq!(matches.len(), 1);
370        assert_eq!(matches[0].0, SymbolId(20));
371        assert_eq!(matches[0].1, 1.0);
372    }
373
374    #[test]
375    fn scope_cache_missing_symbol() {
376        let mut symbol_index = HashMap::new();
377        symbol_index.insert(
378            FileId(0),
379            vec![(SymbolId(10), "foo".into(), LangId::Python, None)],
380        );
381
382        let ctx = ResolutionContext {
383            symbol_index,
384            import_adjacency: HashMap::new(),
385            file_languages: HashMap::from([(FileId(0), LangId::Python)]),
386            file_paths: HashMap::new(),
387        };
388
389        let cache = FlattenedScopeCache::build(&ctx, &mut Vec::new());
390        assert!(cache.resolve(FileId(0), "bar").is_none());
391    }
392
393    #[test]
394    fn scope_cache_cycle_safe() {
395        let mut symbol_index = HashMap::new();
396        symbol_index.insert(
397            FileId(0),
398            vec![(
399                SymbolId(10),
400                "a".into(),
401                LangId::Python,
402                Some(Visibility::Public),
403            )],
404        );
405        symbol_index.insert(
406            FileId(1),
407            vec![(
408                SymbolId(20),
409                "b".into(),
410                LangId::Python,
411                Some(Visibility::Public),
412            )],
413        );
414
415        // Cycle: 0 -> 1 -> 0
416        let ctx = ResolutionContext {
417            symbol_index,
418            import_adjacency: HashMap::from([
419                (FileId(0), vec![FileId(1)]),
420                (FileId(1), vec![FileId(0)]),
421            ]),
422            file_languages: HashMap::from([
423                (FileId(0), LangId::Python),
424                (FileId(1), LangId::Python),
425            ]),
426            file_paths: HashMap::new(),
427        };
428
429        let cache = FlattenedScopeCache::build(&ctx, &mut Vec::new());
430        // Should not infinite loop
431        assert!(cache.resolve(FileId(0), "b").is_some());
432        assert!(cache.resolve(FileId(1), "a").is_some());
433    }
434
435    #[test]
436    fn scope_cache_cross_language_confidence() {
437        let mut symbol_index = HashMap::new();
438        symbol_index.insert(FileId(0), vec![]);
439        symbol_index.insert(
440            FileId(1),
441            vec![(
442                SymbolId(20),
443                "util".into(),
444                LangId::Rust,
445                Some(Visibility::Public),
446            )],
447        );
448
449        let ctx = ResolutionContext {
450            symbol_index,
451            import_adjacency: HashMap::from([(FileId(0), vec![FileId(1)])]),
452            file_languages: HashMap::from([(FileId(0), LangId::Python), (FileId(1), LangId::Rust)]),
453            file_paths: HashMap::new(),
454        };
455
456        let cache = FlattenedScopeCache::build(&ctx, &mut Vec::new());
457        let result = cache.resolve(FileId(0), "util");
458        assert!(result.is_some());
459        assert_eq!(result.unwrap()[0].1, 0.6);
460    }
461
462    #[test]
463    fn resolve_references_creates_edges() {
464        use crate::model::{LineColumn, SourceRange, Symbol, SymbolKind, UnresolvedReference};
465
466        let sym_a = Symbol {
467            id: SymbolId(1),
468            name: "caller".into(),
469            kind: SymbolKind::Function,
470            language: LangId::Python,
471            file_path: PathBuf::from("/proj/a.py"),
472            source_range: SourceRange {
473                byte_start: 0,
474                byte_end: 50,
475                start: LineColumn { line: 0, column: 0 },
476                end: LineColumn { line: 2, column: 0 },
477            },
478            visibility: None,
479            signature: None,
480            docstring: None,
481            is_async: false,
482        };
483
484        let file = FileExtraction {
485            path: PathBuf::from("/proj/a.py"),
486            lang: LangId::Python,
487            symbols: vec![sym_a],
488            imports: vec![],
489            references: vec![UnresolvedReference {
490                name: "helper".into(),
491                range: SourceRange {
492                    byte_start: 20,
493                    byte_end: 26,
494                    start: LineColumn { line: 1, column: 4 },
495                    end: LineColumn {
496                        line: 1,
497                        column: 10,
498                    },
499                },
500            }],
501            diagnostics: vec![],
502            ast_node_count: 0,
503        };
504
505        let mut path_to_file_id = HashMap::new();
506        path_to_file_id.insert(PathBuf::from("/proj/a.py"), FileId(0));
507
508        let mut scopes: HashMap<FileId, ScopeMap> = HashMap::new();
509        let mut scope = HashMap::new();
510        scope.insert("helper".into(), vec![(SymbolId(99), 1.0)]);
511        scopes.insert(FileId(0), scope);
512
513        let cache = FlattenedScopeCache { scopes };
514
515        let edges = resolve_all_references(&[file], &path_to_file_id, &cache, &mut Vec::new());
516        assert_eq!(edges.len(), 1);
517        assert_eq!(edges[0].0, SymbolId(1));
518        assert_eq!(edges[0].1, SymbolId(99));
519        assert_eq!(edges[0].2, 1.0);
520    }
521}