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
34 .parse(&file.content, None)
35 .ok_or_else(|| anyhow::anyhow!("failed to parse {}", file.relative_path))?;
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(
46 &self.language,
47 "(function_item name: (identifier) @name) @fn",
48 ) {
49 let mut cursor = QueryCursor::new();
50 for m in cursor.matches(&query, root, source_bytes) {
51 let Some(name_capture) = m
52 .captures
53 .iter()
54 .find(|c| query.capture_names()[c.index as usize] == "name")
55 else {
56 continue;
57 };
58 let name = node_text(name_capture.node, source_bytes);
59 let start = name_capture.node.start_position();
60 let body_end = m
61 .captures
62 .iter()
63 .find(|c| query.capture_names()[c.index as usize] == "fn")
64 .map(|c| c.node.end_position())
65 .unwrap_or_else(|| name_capture.node.end_position());
66 let id = format!("fn:{}:{}", file.relative_path, name);
67
68 nodes.push(NodeDef {
69 id: id.clone(),
70 kind: NodeKind::Function,
71 name,
72 path: file.relative_path.clone(),
73 line_start: start.row as u32 + 1,
74 line_end: body_end.row as u32 + 1,
75 ..Default::default()
76 });
77
78 edges.push(EdgeDef {
79 src: fp.clone(),
80 dst: id,
81 kind: EdgeKind::Exports,
82 ..Default::default()
83 });
84 }
85 }
86
87 if let Ok(query) = Query::new(
89 &self.language,
90 "(struct_item name: (type_identifier) @name) @s",
91 ) {
92 extract_type_nodes(
93 &mut nodes,
94 &mut edges,
95 &fp,
96 file,
97 &query,
98 root,
99 source_bytes,
100 NodeKind::Class,
101 "cls",
102 );
103 }
104
105 if let Ok(query) = Query::new(
107 &self.language,
108 "(enum_item name: (type_identifier) @name) @e",
109 ) {
110 extract_type_nodes(
111 &mut nodes,
112 &mut edges,
113 &fp,
114 file,
115 &query,
116 root,
117 source_bytes,
118 NodeKind::Class,
119 "cls",
120 );
121 }
122
123 if let Ok(query) = Query::new(
125 &self.language,
126 "(trait_item name: (type_identifier) @name) @t",
127 ) {
128 extract_type_nodes(
129 &mut nodes,
130 &mut edges,
131 &fp,
132 file,
133 &query,
134 root,
135 source_bytes,
136 NodeKind::Class,
137 "cls",
138 );
139 }
140
141 if let Ok(query) = Query::new(
143 &self.language,
144 "(impl_item type: (type_identifier) @type body: (_) @body)",
145 ) {
146 let mut cursor = QueryCursor::new();
147 for m in cursor.matches(&query, root, source_bytes) {
148 if let Some(type_cap) = m
149 .captures
150 .iter()
151 .find(|c| query.capture_names()[c.index as usize] == "type")
152 {
153 let type_name = node_text(type_cap.node, source_bytes);
154 edges.push(EdgeDef {
155 src: fp.clone(),
156 dst: format!("cls:{}:{}", file.relative_path, type_name),
157 kind: EdgeKind::Exports,
158 ..Default::default()
159 });
160 }
161 }
162 }
163
164 if let Ok(query) = Query::new(
166 &self.language,
167 "(use_declaration argument: (scoped_identifier path: (_) @path name: (_)?))",
168 ) {
169 let mut cursor = QueryCursor::new();
170 for m in cursor.matches(&query, root, source_bytes) {
171 if let Some(path_cap) = m
172 .captures
173 .iter()
174 .find(|c| query.capture_names()[c.index as usize] == "path")
175 {
176 let full_path = node_text(path_cap.node, source_bytes);
177 let import_path = if full_path.starts_with("crate::") {
179 format!(
180 "src/{}.rs",
181 full_path.trim_start_matches("crate::").replace("::", "/")
182 )
183 } else {
184 continue;
185 };
186 edges.push(EdgeDef {
187 src: fp.clone(),
188 dst: format!("file:{}", import_path),
189 kind: EdgeKind::Imports,
190 ..Default::default()
191 });
192 }
193 }
194 }
195
196 if let Ok(query) = Query::new(
198 &self.language,
199 "(use_declaration argument: (identifier) @name)",
200 ) {
201 let mut cursor = QueryCursor::new();
202 for m in cursor.matches(&query, root, source_bytes) {
203 if let Some(name_cap) = m
204 .captures
205 .iter()
206 .find(|c| query.capture_names()[c.index as usize] == "name")
207 {
208 let mod_name = node_text(name_cap.node, source_bytes);
209 let import_path = mod_name;
210 edges.push(EdgeDef {
211 src: fp.clone(),
212 dst: format!("file:{}.rs", import_path),
213 kind: EdgeKind::Imports,
214 ..Default::default()
215 });
216 }
217 }
218 }
219
220 Ok(ParseResult {
221 nodes,
222 edges,
223 ..Default::default()
224 })
225 }
226}
227
228#[allow(clippy::too_many_arguments)]
229fn extract_type_nodes(
230 nodes: &mut Vec<NodeDef>,
231 edges: &mut Vec<EdgeDef>,
232 file_id: &str,
233 file: &SourceFile,
234 query: &Query,
235 root: tree_sitter::Node,
236 source_bytes: &[u8],
237 kind: NodeKind,
238 prefix: &str,
239) {
240 let mut cursor = QueryCursor::new();
241 for m in cursor.matches(query, root, source_bytes) {
242 let Some(name_capture) = m
243 .captures
244 .iter()
245 .find(|c| query.capture_names()[c.index as usize] == "name")
246 else {
247 continue;
248 };
249 let name = node_text(name_capture.node, source_bytes);
250 let start = name_capture.node.start_position();
251 let body_end = m
253 .captures
254 .iter()
255 .find(|c| query.capture_names()[c.index as usize] != "name")
256 .map(|c| c.node.end_position())
257 .unwrap_or_else(|| name_capture.node.end_position());
258 let id = format!("{}:{}:{}", prefix, file.relative_path, name);
259
260 nodes.push(NodeDef {
261 id: id.clone(),
262 kind: kind.clone(),
263 name,
264 path: file.relative_path.clone(),
265 line_start: start.row as u32 + 1,
266 line_end: body_end.row as u32 + 1,
267 ..Default::default()
268 });
269
270 edges.push(EdgeDef {
271 src: file_id.to_string(),
272 dst: id,
273 kind: EdgeKind::Exports,
274 ..Default::default()
275 });
276 }
277}
278
279fn node_text(node: tree_sitter::Node, source: &[u8]) -> String {
280 node.utf8_text(source).unwrap_or("").to_string()
281}