Skip to main content

project_map_cli_rust/core/
query_engine.rs

1use crate::core::graph::{ProjectGraph, NodeData, NodeType};
2use crate::error::Result;
3use petgraph::visit::Dfs;
4use std::path::Path;
5
6pub struct QueryEngine {
7    graph: ProjectGraph,
8}
9
10impl QueryEngine {
11    pub fn load(path: &Path) -> Result<Self> {
12        let graph = ProjectGraph::load(path)?;
13        Ok(Self { graph })
14    }
15
16    pub fn find_symbols(&self, query: &str) -> Vec<NodeData> {
17        let keywords: Vec<String> = query.to_lowercase()
18            .split_whitespace()
19            .map(|s| s.to_string())
20            .collect();
21            
22        if keywords.is_empty() {
23            return Vec::new();
24        }
25
26        self.graph.graph.node_weights()
27            .filter(|n| {
28                if n.node_type != NodeType::Symbol {
29                    return false;
30                }
31                
32                let name_lower = n.name.to_lowercase();
33                let doc_lower = n.docstring.as_ref().map(|d| d.to_lowercase()).unwrap_or_default();
34                
35                // Match if ALL keywords are found in either name or docstring
36                keywords.iter().all(|k| name_lower.contains(k) || doc_lower.contains(k))
37            })
38            .cloned()
39            .collect()
40    }
41
42    pub fn get_file_outline(&self, path: &str) -> Vec<NodeData> {
43        if path == "." || path == "./" {
44            return self.graph.graph.node_weights()
45                .filter(|n| n.node_type == NodeType::File)
46                .cloned()
47                .collect();
48        }
49
50        let file_node = self.graph.graph.node_indices()
51            .find(|i| self.graph.graph[*i].node_type == NodeType::File && self.graph.graph[*i].path == path);
52
53        if let Some(idx) = file_node {
54            self.graph.graph.neighbors_directed(idx, petgraph::Direction::Outgoing)
55                .map(|n| self.graph.graph[n].clone())
56                .collect()
57        } else {
58            Vec::new()
59        }
60    }
61
62    pub fn analyze_impact(&self, name: &str) -> Vec<NodeData> {
63        let node_idx = self.graph.graph.node_indices()
64            .find(|i| self.graph.graph[*i].name == name);
65        
66        if let Some(start_node) = node_idx {
67            // Impact: who do I depend on? (Outgoing edges)
68            let mut dfs = Dfs::new(&self.graph.graph, start_node);
69            let mut results = Vec::new();
70            while let Some(nx) = dfs.next(&self.graph.graph) {
71                if nx != start_node {
72                    results.push(self.graph.graph[nx].clone());
73                }
74            }
75            results
76        } else {
77            Vec::new()
78        }
79    }
80
81    pub fn check_blast_radius(&self, path: &str, symbol: &str) -> Vec<NodeData> {
82        let node_idx = self.graph.graph.node_indices()
83            .find(|i| {
84                let node = &self.graph.graph[*i];
85                node.path == path && node.name == symbol && node.node_type == NodeType::Symbol
86            });
87        
88        let start_node = if let Some(idx) = node_idx {
89            idx
90        } else {
91            // Try matching just by file path if symbol not found
92            let file_node = self.graph.graph.node_indices()
93                .find(|i| {
94                    let node = &self.graph.graph[*i];
95                    node.path == path && node.node_type == NodeType::File
96                });
97            if let Some(idx) = file_node { idx } else { return Vec::new(); }
98        };
99
100        // Blast Radius: who depends on me? (Incoming edges)
101        // We need to use a graph traversal that follows edges backwards.
102        // petgraph's Dfs follows outgoing edges. To follow incoming, we can use a custom traversal or reverse the graph.
103        // Alternatively, we can use neighbors_directed with Incoming in a loop.
104        
105        let mut results = Vec::new();
106        let mut stack = vec![start_node];
107        let mut visited = std::collections::HashSet::new();
108        visited.insert(start_node);
109
110        while let Some(current) = stack.pop() {
111            for neighbor in self.graph.graph.neighbors_directed(current, petgraph::Direction::Incoming) {
112                if visited.insert(neighbor) {
113                    results.push(self.graph.graph[neighbor].clone());
114                    stack.push(neighbor);
115                }
116            }
117        }
118        
119        results
120    }
121
122    pub fn find_symbol_in_path(&self, path: &str, name: &str) -> Option<NodeData> {
123        self.graph.graph.node_weights()
124            .find(|n| n.node_type == NodeType::Symbol && n.path == path && n.name == name)
125            .cloned()
126    }
127
128    pub fn find_files(&self, query: &str) -> Vec<NodeData> {
129        let query_lower = query.to_lowercase();
130        self.graph.graph.node_weights()
131            .filter(|n| {
132                n.node_type == NodeType::File && n.path.to_lowercase().contains(&query_lower)
133            })
134            .cloned()
135            .collect()
136    }
137
138    pub fn get_all_file_paths(&self) -> Vec<String> {
139        self.graph.graph.node_weights()
140            .filter(|n| n.node_type == NodeType::File)
141            .map(|n| n.path.clone())
142            .collect()
143    }
144}