1use super::extractor::{
20 call_result_types, collect_assigned_identifiers, collect_function_scope_facts_from_node,
21 collect_scope_facts_from_parsed_source, enclosing_scope_facts, is_declaration_identifier,
22 slice,
23};
24use super::resolver::{
25 annotation_reference_candidates, resolve_callable_parameter_default_types,
26 resolve_constructor_types, resolve_receiver_type,
27};
28use crate::graph::PythonGraphSource;
29use crate::graph_support::PythonUsageSource;
30use crate::imports::resolve_fqn_candidates;
31use crate::usage_index::{usage_resolve_module_files, usage_scope_facts};
32use brokk_bifrost_core::analyzer::symbol_path::parse_symbol_path;
33use brokk_bifrost_core::analyzer::usages::inverted_edges::{
34 FileEdgeScanInput, PerFileEdges, classify_reference_node,
35};
36use brokk_bifrost_core::analyzer::usages::local_inference::LocalBindingsSnapshot;
37use brokk_bifrost_core::analyzer::usages::model::ImportKind;
38use brokk_bifrost_core::analyzer::{CodeUnit, Language, ProjectFile, Range};
39use brokk_bifrost_core::hash::{HashMap, HashSet};
40use std::sync::{Arc, Mutex};
41use tree_sitter::Node;
42
43pub struct PythonEdgeScan<'a> {
50 targets: &'a HashSet<String>,
51 targets_by_terminal: HashMap<String, Vec<String>>,
52 canonical_namespace_candidates: Mutex<HashMap<String, Arc<Vec<String>>>>,
53}
54
55impl<'a> PythonEdgeScan<'a> {
56 pub fn new(nodes: &HashSet<String>, targets: &'a HashSet<String>) -> Self {
59 debug_assert!(targets.is_subset(nodes));
60 let mut targets_by_terminal: HashMap<String, Vec<String>> = HashMap::default();
61 for target in targets {
62 let terminal = parse_symbol_path(Language::Python, target)
67 .pop()
68 .unwrap_or_else(|| target.clone());
69 targets_by_terminal
70 .entry(terminal)
71 .or_default()
72 .push(target.clone());
73 }
74 Self {
75 targets,
76 targets_by_terminal,
77 canonical_namespace_candidates: Mutex::new(HashMap::default()),
78 }
79 }
80
81 pub fn scan_file(
87 &self,
88 graph: &PythonGraphSource<'_>,
89 python: &dyn PythonUsageSource,
90 file: &ProjectFile,
91 input: &FileEdgeScanInput<'_>,
92 ) -> PerFileEdges {
93 let source = input.source;
94
95 let binder = python.import_binder_of(file);
100 let mut named: HashMap<String, String> = HashMap::default();
101 let mut namespace: HashMap<String, NamespaceBinding> = HashMap::default();
102 for (local, binding) in &binder.bindings {
103 match binding.kind {
104 ImportKind::Named => {
105 if let Some(imported) = &binding.imported_name {
106 let module = canonical_import_module_fqn(
107 graph,
108 python,
109 file,
110 &binding.module_specifier,
111 )
112 .unwrap_or_else(|| binding.module_specifier.clone());
113 named.insert(local.clone(), format!("{module}.{imported}"));
114 }
115 }
116 ImportKind::Namespace => {
117 let direct_module = binding.module_specifier.clone();
118 let imported_module = binding
119 .namespace_imported_module
120 .as_deref()
121 .unwrap_or(&direct_module);
122 let module = canonical_import_module_fqn(graph, python, file, imported_module);
123 let workspace_module = module.is_some();
124 let consumed_attributes = module.as_ref().map_or(0, |_| {
125 let imported_segments =
126 parse_symbol_path(Language::Python, imported_module);
127 let bound_segments = parse_symbol_path(Language::Python, &direct_module);
128 imported_segments.len().saturating_sub(bound_segments.len())
129 });
130 namespace.insert(
131 local.clone(),
132 NamespaceBinding {
133 module: module.unwrap_or(direct_module),
134 workspace_module,
135 consumed_attributes,
136 },
137 );
138 }
139 ImportKind::Default | ImportKind::CommonJsRequire | ImportKind::Glob => {}
140 }
141 }
142 let same_file: HashMap<String, String> = graph
143 .index
144 .declarations(file)
145 .into_iter()
146 .map(|unit| (unit.identifier().to_string(), unit.fq_name()))
147 .collect();
148
149 let scope_facts = usage_scope_facts(python, file, || {
153 collect_scope_facts_from_parsed_source(graph, python, file, source, input.root())
154 });
155
156 let mut ctx = PyScan {
157 graph,
158 python,
159 targets: self.targets,
160 targets_by_terminal: &self.targets_by_terminal,
161 file,
162 source,
163 named,
164 namespace,
165 same_file,
166 scope_facts: scope_facts.as_ref(),
167 canonical_namespace_candidates: &self.canonical_namespace_candidates,
168 input,
169 edges: PerFileEdges::default(),
170 };
171 scan_tree(input.root(), &mut ctx);
172 ctx.edges
173 }
174}
175
176fn canonical_import_module_fqn(
177 graph: &PythonGraphSource<'_>,
178 python: &dyn PythonUsageSource,
179 importing_file: &ProjectFile,
180 module_specifier: &str,
181) -> Option<String> {
182 let resolved = usage_resolve_module_files(python, importing_file, module_specifier);
183 let [module_file] = resolved.as_slice() else {
184 return None;
185 };
186 graph
187 .index
188 .declarations(module_file)
189 .into_iter()
190 .find(CodeUnit::is_module)
191 .map(|module| module.fq_name())
192}
193
194struct PyScan<'a> {
195 graph: &'a PythonGraphSource<'a>,
196 python: &'a dyn PythonUsageSource,
197 targets: &'a HashSet<String>,
198 targets_by_terminal: &'a HashMap<String, Vec<String>>,
199 file: &'a ProjectFile,
200 source: &'a str,
201 named: HashMap<String, String>,
202 namespace: HashMap<String, NamespaceBinding>,
203 same_file: HashMap<String, String>,
204 scope_facts: &'a HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
205 canonical_namespace_candidates: &'a Mutex<HashMap<String, Arc<Vec<String>>>>,
206 input: &'a FileEdgeScanInput<'a>,
207 edges: PerFileEdges,
208}
209
210struct NamespaceBinding {
211 module: String,
212 workspace_module: bool,
213 consumed_attributes: usize,
214}
215
216impl PyScan<'_> {
217 fn bare_callee(&self, text: &str) -> Option<String> {
220 if let Some(fqn) = self.named.get(text) {
221 return Some(fqn.clone());
222 }
223 if let Some(fqn) = self.namespace.get(text) {
224 return Some(fqn.module.clone());
225 }
226 if let Some(fqn) = self.same_file.get(text) {
227 return Some(fqn.clone());
228 }
229 None
230 }
231
232 fn receiver_type_fqn(
236 &self,
237 facts: &LocalBindingsSnapshot<String>,
238 receiver: &str,
239 ) -> Option<String> {
240 let resolution = facts.resolution_for(receiver);
241 let type_name = resolution
242 .as_precise()
243 .and_then(|targets| targets.iter().next())?;
244 resolve_receiver_type(self.graph, self.python, self.file, type_name, false)
250 .map(|unit| unit.fq_name())
251 }
252
253 fn record(&mut self, callee: String, node: Node<'_>) {
254 if !self.targets.contains(&callee) {
255 return;
256 }
257 self.edges.record_kind(
258 self.input,
259 callee,
260 classify_reference_node(node),
261 node.start_byte(),
262 node.end_byte(),
263 );
264 }
265
266 fn record_unproven_name(&mut self, name: &str, node: Node<'_>) {
267 let Some(targets) = self.targets_by_terminal.get(name) else {
268 return;
269 };
270 for target in targets {
271 self.edges.record_unproven(
272 self.input,
273 target.clone(),
274 node.start_byte(),
275 node.end_byte(),
276 );
277 }
278 }
279
280 fn canonical_namespace_candidates(&self, direct: &str) -> Arc<Vec<String>> {
281 if let Some(cached) = self
282 .canonical_namespace_candidates
283 .lock()
284 .expect("Python namespace candidate cache mutex poisoned")
285 .get(direct)
286 .cloned()
287 {
288 return cached;
289 }
290
291 let resolved: Arc<Vec<String>> = Arc::new(
292 resolve_fqn_candidates(self.python, direct, |name| {
293 self.graph.index.definitions(name).collect()
294 })
295 .into_iter()
296 .map(|unit| unit.fq_name())
297 .collect(),
298 );
299 self.canonical_namespace_candidates
300 .lock()
301 .expect("Python namespace candidate cache mutex poisoned")
302 .entry(direct.to_string())
303 .or_insert_with(|| resolved.clone())
304 .clone()
305 }
306}
307
308fn scan_tree(root: Node<'_>, ctx: &mut PyScan<'_>) {
309 let mut scopes: Vec<FunctionScope> = Vec::new();
312 walk(root, ctx, &mut scopes, None);
313}
314
315fn walk(
316 node: Node<'_>,
317 ctx: &mut PyScan<'_>,
318 scopes: &mut Vec<FunctionScope>,
319 facts: Option<usize>,
320) {
321 let mut merged_facts = Vec::new();
322 let mut stack = vec![WalkFrame::Enter { node, facts }];
323 while let Some(frame) = stack.pop() {
324 match frame {
325 WalkFrame::Enter { node, facts } => match node.kind() {
326 "import_statement" | "import_from_statement" => {}
327 "function_definition" | "lambda" => {
331 scopes.push(collect_function_scope(node, ctx.source));
332 let scope_facts = merged_enclosing_scope_facts(
333 ctx.graph,
334 ctx.file,
335 ctx.scope_facts,
336 &mut merged_facts,
337 node,
338 ctx.source,
339 facts,
340 );
341 stack.push(WalkFrame::ExitScope);
342 push_function_children(node, facts, scope_facts, &mut stack);
343 }
344 "class_definition" => push_children(node, None, &mut stack),
347 "identifier" => {
348 if !handle_annotation_reference(node, ctx) {
349 handle_identifier(node, ctx, scopes);
350 }
351 push_children(node, facts, &mut stack);
352 }
353 "attribute" => {
354 if handle_annotation_reference(node, ctx) {
355 continue;
356 }
357 let scope_facts = facts.and_then(|id| merged_facts.get(id));
358 handle_attribute(node, ctx, scopes, scope_facts);
359 push_children(node, facts, &mut stack);
360 }
361 "string_content" => {
362 handle_annotation_reference(node, ctx);
363 }
364 "keyword_argument" => {
365 handle_keyword_argument(node, ctx, scopes);
366 if let Some(value) = node.child_by_field_name("value") {
367 stack.push(WalkFrame::Enter { node: value, facts });
368 }
369 }
370 _ => push_children(node, facts, &mut stack),
371 },
372 WalkFrame::ExitScope => {
373 scopes.pop();
374 }
375 }
376 }
377}
378
379enum WalkFrame<'tree> {
380 Enter {
381 node: Node<'tree>,
382 facts: Option<usize>,
383 },
384 ExitScope,
385}
386
387fn push_children<'tree>(
388 node: Node<'tree>,
389 facts: Option<usize>,
390 stack: &mut Vec<WalkFrame<'tree>>,
391) {
392 for index in (0..node.named_child_count()).rev() {
393 if let Some(child) = node.named_child(index) {
394 stack.push(WalkFrame::Enter { node: child, facts });
395 }
396 }
397}
398
399fn push_function_children<'tree>(
400 function: Node<'tree>,
401 enclosing_facts: Option<usize>,
402 body_facts: Option<usize>,
403 stack: &mut Vec<WalkFrame<'tree>>,
404) {
405 let body = function.child_by_field_name("body");
406 for index in (0..function.named_child_count()).rev() {
407 if let Some(child) = function.named_child(index) {
408 let facts = if body == Some(child) {
412 body_facts
413 } else {
414 enclosing_facts
415 };
416 stack.push(WalkFrame::Enter { node: child, facts });
417 }
418 }
419}
420
421fn merged_enclosing_scope_facts(
422 graph: &PythonGraphSource<'_>,
423 file: &ProjectFile,
424 scope_facts: &HashMap<CodeUnit, LocalBindingsSnapshot<String>>,
425 merged_facts: &mut Vec<LocalBindingsSnapshot<String>>,
426 node: Node<'_>,
427 source: &str,
428 inherited: Option<usize>,
429) -> Option<usize> {
430 let structural_local = collect_function_scope_facts_from_node(node, source);
431 let local = if inherited.is_none() && node.kind() == "function_definition" {
436 enclosing_scope_facts(graph.index, file, scope_facts, node)
437 .cloned()
438 .unwrap_or(structural_local)
439 } else {
440 structural_local
441 };
442 match (local, inherited) {
443 (local, Some(inherited_id)) => {
444 let inherited = merged_facts.get(inherited_id)?;
445 let merged = inherited.merged_with_shadowing(&local);
446 let next_id = merged_facts.len();
447 merged_facts.push(merged);
448 Some(next_id)
449 }
450 (local, None) => {
451 let next_id = merged_facts.len();
452 merged_facts.push(local);
453 Some(next_id)
454 }
455 }
456}
457
458#[derive(Default)]
459struct FunctionScope {
460 locals: HashSet<String>,
461 parameters: HashSet<String>,
462}
463
464fn is_shadowed(scopes: &[FunctionScope], name: &str) -> bool {
465 scopes.iter().any(|scope| scope.locals.contains(name))
466}
467
468fn is_receiver_parameter(scopes: &[FunctionScope], name: &str) -> bool {
469 scopes
470 .iter()
471 .rev()
472 .any(|scope| scope.parameters.contains(name))
473}
474
475fn handle_identifier(node: Node<'_>, ctx: &mut PyScan<'_>, scopes: &[FunctionScope]) {
476 if node
478 .parent()
479 .is_some_and(|parent| parent.kind() == "attribute")
480 {
481 return;
482 }
483 if is_declaration_identifier(node) {
484 return;
485 }
486 let text = slice(node, ctx.source);
487 if text.is_empty() || is_shadowed(scopes, text) {
488 return;
489 }
490 if let Some(callee) = ctx.bare_callee(text) {
491 ctx.record(callee, node);
492 }
493}
494
495fn handle_annotation_reference(node: Node<'_>, ctx: &mut PyScan<'_>) -> bool {
496 let Some(candidates) =
497 annotation_reference_candidates(ctx.graph, ctx.python, ctx.file, ctx.source, node, false)
498 else {
499 return false;
500 };
501 let [candidate] = candidates.as_slice() else {
502 return !(node.kind() == "attribute" && candidates.is_empty());
503 };
504
505 let site = if node.kind() == "attribute" {
506 node.child_by_field_name("attribute").unwrap_or(node)
507 } else {
508 node
509 };
510 ctx.record(candidate.fq_name(), site);
511 true
512}
513
514fn handle_attribute(
515 node: Node<'_>,
516 ctx: &mut PyScan<'_>,
517 scopes: &[FunctionScope],
518 facts: Option<&LocalBindingsSnapshot<String>>,
519) {
520 let (Some(object), Some(attribute)) = (
521 node.child_by_field_name("object"),
522 node.child_by_field_name("attribute"),
523 ) else {
524 return;
525 };
526 let object_text = slice(object, ctx.source);
527 let attribute_text = slice(attribute, ctx.source);
528 if object_text.is_empty() || attribute_text.is_empty() {
529 return;
530 }
531 if object.kind() == "call" && ctx.targets_by_terminal.contains_key(attribute_text) {
532 for class in call_result_types(ctx.graph, ctx.python, ctx.file, ctx.source, object, facts) {
533 let direct = format!("{}.{attribute_text}", class.fq_name());
534 if ctx.targets.contains(&direct) {
535 ctx.record(direct, attribute);
536 continue;
537 }
538 if let Some(provider) = ctx.graph.hierarchy {
539 for ancestor in provider.get_ancestors(&class) {
540 let inherited = format!("{}.{attribute_text}", ancestor.fq_name());
541 if ctx.targets.contains(&inherited) {
542 ctx.record(inherited, attribute);
543 }
544 }
545 }
546 }
547 }
548 if let Some((root, attributes)) = attribute_chain(node) {
552 let root_text = slice(root, ctx.source);
553 if !root_text.is_empty()
554 && !is_shadowed(scopes, root_text)
555 && let Some(binding) = ctx.namespace.get(root_text)
556 {
557 let mut direct = binding.module.clone();
558 let workspace_module = binding.workspace_module;
559 let consumed_attributes = binding.consumed_attributes;
560 if object.kind() == "identifier" && ctx.targets.contains(&direct) {
561 ctx.record(direct.clone(), object);
562 }
563 for member in attributes.into_iter().skip(consumed_attributes) {
564 let member_text = slice(member, ctx.source);
565 if member_text.is_empty() {
566 return;
567 }
568 direct.push('.');
569 direct.push_str(member_text);
570 }
571 if ctx.targets.contains(&direct) {
572 ctx.record(direct, attribute);
573 return;
574 }
575 if workspace_module {
581 for resolved in ctx.canonical_namespace_candidates(&direct).iter() {
582 ctx.record(resolved.clone(), attribute);
583 }
584 }
585 return;
586 }
587 }
588
589 if let Some(facts) = facts
595 && ctx.targets_by_terminal.contains_key(attribute_text)
596 {
597 if matches!(object_text, "self" | "cls") {
598 ctx.record_unproven_name(attribute_text, attribute);
604 } else if let Some(type_fqn) = ctx.receiver_type_fqn(facts, object_text) {
605 ctx.record(format!("{type_fqn}.{attribute_text}"), attribute);
606 } else if object.kind() == "identifier" && !ctx.named.contains_key(object_text) {
607 let resolution = facts.resolution_for(object_text);
608 if resolution.is_ambiguous()
609 || (resolution.is_unknown() && is_receiver_parameter(scopes, object_text))
610 {
611 ctx.record_unproven_name(attribute_text, attribute);
612 }
613 }
614 }
615}
616
617fn handle_keyword_argument(node: Node<'_>, ctx: &mut PyScan<'_>, scopes: &[FunctionScope]) {
618 let (Some(name), Some(arguments)) = (node.child_by_field_name("name"), node.parent()) else {
619 return;
620 };
621 if name.kind() != "identifier" || arguments.kind() != "argument_list" {
622 return;
623 }
624 let Some(call) = arguments.parent().filter(|parent| parent.kind() == "call") else {
625 return;
626 };
627 let Some(function) = call.child_by_field_name("function") else {
628 return;
629 };
630 let member = slice(name, ctx.source);
631 if member.is_empty() || !ctx.targets_by_terminal.contains_key(member) {
632 return;
633 }
634 let scoped_class_fqn = if function.kind() == "identifier" {
635 enclosing_scope_facts(ctx.graph.index, ctx.file, ctx.scope_facts, function)
636 .and_then(|facts| ctx.receiver_type_fqn(facts, slice(function, ctx.source)))
637 } else {
638 None
639 };
640 let function_name = (function.kind() == "identifier").then(|| slice(function, ctx.source));
641 let mut default_classes = function_name.map_or_else(Vec::new, |local_name| {
642 resolve_callable_parameter_default_types(
643 ctx.graph, ctx.python, ctx.file, ctx.source, function, local_name,
644 )
645 });
646 let root_shadowed = leftmost_identifier(function)
647 .is_some_and(|root| is_shadowed(scopes, slice(root, ctx.source)));
648 let mut classes = if function_name == Some("cls") {
649 lexical_class(ctx, function).into_iter().collect()
650 } else {
651 if root_shadowed && scoped_class_fqn.is_none() && default_classes.is_empty() {
652 return;
653 }
654 if !root_shadowed {
655 default_classes.extend(resolve_constructor_types(
656 ctx.graph, ctx.python, ctx.file, ctx.source, function,
657 ));
658 }
659 default_classes
660 };
661 if let Some(fqn) = scoped_class_fqn {
662 classes.extend(ctx.graph.index.definitions(&fqn).filter(CodeUnit::is_class));
663 classes.sort();
664 classes.dedup();
665 }
666 for class in classes {
667 let direct = format!("{}.{member}", class.fq_name());
668 if ctx.targets.contains(&direct) {
669 ctx.record(direct, name);
670 continue;
671 }
672 if let Some(provider) = ctx.graph.hierarchy {
673 for ancestor in provider.get_ancestors(&class) {
674 let inherited = format!("{}.{member}", ancestor.fq_name());
675 if ctx.targets.contains(&inherited) {
676 ctx.record(inherited, name);
677 }
678 }
679 }
680 }
681}
682
683fn lexical_class(ctx: &PyScan<'_>, node: Node<'_>) -> Option<CodeUnit> {
684 let range = Range {
685 start_byte: node.start_byte(),
686 end_byte: node.end_byte(),
687 start_line: 0,
688 end_line: 0,
689 };
690 let enclosing = ctx.graph.index.enclosing_code_unit(ctx.file, &range)?;
691 if enclosing.is_class() {
692 Some(enclosing)
693 } else {
694 ctx.graph
695 .index
696 .parent_of(&enclosing)
697 .filter(CodeUnit::is_class)
698 }
699}
700
701fn leftmost_identifier(mut node: Node<'_>) -> Option<Node<'_>> {
702 loop {
703 match node.kind() {
704 "identifier" => return Some(node),
705 "attribute" => node = node.child_by_field_name("object")?,
706 _ => return None,
707 }
708 }
709}
710
711fn attribute_chain<'a>(node: Node<'a>) -> Option<(Node<'a>, Vec<Node<'a>>)> {
712 let mut attributes = Vec::new();
713 let mut current = node;
714 loop {
715 if current.kind() != "attribute" {
716 return None;
717 }
718 attributes.push(current.child_by_field_name("attribute")?);
719 current = current.child_by_field_name("object")?;
720 if current.kind() == "identifier" {
721 attributes.reverse();
722 return Some((current, attributes));
723 }
724 }
725}
726
727fn collect_function_scope(func: Node<'_>, source: &str) -> FunctionScope {
732 let mut scope = FunctionScope::default();
733 if let Some(params) = func.child_by_field_name("parameters") {
734 collect_parameter_names(params, source, &mut scope.parameters);
735 scope.locals.extend(scope.parameters.iter().cloned());
736 }
737 if let Some(body) = func.child_by_field_name("body") {
738 collect_bound_targets(body, source, &mut scope.locals);
739 }
740 scope
741}
742
743fn collect_parameter_names(params: Node<'_>, source: &str, out: &mut HashSet<String>) {
744 let mut cursor = params.walk();
745 for child in params.named_children(&mut cursor) {
746 let name = match child.kind() {
747 "identifier" => Some(child),
748 _ => child
751 .child_by_field_name("name")
752 .or_else(|| child.named_child(0).filter(|n| n.kind() == "identifier")),
753 };
754 if let Some(name) = name {
755 let text = slice(name, source).trim();
756 if !text.is_empty() {
757 out.insert(text.to_string());
758 }
759 }
760 }
761}
762
763fn collect_bound_targets(node: Node<'_>, source: &str, out: &mut HashSet<String>) {
767 let mut stack = vec![node];
768 while let Some(node) = stack.pop() {
769 match node.kind() {
770 "function_definition" | "class_definition" => {
771 if let Some(name) = node.child_by_field_name("name") {
772 let text = slice(name, source).trim();
773 if !text.is_empty() {
774 out.insert(text.to_string());
775 }
776 }
777 continue;
778 }
779 "lambda" => continue,
780 "assignment" | "augmented_assignment" | "for_statement" | "for_in_clause" => {
781 if let Some(left) = node.child_by_field_name("left") {
782 collect_assigned_identifiers(left, source, out);
783 }
784 }
785 "named_expression" => {
786 if let Some(name) = node.child_by_field_name("name") {
787 collect_assigned_identifiers(name, source, out);
788 }
789 }
790 _ => {}
791 }
792 let mut cursor = node.walk();
793 let mut children: Vec<Node<'_>> = node.named_children(&mut cursor).collect();
794 children.reverse();
795 stack.extend(children);
796 }
797}