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 { nodes, edges })
221 }
222}
223
224#[allow(clippy::too_many_arguments)]
225fn extract_type_nodes(
226 nodes: &mut Vec<NodeDef>,
227 edges: &mut Vec<EdgeDef>,
228 file_id: &str,
229 file: &SourceFile,
230 query: &Query,
231 root: tree_sitter::Node,
232 source_bytes: &[u8],
233 kind: NodeKind,
234 prefix: &str,
235) {
236 let mut cursor = QueryCursor::new();
237 for m in cursor.matches(query, root, source_bytes) {
238 let Some(name_capture) = m
239 .captures
240 .iter()
241 .find(|c| query.capture_names()[c.index as usize] == "name")
242 else {
243 continue;
244 };
245 let name = node_text(name_capture.node, source_bytes);
246 let start = name_capture.node.start_position();
247 let body_end = m
249 .captures
250 .iter()
251 .find(|c| query.capture_names()[c.index as usize] != "name")
252 .map(|c| c.node.end_position())
253 .unwrap_or_else(|| name_capture.node.end_position());
254 let id = format!("{}:{}:{}", prefix, file.relative_path, name);
255
256 nodes.push(NodeDef {
257 id: id.clone(),
258 kind: kind.clone(),
259 name,
260 path: file.relative_path.clone(),
261 line_start: start.row as u32 + 1,
262 line_end: body_end.row as u32 + 1,
263 ..Default::default()
264 });
265
266 edges.push(EdgeDef {
267 src: file_id.to_string(),
268 dst: id,
269 kind: EdgeKind::Exports,
270 ..Default::default()
271 });
272 }
273}
274
275fn node_text(node: tree_sitter::Node, source: &[u8]) -> String {
276 node.utf8_text(source).unwrap_or("").to_string()
277}