llmy_codegraph/
rust_lang.rs

1//! Rust contract extraction over tree-sitter-rust. One module per file.
2//! State is recognized for the two supported contract ecosystems:
3//! Anchor (`#[account]` structs, expanded through `Context<T>` /
4//! `#[derive(Accounts)]` containers) and CosmWasm (`Item` / `Map` storage
5//! declarations with their `save`/`load`-style accesses). Plain Rust still
6//! gets modules, functions and call edges.
7
8use std::collections::BTreeMap;
9
10use llmy_types::error::LLMYError;
11use tree_sitter::Node;
12
13use crate::extract::{
14    FileExtraction, GrammarSet, NodeUtil, RawCallSite, RawCallable, RawModule, RawState,
15    RawStateRef, SourceFile,
16};
17use crate::model::{CallableKind, Language, ModuleKind, StateKind};
18
19const CW_WRITE_METHODS: [&str; 4] = ["save", "update", "remove", "replace"];
20const CW_READ_METHODS: [&str; 7] = [
21    "load", "may_load", "has", "query", "range", "keys", "prefix",
22];
23
24/// One field of an Anchor `#[derive(Accounts)]` container: the inner account
25/// type plus whether the field is `#[account(mut)]`.
26#[derive(Debug, Clone)]
27struct AccountsField {
28    type_name: String,
29    mutable: bool,
30}
31
32pub struct RustExtractor;
33
34impl RustExtractor {
35    pub fn extract(file: &SourceFile) -> Result<FileExtraction, LLMYError> {
36        let tree = GrammarSet::parse(Language::Rust, &file.content)?;
37        let root = tree.root_node();
38        let parse_errors = GrammarSet::count_errors(root);
39        let source = &file.content;
40
41        let module_name = file
42            .relative
43            .with_extension("")
44            .display()
45            .to_string()
46            .replace('\\', "/");
47
48        let mut states = vec![];
49        // Accounts containers: name -> fields, used to expand Context<T>.
50        let mut containers: BTreeMap<String, Vec<AccountsField>> = BTreeMap::new();
51
52        for item in root.descendants_of_kinds(&["struct_item", "const_item", "static_item"], false)
53        {
54            match item.kind() {
55                "struct_item" => {
56                    let Some(name) = item
57                        .child_of_kind("type_identifier")
58                        .map(|n| n.text_of(source))
59                    else {
60                        continue;
61                    };
62                    let attributes = Self::preceding_attributes(item, source);
63                    if attributes.iter().any(|a| a == "account") {
64                        states.push(RawState {
65                            name,
66                            kind: StateKind::AnchorAccount,
67                            type_text: "#[account] struct".to_string(),
68                            span: item.line_span(),
69                        });
70                    } else if attributes.iter().any(|a| a.contains("Accounts")) {
71                        containers.insert(name, Self::container_fields(item, source));
72                    }
73                }
74                _ => {
75                    let Some(name) = item.child_of_kind("identifier").map(|n| n.text_of(source))
76                    else {
77                        continue;
78                    };
79                    let Some(type_node) = item.child_of_kind("generic_type") else {
80                        continue;
81                    };
82                    let head = type_node
83                        .child_of_kind("type_identifier")
84                        .map(|n| n.text_of(source))
85                        .unwrap_or_default();
86                    let kind = match head.as_str() {
87                        "Item" => StateKind::CwItem,
88                        "Map" => StateKind::CwMap,
89                        _ => continue,
90                    };
91                    states.push(RawState {
92                        name,
93                        kind,
94                        type_text: type_node.text_of(source),
95                        span: item.line_span(),
96                    });
97                }
98            }
99        }
100
101        // Modules under a #[program] attribute mark Anchor instruction
102        // handlers.
103        let program_mods: Vec<Node<'_>> = root
104            .descendants_of_kinds(&["mod_item"], true)
105            .into_iter()
106            .filter(|m| {
107                Self::preceding_attributes(*m, source)
108                    .iter()
109                    .any(|a| a == "program")
110            })
111            .collect();
112
113        let cw_state_names: Vec<String> = states
114            .iter()
115            .filter(|s| matches!(s.kind, StateKind::CwItem | StateKind::CwMap))
116            .map(|s| s.name.clone())
117            .collect();
118
119        let mut callables = vec![];
120        for function in root.descendants_of_kinds(&["function_item"], false) {
121            callables.push(Self::extract_function(
122                function,
123                source,
124                &program_mods,
125                &containers,
126                &cw_state_names,
127            ));
128        }
129
130        Ok(FileExtraction {
131            file: file.relative.clone(),
132            language: Language::Rust,
133            modules: vec![RawModule {
134                name: module_name,
135                kind: ModuleKind::Module,
136                span: root.line_span(),
137                parents: vec![],
138                callables,
139                states,
140            }],
141            parse_errors,
142        })
143    }
144
145    /// The attribute names written directly above an item (`#[account]` ->
146    /// "account", `#[derive(Accounts)]` -> "derive(Accounts)").
147    fn preceding_attributes(item: Node<'_>, source: &str) -> Vec<String> {
148        let mut out = vec![];
149        let mut current = item;
150        while let Some(previous) = current.prev_named_sibling() {
151            if previous.kind() != "attribute_item" {
152                break;
153            }
154            if let Some(attribute) = previous.child_of_kind("attribute") {
155                out.push(attribute.text_of(source));
156            }
157            current = previous;
158        }
159        out
160    }
161
162    fn container_fields(item: Node<'_>, source: &str) -> Vec<AccountsField> {
163        let Some(fields) = item.child_of_kind("field_declaration_list") else {
164            return vec![];
165        };
166        let mut out = vec![];
167        for field in fields.children_of_kind("field_declaration") {
168            let mutable = Self::preceding_attributes(field, source)
169                .iter()
170                .any(|a| a.starts_with("account") && a.contains("mut"));
171            // The inner account type is the last type_identifier of the
172            // field's type (e.g. `Account<'info, Counter>` -> `Counter`).
173            let type_names = field.descendants_of_kinds(&["type_identifier"], true);
174            if let Some(inner) = type_names.last() {
175                out.push(AccountsField {
176                    type_name: inner.text_of(source),
177                    mutable,
178                });
179            }
180        }
181        out
182    }
183
184    fn extract_function(
185        function: Node<'_>,
186        source: &str,
187        program_mods: &[Node<'_>],
188        containers: &BTreeMap<String, Vec<AccountsField>>,
189        cw_state_names: &[String],
190    ) -> RawCallable {
191        let name = function
192            .child_of_kind("identifier")
193            .map(|n| n.text_of(source))
194            .unwrap_or_else(|| "<function>".to_string());
195        let signature = function.signature_head(&["block"], source);
196
197        let attributes = Self::preceding_attributes(function, source);
198        let in_program_mod = program_mods.iter().any(|m| {
199            m.start_byte() <= function.start_byte() && function.end_byte() <= m.end_byte()
200        });
201        let kind = if in_program_mod || attributes.iter().any(|a| a.contains("entry_point")) {
202            CallableKind::Entry
203        } else {
204            CallableKind::Function
205        };
206
207        let mut calls = vec![];
208        let mut state_refs = vec![];
209
210        // Context<T> parameters expand through the Accounts container into
211        // per-account state references.
212        if let Some(parameters) = function.child_of_kind("parameters") {
213            for generic in parameters.descendants_of_kinds(&["generic_type"], true) {
214                let head = generic
215                    .child_of_kind("type_identifier")
216                    .map(|n| n.text_of(source))
217                    .unwrap_or_default();
218                if head != "Context" {
219                    continue;
220                }
221                let Some(arguments) = generic.child_of_kind("type_arguments") else {
222                    continue;
223                };
224                let Some(container_name) = arguments
225                    .descendants_of_kinds(&["type_identifier"], false)
226                    .first()
227                    .map(|n| n.text_of(source))
228                else {
229                    continue;
230                };
231                for field in containers.get(&container_name).into_iter().flatten() {
232                    state_refs.push(RawStateRef {
233                        name: field.type_name.clone(),
234                        write: field.mutable,
235                        line: function.line_span().start_line,
236                    });
237                }
238            }
239        }
240
241        if let Some(body) = function.child_of_kind("block") {
242            for call in body.descendants_of_kinds(&["call_expression"], true) {
243                let Some(target) = call.named_child(0) else {
244                    continue;
245                };
246                let line = call.line_span().start_line;
247                match target.kind() {
248                    "identifier" => {
249                        let call_name = target.text_of(source);
250                        calls.push(RawCallSite {
251                            text: call_name.clone(),
252                            name: call_name,
253                            qualifier: None,
254                            line,
255                        });
256                    }
257                    "scoped_identifier" => {
258                        let mut cursor = target.walk();
259                        let parts: Vec<String> = target
260                            .children(&mut cursor)
261                            .filter(|c| c.is_named())
262                            .map(|c| c.text_of(source))
263                            .collect();
264                        let Some(call_name) = parts.last().cloned() else {
265                            continue;
266                        };
267                        let qualifier = (parts.len() > 1).then(|| parts[0].clone());
268                        calls.push(RawCallSite {
269                            text: target.text_of(source),
270                            name: call_name,
271                            qualifier,
272                            line,
273                        });
274                    }
275                    "field_expression" => {
276                        let Some(method) = target
277                            .child_of_kind("field_identifier")
278                            .map(|n| n.text_of(source))
279                        else {
280                            continue;
281                        };
282                        let receiver = target
283                            .named_child(0)
284                            .map(|n| n.text_of(source))
285                            .unwrap_or_default();
286
287                        // CosmWasm storage access: STATE.save(...) etc.
288                        if cw_state_names.contains(&receiver) {
289                            if CW_WRITE_METHODS.contains(&method.as_str()) {
290                                state_refs.push(RawStateRef {
291                                    name: receiver.clone(),
292                                    write: true,
293                                    line,
294                                });
295                                continue;
296                            }
297                            if CW_READ_METHODS.contains(&method.as_str()) {
298                                state_refs.push(RawStateRef {
299                                    name: receiver.clone(),
300                                    write: false,
301                                    line,
302                                });
303                                continue;
304                            }
305                        }
306
307                        calls.push(RawCallSite {
308                            text: target.text_of(source),
309                            name: method,
310                            qualifier: Some(receiver),
311                            line,
312                        });
313                    }
314                    _ => {}
315                }
316            }
317        }
318
319        RawCallable {
320            name,
321            kind,
322            signature,
323            span: function.line_span(),
324            calls,
325            state_refs,
326        }
327    }
328}
329
330#[cfg(test)]
331mod tests {
332    use super::*;
333    use std::path::PathBuf;
334
335    fn extract(source: &str) -> FileExtraction {
336        RustExtractor::extract(&SourceFile {
337            relative: PathBuf::from("src/lib.rs"),
338            content: source.to_string(),
339        })
340        .expect("extract")
341    }
342
343    const COSMWASM: &str = r#"
344use cw_storage_plus::{Item, Map};
345const STATE: Item<State> = Item::new("state");
346const BALANCES: Map<&Addr, u128> = Map::new("balances");
347
348#[entry_point]
349pub fn execute(deps: DepsMut) -> Result<Response, ContractError> {
350    let mut state = STATE.load(deps.storage)?;
351    state.count += 1;
352    STATE.save(deps.storage, &state)?;
353    helper(&state);
354    Ok(Response::new())
355}
356
357fn helper(state: &State) -> u64 { state.count }
358"#;
359
360    const ANCHOR: &str = r#"
361#[program]
362pub mod counter {
363    use super::*;
364    pub fn increment(ctx: Context<Increment>) -> Result<()> {
365        ctx.accounts.counter.count += 1;
366        Ok(())
367    }
368}
369
370#[account]
371pub struct Counter { pub count: u64 }
372
373#[derive(Accounts)]
374pub struct Increment<'info> {
375    #[account(mut)]
376    pub counter: Account<'info, Counter>,
377    pub user: Signer<'info>,
378}
379"#;
380
381    #[test]
382    fn cosmwasm_state_and_accesses_are_extracted() {
383        let extraction = extract(COSMWASM);
384        assert_eq!(extraction.parse_errors, 0);
385        let module = &extraction.modules[0];
386        assert_eq!(module.name, "src/lib");
387
388        let kinds: Vec<_> = module
389            .states
390            .iter()
391            .map(|s| (s.name.as_str(), s.kind))
392            .collect();
393        assert_eq!(
394            kinds,
395            vec![("STATE", StateKind::CwItem), ("BALANCES", StateKind::CwMap)]
396        );
397
398        let execute = &module.callables[0];
399        assert_eq!(execute.kind, CallableKind::Entry);
400        let refs: Vec<_> = execute
401            .state_refs
402            .iter()
403            .map(|r| (r.name.as_str(), r.write))
404            .collect();
405        assert!(refs.contains(&("STATE", false)), "{refs:?}");
406        assert!(refs.contains(&("STATE", true)), "{refs:?}");
407        assert!(execute.calls.iter().any(|c| c.name == "helper"));
408        // Storage method calls became state refs, not call edges.
409        assert!(!execute.calls.iter().any(|c| c.name == "save"));
410    }
411
412    #[test]
413    fn anchor_accounts_expand_through_context() {
414        let extraction = extract(ANCHOR);
415        assert_eq!(extraction.parse_errors, 0);
416        let module = &extraction.modules[0];
417
418        assert_eq!(module.states.len(), 1);
419        assert_eq!(module.states[0].name, "Counter");
420        assert_eq!(module.states[0].kind, StateKind::AnchorAccount);
421
422        let increment = &module.callables[0];
423        assert_eq!(increment.name, "increment");
424        assert_eq!(increment.kind, CallableKind::Entry);
425        let refs: Vec<_> = increment
426            .state_refs
427            .iter()
428            .map(|r| (r.name.as_str(), r.write))
429            .collect();
430        assert!(refs.contains(&("Counter", true)), "{refs:?}");
431        // Signer is also expanded but resolves to no state item later.
432        assert!(refs.contains(&("Signer", false)), "{refs:?}");
433    }
434}
llmy_codegraph/rust_lang.rs

llmy_codegraph/
rust_lang.rs
