lean_ctx/core/graph_expand/
partial.rs

1//! Partial call-graph representation for bounded neighborhood expansion.
2//!
3//! Stores nodes, directed edges, and hop depth from a center symbol for
4//! compact LLM context injection.
5
6use std::collections::HashMap;
7
8/// A partial subgraph centered on a target symbol.
9#[derive(Debug, Clone, Default)]
10pub struct PartialGraph {
11    /// Nodes in the subgraph, keyed by symbol name.
12    pub nodes: HashMap<String, NodeInfo>,
13    /// Directed edges represented as source, target, and relation type.
14    pub edges: Vec<(String, String, EdgeKind)>,
15    /// The center symbol this graph was expanded from.
16    pub center: String,
17    /// Maximum depth requested during expansion.
18    pub max_depth: usize,
19}
20
21/// Metadata for a symbol included in a partial graph.
22#[derive(Debug, Clone)]
23pub struct NodeInfo {
24    /// File containing the symbol.
25    pub file: String,
26    /// Symbol kind, such as function, method, or type.
27    pub kind: String,
28    /// Shortest hop distance from the center symbol.
29    pub depth: usize,
30}
31
32/// Relationship represented by a directed graph edge.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
34pub enum EdgeKind {
35    /// The source symbol calls the target symbol.
36    Calls,
37    /// The source symbol is called by the target symbol.
38    CalledBy,
39    /// The source symbol imports the target symbol.
40    Imports,
41    /// The source symbol implements the target symbol.
42    Implements,
43}
44
45impl PartialGraph {
46    /// Render the subgraph as deterministic compact text for LLM context injection.
47    pub fn render(&self) -> String {
48        let Some(center) = self.nodes.get(&self.center) else {
49            return String::new();
50        };
51
52        let mut output = format!("Center: {} ({})", self.center, center.file);
53        for depth in 1..=self.max_depth {
54            let mut nodes: Vec<(&str, &NodeInfo)> = self
55                .nodes
56                .iter()
57                .filter(|(_, info)| info.depth == depth)
58                .map(|(name, info)| (name.as_str(), info))
59                .collect();
60            nodes.sort_unstable_by_key(|(name, _)| *name);
61
62            if nodes.is_empty() {
63                continue;
64            }
65
66            output.push_str(&format!("\nDepth {depth}: "));
67            let rendered = nodes
68                .into_iter()
69                .map(|(name, info)| {
70                    let direction = self.direction_for(name, depth);
71                    format!("{direction}{name} ({})", info.file)
72                })
73                .collect::<Vec<_>>()
74                .join(", ");
75            output.push_str(&rendered);
76        }
77        output
78    }
79
80    /// Return the count of nodes in the subgraph.
81    pub fn node_count(&self) -> usize {
82        self.nodes.len()
83    }
84
85    /// Return symbols at a specific depth in deterministic name order.
86    pub fn at_depth(&self, depth: usize) -> Vec<&str> {
87        let mut symbols: Vec<&str> = self
88            .nodes
89            .iter()
90            .filter(|(_, info)| info.depth == depth)
91            .map(|(name, _)| name.as_str())
92            .collect();
93        symbols.sort_unstable();
94        symbols
95    }
96
97    fn direction_for(&self, name: &str, depth: usize) -> String {
98        let relation = self.edges.iter().find_map(|(from, to, relation)| {
99            if to == name {
100                self.nodes
101                    .get(from)
102                    .filter(|info| info.depth + 1 == depth)
103                    .map(|_| *relation)
104            } else {
105                None
106            }
107        });
108        let arrow = match relation {
109            Some(EdgeKind::CalledBy) => "← ",
110            Some(EdgeKind::Calls | EdgeKind::Imports | EdgeKind::Implements) | None => "→ ",
111        };
112        arrow.repeat(depth)
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::{NodeInfo, PartialGraph};
119
120    #[test]
121    fn empty_graph_renders_empty() {
122        assert_eq!(PartialGraph::default().render(), "");
123    }
124
125    #[test]
126    fn single_node_graph() {
127        let mut graph = PartialGraph {
128            center: "root".to_string(),
129            max_depth: 2,
130            ..PartialGraph::default()
131        };
132        graph.nodes.insert(
133            "root".to_string(),
134            NodeInfo {
135                file: "src/root.rs".to_string(),
136                kind: "function".to_string(),
137                depth: 0,
138            },
139        );
140
141        assert_eq!(graph.node_count(), 1);
142        assert_eq!(graph.render(), "Center: root (src/root.rs)");
143    }
144
145    #[test]
146    fn at_depth_filters_correctly() {
147        let mut graph = PartialGraph::default();
148        for (name, depth) in [("root", 0), ("beta", 1), ("alpha", 1), ("leaf", 2)] {
149            graph.nodes.insert(
150                name.to_string(),
151                NodeInfo {
152                    file: format!("{name}.rs"),
153                    kind: "function".to_string(),
154                    depth,
155                },
156            );
157        }
158
159        assert_eq!(graph.at_depth(1), vec!["alpha", "beta"]);
160        assert_eq!(graph.at_depth(2), vec!["leaf"]);
161        assert!(graph.at_depth(3).is_empty());
162    }
163}
lean_ctx/core/graph_expand/partial.rs

lean_ctx/core/graph_expand/
partial.rs
