dogfood
Jump to…
snowcorepublic2 bookmarks
snowinitial commitqoxwzsukwmkx1mo
mainroot / crates / df-render / src / symbols.rsBlameRaw
1//! Symbol extraction from source files.
2//!
3//! Extracts top-level symbols (functions, structs, classes, interfaces, etc.)
4//! from source code using tree-sitter grammars. This powers the symbol outline
5//! sidebar in the code view.
6//!
7//! The extraction is best-effort: languages without a grammar or files that fail
8//! to parse simply return an empty list, never an error.
9
10use tree_sitter::{Parser, Tree};
11
12/// A symbol extracted from a source file.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct Symbol {
15 /// The symbol name (e.g. `main`, `MyStruct`, `handle_request`).
16 pub name: String,
17 /// What kind of symbol this is.
18 pub kind: SymbolKind,
19 /// 1-based line number where the symbol starts.
20 pub line: usize,
21}
22
23/// The kind of symbol.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum SymbolKind {
26 Function,
27 Method,
28 Class,
29 Struct,
30 Enum,
31 Interface,
32 Trait,
33 Constant,
34 Module,
35 Type,
36}
37
38impl SymbolKind {
39 /// Short label for the UI.
40 pub fn label(&self) -> &'static str {
41 match self {
42 SymbolKind::Function => "fn",
43 SymbolKind::Method => "method",
44 SymbolKind::Class => "class",
45 SymbolKind::Struct => "struct",
46 SymbolKind::Enum => "enum",
47 SymbolKind::Interface => "iface",
48 SymbolKind::Trait => "trait",
49 SymbolKind::Constant => "const",
50 SymbolKind::Module => "mod",
51 SymbolKind::Type => "type",
52 }
53 }
54
55 /// CSS class suffix for styling.
56 pub fn css_class(&self) -> &'static str {
57 match self {
58 SymbolKind::Function | SymbolKind::Method => "fn",
59 SymbolKind::Class | SymbolKind::Struct => "type",
60 SymbolKind::Enum => "enum",
61 SymbolKind::Interface | SymbolKind::Trait => "trait",
62 SymbolKind::Constant => "const",
63 SymbolKind::Module => "mod",
64 SymbolKind::Type => "type",
65 }
66 }
67}
68
69/// Extract symbols from the given source code.
70///
71/// `path` is used to determine the language (by extension). Returns an empty
72/// vec for unsupported languages or if parsing fails.
73pub fn extract_symbols(path: &str, source: &str) -> Vec<Symbol> {
74 // Don't attempt symbol extraction on very large files.
75 if source.len() > 1024 * 1024 {
76 return Vec::new();
77 }
78
79 let lang = match language_for_symbols(path) {
80 Some(l) => l,
81 None => return Vec::new(),
82 };
83
84 let mut parser = Parser::new();
85 if parser.set_language(&lang.ts_language).is_err() {
86 return Vec::new();
87 }
88
89 let tree = match parser.parse(source, None) {
90 Some(t) => t,
91 None => return Vec::new(),
92 };
93
94 extract_from_tree(&tree, source, lang.node_kinds)
95}
96
97/// Walk the tree and collect symbols based on language-specific node kinds.
98fn extract_from_tree(tree: &Tree, source: &str, node_kinds: &[NodeKindMapping]) -> Vec<Symbol> {
99 let mut symbols = Vec::new();
100 let root = tree.root_node();
101
102 let mut cursor = root.walk();
103 // Only walk top-level and one level deep (for impl blocks, class bodies).
104 for node in root.children(&mut cursor) {
105 extract_node(&node, source, node_kinds, &mut symbols, 0);
106 }
107
108 symbols
109}
110
111/// Maximum nesting depth for symbol extraction.
112const MAX_DEPTH: usize = 2;
113
114fn extract_node(
115 node: &tree_sitter::Node,
116 source: &str,
117 node_kinds: &[NodeKindMapping],
118 symbols: &mut Vec<Symbol>,
119 depth: usize,
120) {
121 let kind = node.kind();
122
123 // Check if this node kind maps to a symbol.
124 for mapping in node_kinds {
125 if mapping.node_kind == kind {
126 if let Some(name) = find_name_child(node, source, mapping.name_field) {
127 symbols.push(Symbol {
128 name,
129 kind: mapping.symbol_kind,
130 line: node.start_position().row + 1,
131 });
132 }
133 break;
134 }
135 }
136
137 // Recurse into container nodes (impl blocks, class bodies) but not too deep.
138 if depth < MAX_DEPTH && is_container_node(kind) {
139 let mut cursor = node.walk();
140 for child in node.children(&mut cursor) {
141 extract_node(&child, source, node_kinds, symbols, depth + 1);
142 }
143 }
144}
145
146/// Find the name of a node by looking for a specific child field.
147fn find_name_child(
148 node: &tree_sitter::Node,
149 source: &str,
150 name_field: &str,
151) -> Option<String> {
152 let child = node.child_by_field_name(name_field)?;
153 let text = child.utf8_text(source.as_bytes()).ok()?;
154 // Limit name length to avoid absurd entries.
155 if text.len() > 200 {
156 return None;
157 }
158 Some(text.to_string())
159}
160
161/// Whether a node is a "container" that might have nested symbols.
162fn is_container_node(kind: &str) -> bool {
163 matches!(
164 kind,
165 "impl_item"
166 | "class_declaration"
167 | "class_definition"
168 | "class_body"
169 | "module"
170 | "namespace_declaration"
171 | "object_declaration"
172 | "trait_definition"
173 | "interface_declaration"
174 | "declaration_list"
175 | "block"
176 | "program"
177 )
178}
179
180// ─── Language-specific configuration ─────────────────────────────────────────
181
182struct NodeKindMapping {
183 node_kind: &'static str,
184 name_field: &'static str,
185 symbol_kind: SymbolKind,
186}
187
188struct SymbolLanguage {
189 ts_language: tree_sitter::Language,
190 node_kinds: &'static [NodeKindMapping],
191}
192
193const RUST_KINDS: &[NodeKindMapping] = &[
194 NodeKindMapping { node_kind: "function_item", name_field: "name", symbol_kind: SymbolKind::Function },
195 NodeKindMapping { node_kind: "struct_item", name_field: "name", symbol_kind: SymbolKind::Struct },
196 NodeKindMapping { node_kind: "enum_item", name_field: "name", symbol_kind: SymbolKind::Enum },
197 NodeKindMapping { node_kind: "trait_item", name_field: "name", symbol_kind: SymbolKind::Trait },
198 NodeKindMapping { node_kind: "type_item", name_field: "name", symbol_kind: SymbolKind::Type },
199 NodeKindMapping { node_kind: "const_item", name_field: "name", symbol_kind: SymbolKind::Constant },
200 NodeKindMapping { node_kind: "mod_item", name_field: "name", symbol_kind: SymbolKind::Module },
201 NodeKindMapping { node_kind: "impl_item", name_field: "type", symbol_kind: SymbolKind::Struct },
202];
203
204const PYTHON_KINDS: &[NodeKindMapping] = &[
205 NodeKindMapping { node_kind: "function_definition", name_field: "name", symbol_kind: SymbolKind::Function },
206 NodeKindMapping { node_kind: "class_definition", name_field: "name", symbol_kind: SymbolKind::Class },
207];
208
209const JS_KINDS: &[NodeKindMapping] = &[
210 NodeKindMapping { node_kind: "function_declaration", name_field: "name", symbol_kind: SymbolKind::Function },
211 NodeKindMapping { node_kind: "class_declaration", name_field: "name", symbol_kind: SymbolKind::Class },
212 NodeKindMapping { node_kind: "method_definition", name_field: "name", symbol_kind: SymbolKind::Method },
213];
214
215const GO_KINDS: &[NodeKindMapping] = &[
216 NodeKindMapping { node_kind: "function_declaration", name_field: "name", symbol_kind: SymbolKind::Function },
217 NodeKindMapping { node_kind: "method_declaration", name_field: "name", symbol_kind: SymbolKind::Method },
218 NodeKindMapping { node_kind: "type_declaration", name_field: "name", symbol_kind: SymbolKind::Type },
219];
220
221const JAVA_KINDS: &[NodeKindMapping] = &[
222 NodeKindMapping { node_kind: "class_declaration", name_field: "name", symbol_kind: SymbolKind::Class },
223 NodeKindMapping { node_kind: "method_declaration", name_field: "name", symbol_kind: SymbolKind::Method },
224 NodeKindMapping { node_kind: "interface_declaration", name_field: "name", symbol_kind: SymbolKind::Interface },
225 NodeKindMapping { node_kind: "enum_declaration", name_field: "name", symbol_kind: SymbolKind::Enum },
226];
227
228const C_KINDS: &[NodeKindMapping] = &[
229 NodeKindMapping { node_kind: "function_definition", name_field: "declarator", symbol_kind: SymbolKind::Function },
230 NodeKindMapping { node_kind: "struct_specifier", name_field: "name", symbol_kind: SymbolKind::Struct },
231 NodeKindMapping { node_kind: "enum_specifier", name_field: "name", symbol_kind: SymbolKind::Enum },
232];
233
234const RUBY_KINDS: &[NodeKindMapping] = &[
235 NodeKindMapping { node_kind: "method", name_field: "name", symbol_kind: SymbolKind::Method },
236 NodeKindMapping { node_kind: "class", name_field: "name", symbol_kind: SymbolKind::Class },
237 NodeKindMapping { node_kind: "module", name_field: "name", symbol_kind: SymbolKind::Module },
238];
239
240fn language_for_symbols(path: &str) -> Option<SymbolLanguage> {
241 let name = path.rsplit('/').next().unwrap_or(path);
242 let ext = name.rsplit_once('.').map(|(_, e)| e.to_ascii_lowercase())?;
243
244 match ext.as_str() {
245 "rs" => Some(SymbolLanguage {
246 ts_language: tree_sitter_rust::LANGUAGE.into(),
247 node_kinds: RUST_KINDS,
248 }),
249 "py" | "pyi" | "pyw" => Some(SymbolLanguage {
250 ts_language: tree_sitter_python::LANGUAGE.into(),
251 node_kinds: PYTHON_KINDS,
252 }),
253 "js" | "mjs" | "cjs" | "jsx" => Some(SymbolLanguage {
254 ts_language: tree_sitter_javascript::LANGUAGE.into(),
255 node_kinds: JS_KINDS,
256 }),
257 "ts" | "mts" | "cts" => Some(SymbolLanguage {
258 ts_language: tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
259 node_kinds: JS_KINDS,
260 }),
261 "tsx" => Some(SymbolLanguage {
262 ts_language: tree_sitter_typescript::LANGUAGE_TSX.into(),
263 node_kinds: JS_KINDS,
264 }),
265 "go" => Some(SymbolLanguage {
266 ts_language: tree_sitter_go::LANGUAGE.into(),
267 node_kinds: GO_KINDS,
268 }),
269 "c" | "h" => Some(SymbolLanguage {
270 ts_language: tree_sitter_c::LANGUAGE.into(),
271 node_kinds: C_KINDS,
272 }),
273 "cc" | "cpp" | "cxx" | "hpp" | "hh" | "hxx" => Some(SymbolLanguage {
274 ts_language: tree_sitter_cpp::LANGUAGE.into(),
275 node_kinds: C_KINDS,
276 }),
277 "java" => Some(SymbolLanguage {
278 ts_language: tree_sitter_java::LANGUAGE.into(),
279 node_kinds: JAVA_KINDS,
280 }),
281 "rb" | "rake" | "gemspec" => Some(SymbolLanguage {
282 ts_language: tree_sitter_ruby::LANGUAGE.into(),
283 node_kinds: RUBY_KINDS,
284 }),
285 _ => None,
286 }
287}
288
289#[cfg(test)]
290mod tests {
291 use super::*;
292
293 #[test]
294 fn extracts_rust_symbols() {
295 let source = r#"
296fn main() {
297 println!("hello");
298}
299
300struct Foo {
301 x: i32,
302}
303
304enum Bar {
305 A,
306 B,
307}
308
309trait Baz {
310 fn method(&self);
311}
312"#;
313 let symbols = extract_symbols("src/main.rs", source);
314 let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
315 assert!(names.contains(&"main"));
316 assert!(names.contains(&"Foo"));
317 assert!(names.contains(&"Bar"));
318 assert!(names.contains(&"Baz"));
319 }
320
321 #[test]
322 fn extracts_python_symbols() {
323 let source = r#"
324def hello():
325 pass
326
327class World:
328 def method(self):
329 pass
330"#;
331 let symbols = extract_symbols("app.py", source);
332 let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect();
333 assert!(names.contains(&"hello"));
334 assert!(names.contains(&"World"));
335 }
336
337 #[test]
338 fn unsupported_language_returns_empty() {
339 let symbols = extract_symbols("data.csv", "a,b,c\n1,2,3\n");
340 assert!(symbols.is_empty());
341 }
342
343 #[test]
344 fn empty_source_returns_empty() {
345 let symbols = extract_symbols("main.rs", "");
346 assert!(symbols.is_empty());
347 }
348}

348 lines · Rust