1use 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
20pub 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 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
54pub struct FlattenedScopeCache {
59 scopes: HashMap<FileId, ScopeMap>,
60}
61
62impl FlattenedScopeCache {
63 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(¤t) {
104 let default_vis = ctx
105 .file_languages
106 .get(¤t)
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(¤t) {
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(¤t)
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 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 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 pub fn len(&self) -> usize {
194 self.scopes.len()
195 }
196
197 pub fn is_empty(&self) -> bool {
199 self.scopes.is_empty()
200 }
201}
202
203pub 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 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 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 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
279pub 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 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 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}