use std::collections::{HashMap, HashSet}; use crate::{CodeElement, ModuleName, Relationship, RelationshipKind}; #[derive(Debug, Clone)] pub struct CodeGraph { elements: Vec, relationships: Vec, } impl Default for CodeGraph { fn default() -> Self { Self::new() } } impl CodeGraph { pub fn new() -> Self { Self { elements: Vec::new(), relationships: Vec::new(), } } pub fn add_element(&mut self, element: CodeElement) { self.elements.push(element); } pub fn add_relationship(&mut self, relationship: Relationship) { self.relationships.push(relationship); } pub fn elements(&self) -> &[CodeElement] { &self.elements } pub fn relationships(&self) -> &[Relationship] { &self.relationships } pub fn modules(&self) -> Vec { let mut seen = HashSet::new(); let mut modules = Vec::new(); for element in &self.elements { if let Some(module) = element.module() && seen.insert(module.as_str().to_string()) { modules.push(module.clone()); } } modules } pub fn elements_by_module(&self) -> (HashMap>, Vec<&CodeElement>) { let mut grouped: HashMap> = HashMap::new(); let mut ungrouped: Vec<&CodeElement> = Vec::new(); for element in &self.elements { if let Some(module) = element.module() { grouped .entry(module.as_str().to_string()) .or_default() .push(element); } else { ungrouped.push(element); } } (grouped, ungrouped) } pub fn resolve_relationships(self) -> CodeGraph { let qualified_names: HashSet<&str> = self.elements.iter().map(|e| e.qualified_name()).collect(); // Also keep bare name lookup for import relationships and unqualified fallback let all_bare_names: HashSet<&str> = self.elements.iter().map(|e| e.name()).collect(); let mut resolved = CodeGraph::new(); for element in &self.elements { resolved.add_element(element.clone()); } for rel in &self.relationships { match rel.kind() { RelationshipKind::Import => { resolved.add_relationship(rel.clone()); } _ => { let src_ok = qualified_names.contains(rel.source()) || all_bare_names.contains(rel.source()); let tgt_ok = qualified_names.contains(rel.target()) || all_bare_names.contains(rel.target()); if src_ok && tgt_ok { resolved.add_relationship(rel.clone()); } } } } resolved } pub fn filter_external_imports(self, known_modules: &HashSet) -> CodeGraph { let module_names: HashSet = self .modules() .iter() .map(|m| m.as_str().to_lowercase()) .collect(); let all_known: HashSet<&str> = known_modules .iter() .map(|s| s.as_str()) .chain(module_names.iter().map(|s| s.as_str())) .collect(); let mut filtered = CodeGraph::new(); for element in &self.elements { filtered.add_element(element.clone()); } for rel in &self.relationships { if rel.kind() == RelationshipKind::Import { let target_top = rel.target().split('.').next().unwrap_or("").to_lowercase(); if !all_known.contains(target_top.as_str()) { continue; } } filtered.add_relationship(rel.clone()); } filtered } pub fn qualify(self) -> CodeGraph { // Build lookup: bare name -> Vec (for disambiguation) let mut name_to_qualified: HashMap<&str, Vec> = HashMap::new(); for element in &self.elements { let qn = match element.module() { Some(m) => format!("{}::{}", m.as_str(), element.name()), None => element.name().to_string(), }; name_to_qualified .entry(element.name()) .or_default() .push(qn); } // Build lookup: file_path -> qualified source names in that file let mut file_to_qualified: HashMap<&str, Vec> = HashMap::new(); for element in &self.elements { let qn = match element.module() { Some(m) => format!("{}::{}", m.as_str(), element.name()), None => element.name().to_string(), }; file_to_qualified .entry(element.file_path().as_str()) .or_default() .push(qn); } let mut qualified = CodeGraph::new(); // 1. Qualify element names for element in &self.elements { let qn = match element.module() { Some(m) => format!("{}::{}", m.as_str(), element.name()), None => element.name().to_string(), }; qualified.add_element(element.clone().with_qualified_name(qn)); } // 2. Rewrite relationship source/target for rel in &self.relationships { // Qualify source: find the qualified name of the source in its file let src_qualified = rel .source_file() .and_then(|f| file_to_qualified.get(f.as_str())) .and_then(|qns| { qns.iter().find(|qn| { qn.ends_with(&format!("::{}", rel.source())) || *qn == rel.source() }) }) .cloned() .unwrap_or_else(|| { // Fall back: unambiguous lookup name_to_qualified .get(rel.source()) .filter(|v| v.len() == 1) .and_then(|v| v.first()) .cloned() .unwrap_or_else(|| rel.source().to_string()) }); // Qualify target: prefer same module as source when ambiguous let src_module = src_qualified.split("::").next().unwrap_or(""); let tgt_qualified = match name_to_qualified.get(rel.target()) { Some(candidates) if candidates.len() == 1 => candidates[0].clone(), Some(candidates) => { // Prefer same module as source candidates .iter() .find(|qn| qn.starts_with(&format!("{}::", src_module))) .cloned() .unwrap_or_else(|| rel.target().to_string()) } None => rel.target().to_string(), }; let new_rel = Relationship::new(&src_qualified, &tgt_qualified, rel.kind()) .unwrap_or_else(|_| rel.clone()); let new_rel = if let Some(f) = rel.source_file() { new_rel.with_source_file(f.clone()) } else { new_rel }; qualified.add_relationship(new_rel); } qualified } pub fn cross_module_deps_for(&self, module: &ModuleName) -> Vec<(ModuleName, usize)> { let module_element_qnames: HashSet<&str> = self .elements .iter() .filter(|e| e.module().is_some_and(|m| m == module)) .map(|e| e.qualified_name()) .collect(); let target_module_of: HashMap<&str, Option<&ModuleName>> = self .elements .iter() .map(|e| (e.qualified_name(), e.module())) .collect(); let mut counts: HashMap<&str, usize> = HashMap::new(); for rel in &self.relationships { if !module_element_qnames.contains(rel.source()) { continue; } if module_element_qnames.contains(rel.target()) { continue; } if let Some(Some(target_mod)) = target_module_of.get(rel.target()) { *counts.entry(target_mod.as_str()).or_insert(0) += 1; } } let mut result: Vec<(ModuleName, usize)> = counts .into_iter() .filter_map(|(name, count)| ModuleName::new(name).ok().map(|m| (m, count))) .collect(); result.sort_by(|a, b| a.0.as_str().cmp(b.0.as_str())); result } pub fn subgraph_by_module(&self, module: &ModuleName) -> CodeGraph { let filtered_elements: Vec = self .elements .iter() .filter(|e| e.module().is_some_and(|m| m == module)) .cloned() .collect(); let element_qnames: HashSet<&str> = filtered_elements .iter() .map(|e| e.qualified_name()) .collect(); let filtered_relationships: Vec = self .relationships .iter() .filter(|r| element_qnames.contains(r.source()) && element_qnames.contains(r.target())) .cloned() .collect(); CodeGraph { elements: filtered_elements, relationships: filtered_relationships, } } }