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 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 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 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 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 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 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 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 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 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}