Skip to main content

cgx_engine/parsers/
rust.rs

1use tree_sitter::{Parser, Query, QueryCursor};
2
3use crate::parser::{EdgeDef, EdgeKind, LanguageParser, NodeDef, NodeKind, ParseResult};
4use crate::walker::SourceFile;
5
6pub struct RustParser {
7    language: tree_sitter::Language,
8}
9
10impl RustParser {
11    pub fn new() -> Self {
12        Self {
13            language: tree_sitter_rust::language(),
14        }
15    }
16}
17
18impl Default for RustParser {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl LanguageParser for RustParser {
25    fn extensions(&self) -> &[&str] {
26        &["rs"]
27    }
28
29    fn extract(&self, file: &SourceFile) -> anyhow::Result<ParseResult> {
30        let mut parser = Parser::new();
31        parser.set_language(&self.language)?;
32
33        let tree = parser.parse(&file.content, None).ok_or_else(|| {
34            anyhow::anyhow!("failed to parse {}", file.relative_path)
35        })?;
36
37        let source_bytes = file.content.as_bytes();
38        let root = tree.root_node();
39        let mut nodes = Vec::new();
40        let mut edges = Vec::new();
41
42        let fp = format!("file:{}", file.relative_path);
43
44        // Function definitions
45        if let Ok(query) = Query::new(&self.language, "(function_item name: (identifier) @name) @fn") {
46            let mut cursor = QueryCursor::new();
47            for m in cursor.matches(&query, root, source_bytes) {
48                let Some(name_capture) = m
49                    .captures
50                    .iter()
51                    .find(|c| query.capture_names()[c.index as usize] == "name")
52                else {
53                    continue;
54                };
55                let name = node_text(name_capture.node, source_bytes);
56                let start = name_capture.node.start_position();
57                let body_end = m.captures.iter()
58                    .find(|c| query.capture_names()[c.index as usize] == "fn")
59                    .map(|c| c.node.end_position())
60                    .unwrap_or_else(|| name_capture.node.end_position());
61                let id = format!("fn:{}:{}", file.relative_path, name);
62
63                nodes.push(NodeDef {
64                    id: id.clone(),
65                    kind: NodeKind::Function,
66                    name,
67                    path: file.relative_path.clone(),
68                    line_start: start.row as u32 + 1,
69                    line_end: body_end.row as u32 + 1,
70                    ..Default::default()
71                });
72
73                edges.push(EdgeDef {
74                    src: fp.clone(),
75                    dst: id,
76                    kind: EdgeKind::Exports,
77                    ..Default::default()
78                });
79            }
80        }
81
82        // Struct definitions
83        if let Ok(query) = Query::new(&self.language, "(struct_item name: (type_identifier) @name) @s") {
84            extract_type_nodes(
85                &mut nodes,
86                &mut edges,
87                &fp,
88                file,
89                &query,
90                root,
91                source_bytes,
92                NodeKind::Class,
93                "cls",
94            );
95        }
96
97        // Enum definitions
98        if let Ok(query) = Query::new(&self.language, "(enum_item name: (type_identifier) @name) @e") {
99            extract_type_nodes(
100                &mut nodes,
101                &mut edges,
102                &fp,
103                file,
104                &query,
105                root,
106                source_bytes,
107                NodeKind::Class,
108                "cls",
109            );
110        }
111
112        // Trait definitions
113        if let Ok(query) = Query::new(&self.language, "(trait_item name: (type_identifier) @name) @t") {
114            extract_type_nodes(
115                &mut nodes,
116                &mut edges,
117                &fp,
118                file,
119                &query,
120                root,
121                source_bytes,
122                NodeKind::Class,
123                "cls",
124            );
125        }
126
127        // Impl blocks — add edges for impl'd struct/trait methods
128        if let Ok(query) = Query::new(
129            &self.language,
130            "(impl_item type: (type_identifier) @type body: (_) @body)",
131        ) {
132            let mut cursor = QueryCursor::new();
133            for m in cursor.matches(&query, root, source_bytes) {
134                if let Some(type_cap) = m
135                    .captures
136                    .iter()
137                    .find(|c| query.capture_names()[c.index as usize] == "type")
138                {
139                    let type_name = node_text(type_cap.node, source_bytes);
140                    edges.push(EdgeDef {
141                        src: fp.clone(),
142                        dst: format!("cls:{}:{}", file.relative_path, type_name),
143                        kind: EdgeKind::Exports,
144                        ..Default::default()
145                    });
146                }
147            }
148        }
149
150        // Use statements
151        if let Ok(query) = Query::new(
152            &self.language,
153            "(use_declaration argument: (scoped_identifier path: (_) @path name: (_)?))",
154        ) {
155            let mut cursor = QueryCursor::new();
156            for m in cursor.matches(&query, root, source_bytes) {
157                if let Some(path_cap) = m
158                    .captures
159                    .iter()
160                    .find(|c| query.capture_names()[c.index as usize] == "path")
161                {
162                    let full_path = node_text(path_cap.node, source_bytes);
163                    // Simple case: use crate::foo::bar -> file path is src/foo/bar.rs
164                    let import_path = if full_path.starts_with("crate::") {
165                        format!(
166                            "src/{}.rs",
167                            full_path.trim_start_matches("crate::").replace("::", "/")
168                        )
169                    } else {
170                        continue;
171                    };
172                    edges.push(EdgeDef {
173                        src: fp.clone(),
174                        dst: format!("file:{}", import_path),
175                        kind: EdgeKind::Imports,
176                        ..Default::default()
177                    });
178                }
179            }
180        }
181
182        // Simpler use declarations (use foo::Bar)
183        if let Ok(query) = Query::new(&self.language, "(use_declaration argument: (identifier) @name)") {
184            let mut cursor = QueryCursor::new();
185            for m in cursor.matches(&query, root, source_bytes) {
186                if let Some(name_cap) = m
187                    .captures
188                    .iter()
189                    .find(|c| query.capture_names()[c.index as usize] == "name")
190                {
191                    let mod_name = node_text(name_cap.node, source_bytes);
192                    let import_path = mod_name;
193                    edges.push(EdgeDef {
194                        src: fp.clone(),
195                        dst: format!("file:{}.rs", import_path),
196                        kind: EdgeKind::Imports,
197                        ..Default::default()
198                    });
199                }
200            }
201        }
202
203        Ok(ParseResult { nodes, edges })
204    }
205}
206
207#[allow(clippy::too_many_arguments)]
208fn extract_type_nodes(
209    nodes: &mut Vec<NodeDef>,
210    edges: &mut Vec<EdgeDef>,
211    file_id: &str,
212    file: &SourceFile,
213    query: &Query,
214    root: tree_sitter::Node,
215    source_bytes: &[u8],
216    kind: NodeKind,
217    prefix: &str,
218) {
219    let mut cursor = QueryCursor::new();
220    for m in cursor.matches(query, root, source_bytes) {
221        let Some(name_capture) = m
222            .captures
223            .iter()
224            .find(|c| query.capture_names()[c.index as usize] == "name")
225        else {
226            continue;
227        };
228        let name = node_text(name_capture.node, source_bytes);
229        let start = name_capture.node.start_position();
230        // Use the body node for end position; fall back to name node if no body capture
231        let body_end = m.captures.iter()
232            .find(|c| query.capture_names()[c.index as usize] != "name")
233            .map(|c| c.node.end_position())
234            .unwrap_or_else(|| name_capture.node.end_position());
235        let id = format!("{}:{}:{}", prefix, file.relative_path, name);
236
237        nodes.push(NodeDef {
238            id: id.clone(),
239            kind: kind.clone(),
240            name,
241            path: file.relative_path.clone(),
242            line_start: start.row as u32 + 1,
243            line_end: body_end.row as u32 + 1,
244            ..Default::default()
245        });
246
247        edges.push(EdgeDef {
248            src: file_id.to_string(),
249            dst: id,
250            kind: EdgeKind::Exports,
251            ..Default::default()
252        });
253    }
254}
255
256fn node_text(node: tree_sitter::Node, source: &[u8]) -> String {
257    node.utf8_text(source).unwrap_or("").to_string()
258}