llmy_codegraph/
model.rs

1//! Language-neutral code graph model shared by every extractor: containers
2//! (contracts / modules), callables (functions / modifiers / entries), state
3//! items (storage variables / accounts / resources / objects) and the edges
4//! between them. Call edges resolved purely from syntax keep their ambiguity
5//! explicit instead of pretending precision the parser does not have.
6
7use std::collections::{BTreeMap, BTreeSet};
8use std::path::{Path, PathBuf};
9
10use color_eyre::eyre::eyre;
11use llmy_types::error::LLMYError;
12use serde::{Deserialize, Serialize};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
15#[serde(rename_all = "snake_case")]
16pub enum Language {
17    Solidity,
18    Rust,
19    MoveAptos,
20    MoveSui,
21}
22
23impl Language {
24    pub fn render(&self) -> &'static str {
25        match self {
26            Self::Solidity => "Solidity",
27            Self::Rust => "Rust",
28            Self::MoveAptos => "Move (Aptos)",
29            Self::MoveSui => "Move (Sui)",
30        }
31    }
32}
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum ModuleKind {
37    Contract,
38    Interface,
39    Library,
40    /// A Rust or Move module.
41    Module,
42}
43
44impl ModuleKind {
45    pub fn render(&self) -> &'static str {
46        match self {
47            Self::Contract => "contract",
48            Self::Interface => "interface",
49            Self::Library => "library",
50            Self::Module => "module",
51        }
52    }
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
56#[serde(rename_all = "snake_case")]
57pub enum CallableKind {
58    Function,
59    Constructor,
60    Modifier,
61    Fallback,
62    Receive,
63    /// A Move `entry` function or an Anchor instruction handler — the
64    /// externally reachable surface.
65    Entry,
66}
67
68impl CallableKind {
69    pub fn render(&self) -> &'static str {
70        match self {
71            Self::Function => "function",
72            Self::Constructor => "constructor",
73            Self::Modifier => "modifier",
74            Self::Fallback => "fallback",
75            Self::Receive => "receive",
76            Self::Entry => "entry",
77        }
78    }
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
82#[serde(rename_all = "snake_case")]
83pub enum StateKind {
84    /// A Solidity storage variable.
85    StateVariable,
86    /// An Anchor `#[account]` data struct.
87    AnchorAccount,
88    /// A CosmWasm `Item` storage slot.
89    CwItem,
90    /// A CosmWasm `Map` storage collection.
91    CwMap,
92    /// An Aptos global resource (struct with `key`).
93    MoveResource,
94    /// A Sui object (struct with `key`).
95    SuiObject,
96}
97
98impl StateKind {
99    pub fn render(&self) -> &'static str {
100        match self {
101            Self::StateVariable => "state variable",
102            Self::AnchorAccount => "anchor account",
103            Self::CwItem => "cw item",
104            Self::CwMap => "cw map",
105            Self::MoveResource => "move resource",
106            Self::SuiObject => "sui object",
107        }
108    }
109}
110
111#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
112#[serde(rename_all = "snake_case")]
113pub enum AccessKind {
114    Read,
115    Write,
116}
117
118impl AccessKind {
119    pub fn render(&self) -> &'static str {
120        match self {
121            Self::Read => "read",
122            Self::Write => "write",
123        }
124    }
125}
126
127/// 1-based inclusive line span inside a source file.
128#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
129pub struct LineSpan {
130    pub start_line: usize,
131    pub end_line: usize,
132}
133
134#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct Module {
136    pub id: i64,
137    pub language: Language,
138    pub kind: ModuleKind,
139    pub name: String,
140    /// Path relative to the indexed root.
141    pub file: PathBuf,
142    pub span: LineSpan,
143}
144
145#[derive(Debug, Clone, Serialize, Deserialize)]
146pub struct Callable {
147    pub id: i64,
148    pub module_id: i64,
149    pub kind: CallableKind,
150    pub name: String,
151    /// Signature head as written in the source (cut before the body).
152    pub signature: String,
153    pub file: PathBuf,
154    pub span: LineSpan,
155}
156
157#[derive(Debug, Clone, Serialize, Deserialize)]
158pub struct StateItem {
159    pub id: i64,
160    pub module_id: i64,
161    pub kind: StateKind,
162    pub name: String,
163    /// Declared type, verbatim.
164    pub type_text: String,
165    pub file: PathBuf,
166    pub span: LineSpan,
167}
168
169/// A syntactically resolved callee. Name-based resolution over tree-sitter
170/// cannot always pick a single target, so ambiguity is first-class.
171#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
172#[serde(rename_all = "snake_case")]
173pub enum CalleeRef {
174    Resolved(i64),
175    Ambiguous(Vec<i64>),
176    /// Not defined inside the indexed project (external library, builtin, or
177    /// an interface without an in-project implementation).
178    External(String),
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize)]
182pub struct CallEdge {
183    pub caller_id: i64,
184    pub callee: CalleeRef,
185    /// The call site as written (e.g. `token.transfer`), for display.
186    pub callee_text: String,
187    pub line: usize,
188}
189
190#[derive(Debug, Clone, Serialize, Deserialize)]
191pub struct StateEdge {
192    pub callable_id: i64,
193    pub state_id: i64,
194    pub access: AccessKind,
195    pub line: usize,
196}
197
198#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
199#[serde(rename_all = "snake_case")]
200pub enum ParentRef {
201    Resolved(i64),
202    External(String),
203}
204
205#[derive(Debug, Clone, Serialize, Deserialize)]
206pub struct InheritEdge {
207    pub module_id: i64,
208    pub parent: ParentRef,
209}
210
211/// The complete index over one project root.
212#[derive(Debug, Clone, Default, Serialize, Deserialize)]
213pub struct CodeGraph {
214    pub modules: BTreeMap<i64, Module>,
215    pub callables: BTreeMap<i64, Callable>,
216    pub states: BTreeMap<i64, StateItem>,
217    pub call_edges: Vec<CallEdge>,
218    pub state_edges: Vec<StateEdge>,
219    pub inherit_edges: Vec<InheritEdge>,
220}
221
222impl CodeGraph {
223    pub fn is_empty(&self) -> bool {
224        self.modules.is_empty()
225    }
226
227    pub fn modules_by_name(&self, name: &str) -> Vec<&Module> {
228        self.modules.values().filter(|m| m.name == name).collect()
229    }
230
231    pub fn callables_of_module(&self, module_id: i64) -> Vec<&Callable> {
232        self.callables
233            .values()
234            .filter(|c| c.module_id == module_id)
235            .collect()
236    }
237
238    pub fn states_of_module(&self, module_id: i64) -> Vec<&StateItem> {
239        self.states
240            .values()
241            .filter(|s| s.module_id == module_id)
242            .collect()
243    }
244
245    /// Find callables by name, optionally scoped to a module name.
246    pub fn find_callables(&self, module: Option<&str>, name: &str) -> Vec<&Callable> {
247        self.callables
248            .values()
249            .filter(|c| c.name == name)
250            .filter(|c| match module {
251                Some(module) => self
252                    .modules
253                    .get(&c.module_id)
254                    .map(|m| m.name == module)
255                    .unwrap_or(false),
256                None => true,
257            })
258            .collect()
259    }
260
261    /// Find state items by name, optionally scoped to a module name.
262    pub fn find_states(&self, module: Option<&str>, name: &str) -> Vec<&StateItem> {
263        self.states
264            .values()
265            .filter(|s| s.name == name)
266            .filter(|s| match module {
267                Some(module) => self
268                    .modules
269                    .get(&s.module_id)
270                    .map(|m| m.name == module)
271                    .unwrap_or(false),
272                None => true,
273            })
274            .collect()
275    }
276
277    pub fn outgoing_calls(&self, caller_id: i64) -> Vec<&CallEdge> {
278        self.call_edges
279            .iter()
280            .filter(|e| e.caller_id == caller_id)
281            .collect()
282    }
283
284    pub fn incoming_calls(&self, callee_id: i64) -> Vec<&CallEdge> {
285        self.call_edges
286            .iter()
287            .filter(|e| match &e.callee {
288                CalleeRef::Resolved(id) => *id == callee_id,
289                CalleeRef::Ambiguous(ids) => ids.contains(&callee_id),
290                CalleeRef::External(_) => false,
291            })
292            .collect()
293    }
294
295    pub fn state_accesses_of(&self, callable_id: i64) -> Vec<&StateEdge> {
296        self.state_edges
297            .iter()
298            .filter(|e| e.callable_id == callable_id)
299            .collect()
300    }
301
302    pub fn accessors_of_state(&self, state_id: i64) -> Vec<&StateEdge> {
303        self.state_edges
304            .iter()
305            .filter(|e| e.state_id == state_id)
306            .collect()
307    }
308
309    /// Transitive ancestors of a module through inheritance edges,
310    /// cycle-safe.
311    pub fn ancestors_of(&self, module_id: i64) -> BTreeSet<i64> {
312        let mut out = BTreeSet::new();
313        let mut frontier = vec![module_id];
314        while let Some(current) = frontier.pop() {
315            for edge in self.inherit_edges.iter().filter(|e| e.module_id == current) {
316                if let ParentRef::Resolved(parent) = edge.parent
317                    && out.insert(parent)
318                {
319                    frontier.push(parent);
320                }
321            }
322        }
323        out
324    }
325
326    /// Transitive descendants of a module through inheritance edges,
327    /// cycle-safe.
328    pub fn descendants_of(&self, module_id: i64) -> BTreeSet<i64> {
329        let mut out = BTreeSet::new();
330        let mut frontier = vec![module_id];
331        while let Some(current) = frontier.pop() {
332            for edge in self.inherit_edges.iter() {
333                if edge.parent == ParentRef::Resolved(current) && out.insert(edge.module_id) {
334                    frontier.push(edge.module_id);
335                }
336            }
337        }
338        out
339    }
340
341    /// State items visible to a module: its own plus everything declared by
342    /// its ancestors.
343    pub fn visible_states(&self, module_id: i64) -> Vec<&StateItem> {
344        let mut scope = self.ancestors_of(module_id);
345        scope.insert(module_id);
346        self.states
347            .values()
348            .filter(|s| scope.contains(&s.module_id))
349            .collect()
350    }
351
352    pub fn render_callee(&self, callee: &CalleeRef) -> String {
353        match callee {
354            CalleeRef::Resolved(id) => self.render_callable_ref(*id),
355            CalleeRef::Ambiguous(ids) => {
356                let rendered = ids
357                    .iter()
358                    .map(|id| self.render_callable_ref(*id))
359                    .collect::<Vec<_>>()
360                    .join(" | ");
361                format!("ambiguous({rendered})")
362            }
363            CalleeRef::External(name) => format!("external({name})"),
364        }
365    }
366
367    pub fn render_callable_ref(&self, callable_id: i64) -> String {
368        match self.callables.get(&callable_id) {
369            Some(callable) => {
370                let module = self
371                    .modules
372                    .get(&callable.module_id)
373                    .map(|m| m.name.as_str())
374                    .unwrap_or("?");
375                format!("{}.{}", module, callable.name)
376            }
377            None => format!("#{callable_id}"),
378        }
379    }
380
381    /// Read the source lines of a span from disk, relative to `root`.
382    pub async fn read_span(
383        &self,
384        root: &Path,
385        file: &Path,
386        span: LineSpan,
387    ) -> Result<String, LLMYError> {
388        let content = tokio::fs::read_to_string(root.join(file)).await?;
389        let start = span.start_line.saturating_sub(1);
390        let count = span.end_line.saturating_sub(start);
391        let selected = content
392            .lines()
393            .skip(start)
394            .take(count)
395            .collect::<Vec<_>>()
396            .join("\n");
397        if selected.is_empty() {
398            return Err(eyre!(
399                "span {}..{} of {} is empty or out of range",
400                span.start_line,
401                span.end_line,
402                file.display()
403            )
404            .into());
405        }
406        Ok(selected)
407    }
408
409    pub fn counts(&self) -> String {
410        format!(
411            "{} modules, {} callables, {} state items, {} call edges, {} state edges, {} inheritance edges",
412            self.modules.len(),
413            self.callables.len(),
414            self.states.len(),
415            self.call_edges.len(),
416            self.state_edges.len(),
417            self.inherit_edges.len()
418        )
419    }
420}
421
422#[cfg(test)]
423mod tests {
424    use super::*;
425
426    fn graph_with_inheritance() -> CodeGraph {
427        let mut graph = CodeGraph::default();
428        for (id, name) in [(1, "Base"), (2, "Mid"), (3, "Leaf")] {
429            graph.modules.insert(
430                id,
431                Module {
432                    id,
433                    language: Language::Solidity,
434                    kind: ModuleKind::Contract,
435                    name: name.to_string(),
436                    file: PathBuf::from("a.sol"),
437                    span: LineSpan {
438                        start_line: 1,
439                        end_line: 10,
440                    },
441                },
442            );
443        }
444        graph.inherit_edges.push(InheritEdge {
445            module_id: 2,
446            parent: ParentRef::Resolved(1),
447        });
448        graph.inherit_edges.push(InheritEdge {
449            module_id: 3,
450            parent: ParentRef::Resolved(2),
451        });
452        graph
453    }
454
455    #[test]
456    fn ancestors_and_descendants_are_transitive() {
457        let graph = graph_with_inheritance();
458        assert_eq!(graph.ancestors_of(3), BTreeSet::from([1, 2]));
459        assert_eq!(graph.descendants_of(1), BTreeSet::from([2, 3]));
460        assert!(graph.ancestors_of(1).is_empty());
461    }
462
463    #[test]
464    fn visible_states_include_inherited_ones() {
465        let mut graph = graph_with_inheritance();
466        graph.states.insert(
467            1,
468            StateItem {
469                id: 1,
470                module_id: 1,
471                kind: StateKind::StateVariable,
472                name: "owner".to_string(),
473                type_text: "address".to_string(),
474                file: PathBuf::from("a.sol"),
475                span: LineSpan {
476                    start_line: 2,
477                    end_line: 2,
478                },
479            },
480        );
481        let visible = graph.visible_states(3);
482        assert_eq!(visible.len(), 1);
483        assert_eq!(visible[0].name, "owner");
484        assert!(graph.visible_states(1).len() == 1);
485    }
486
487    #[test]
488    fn incoming_calls_match_resolved_and_ambiguous() {
489        let mut graph = graph_with_inheritance();
490        for (id, name) in [(10, "a"), (11, "b"), (12, "c")] {
491            graph.callables.insert(
492                id,
493                Callable {
494                    id,
495                    module_id: 1,
496                    kind: CallableKind::Function,
497                    name: name.to_string(),
498                    signature: format!("function {name}()"),
499                    file: PathBuf::from("a.sol"),
500                    span: LineSpan {
501                        start_line: 3,
502                        end_line: 5,
503                    },
504                },
505            );
506        }
507        graph.call_edges.push(CallEdge {
508            caller_id: 10,
509            callee: CalleeRef::Resolved(11),
510            callee_text: "b".to_string(),
511            line: 4,
512        });
513        graph.call_edges.push(CallEdge {
514            caller_id: 12,
515            callee: CalleeRef::Ambiguous(vec![10, 11]),
516            callee_text: "ab".to_string(),
517            line: 4,
518        });
519
520        assert_eq!(graph.incoming_calls(11).len(), 2);
521        assert_eq!(graph.incoming_calls(10).len(), 1);
522        assert_eq!(graph.outgoing_calls(10).len(), 1);
523    }
524}
llmy_codegraph/model.rs

llmy_codegraph/
model.rs
