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