1use super::{add_contains_edge, add_node_if_missing, make_file_node, run_query_named};
2use crate::error::Result;
3use crate::extract::{make_id, ImportNode, LanguageExtractor};
4use crate::types::{Edge, ExtractionFragment, Node};
5use std::collections::HashMap;
6use std::path::Path;
7use tree_sitter::{Node as TsNode, Parser};
8
9pub struct TsElixirExtractor;
10
11fn ex_text(source: &[u8], node: &TsNode<'_>) -> String {
12 std::str::from_utf8(&source[node.start_byte()..node.end_byte()])
13 .unwrap_or("")
14 .trim()
15 .to_string()
16}
17
18const ELIXIR_KEYWORDS: &[&str] = &[
19 "def",
20 "defp",
21 "defmodule",
22 "defmacro",
23 "defmacrop",
24 "defstruct",
25 "defprotocol",
26 "defimpl",
27 "defguard",
28 "alias",
29 "import",
30 "require",
31 "use",
32 "if",
33 "unless",
34 "case",
35 "cond",
36 "with",
37 "for",
38];
39
40impl TsElixirExtractor {
41 pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
42 let (file_id, _, file_node) = make_file_node(path);
43 let mut fragment = ExtractionFragment {
44 nodes: vec![file_node],
45 edges: vec![],
46 };
47
48 let lang = tree_sitter_elixir::LANGUAGE.into();
49
50 let module_query = r#"
52 (call
53 target: (identifier) @kw
54 (arguments (alias) @module.name)
55 (#eq? @kw "defmodule"))
56 "#;
57 if let Ok(mut captures) = run_query_named(source, &lang, module_query) {
58 for name in captures.remove("module.name").unwrap_or_default() {
59 let id = make_id(&[&file_id, &name]);
60 fragment.nodes.push(Node {
61 id: id.clone(),
62 label: name,
63 file_type: "module".to_string(),
64 source_file: path.to_string_lossy().to_string(),
65 source_location: None,
66 community: None,
67 rationale: None,
68 docstring: None,
69 metadata: HashMap::new(),
70 });
71 add_contains_edge(&mut fragment, &file_id, id, path);
72 }
73 }
74
75 let func_query = r#"
77 (call
78 target: (identifier) @kw
79 (arguments (call target: (identifier) @func.name))
80 (#match? @kw "^defp?$"))
81 "#;
82 if let Ok(mut captures) = run_query_named(source, &lang, func_query) {
83 for name in captures.remove("func.name").unwrap_or_default() {
84 let label = format!("{}()", name);
85 let id = make_id(&[&file_id, &name]);
86 add_node_if_missing(
87 &mut fragment,
88 Node {
89 id: id.clone(),
90 label,
91 file_type: "function".to_string(),
92 source_file: path.to_string_lossy().to_string(),
93 source_location: None,
94 community: None,
95 rationale: None,
96 docstring: None,
97 metadata: HashMap::new(),
98 },
99 );
100 add_contains_edge(&mut fragment, &file_id, id, path);
101 }
102 }
103
104 let import_query = r#"
106 (call
107 target: (identifier) @kw
108 (arguments (alias) @import.name)
109 (#match? @kw "^(alias|import|require|use)$"))
110 "#;
111 if let Ok(mut captures) = run_query_named(source, &lang, import_query) {
112 for name in captures.remove("import.name").unwrap_or_default() {
113 let mod_id = make_id(&[&name]);
114 add_node_if_missing(
115 &mut fragment,
116 Node {
117 id: mod_id.clone(),
118 label: name.clone(),
119 file_type: "module".to_string(),
120 source_file: path.to_string_lossy().to_string(),
121 source_location: None,
122 community: None,
123 rationale: None,
124 docstring: None,
125 metadata: HashMap::new(),
126 },
127 );
128 fragment.edges.push(Edge {
129 source: file_id.clone(),
130 target: mod_id,
131 relation: "imports".to_string(),
132 confidence: "EXTRACTED".to_string(),
133 source_file: Some(path.to_string_lossy().to_string()),
134 weight: 1.0,
135 context: Some("import".to_string()),
136 });
137 }
138 }
139
140 Self::extract_calls(source, path, &file_id, &lang, &mut fragment);
142
143 Ok(fragment)
144 }
145
146 fn extract_calls(
147 source: &[u8],
148 path: &Path,
149 file_id: &str,
150 lang: &tree_sitter::Language,
151 fragment: &mut ExtractionFragment,
152 ) {
153 let mut parser = Parser::new();
154 if parser.set_language(lang).is_err() {
155 return;
156 }
157 let tree = match parser.parse(source, None) {
158 Some(t) => t,
159 None => return,
160 };
161 let root = tree.root_node();
162
163 let label_to_id: HashMap<String, String> = fragment
164 .nodes
165 .iter()
166 .filter(|n| n.file_type == "function")
167 .map(|n| {
168 let label = n.label.trim_end_matches("()").to_string();
169 (label, n.id.clone())
170 })
171 .collect();
172
173 let mut seen: std::collections::HashSet<(String, String)> = Default::default();
174 Self::walk_elixir_calls(
175 root,
176 source,
177 path,
178 file_id,
179 fragment,
180 None,
181 &label_to_id,
182 &mut seen,
183 );
184 }
185
186 #[allow(clippy::only_used_in_recursion, clippy::too_many_arguments)]
187 fn walk_elixir_calls(
188 node: TsNode<'_>,
189 source: &[u8],
190 path: &Path,
191 file_id: &str,
192 fragment: &mut ExtractionFragment,
193 caller_id: Option<&str>,
194 label_to_id: &HashMap<String, String>,
195 seen: &mut std::collections::HashSet<(String, String)>,
196 ) {
197 if node.kind() != "call" {
198 for i in 0..node.child_count() {
199 if let Some(child) = node.child(i) {
200 Self::walk_elixir_calls(
201 child,
202 source,
203 path,
204 file_id,
205 fragment,
206 caller_id,
207 label_to_id,
208 seen,
209 );
210 }
211 }
212 return;
213 }
214
215 let mut kw = None;
217 let mut args_node = None;
218 let mut do_block = None;
219 for i in 0..node.child_count() {
220 if let Some(child) = node.child(i) {
221 match child.kind() {
222 "identifier" if kw.is_none() => {
223 kw = Some(ex_text(source, &child));
224 }
225 "arguments" => args_node = Some(child),
226 "do_block" => do_block = Some(child),
227 _ => {}
228 }
229 }
230 }
231
232 let kw_str = kw.as_deref().unwrap_or("");
233
234 if kw_str == "def" || kw_str == "defp" {
235 let func_name = args_node.and_then(|args| {
237 for i in 0..args.child_count() {
239 if let Some(child) = args.child(i) {
240 if child.kind() == "call" {
241 for j in 0..child.child_count() {
242 if let Some(id_node) = child.child(j) {
243 if id_node.kind() == "identifier" {
244 return Some(ex_text(source, &id_node));
245 }
246 }
247 }
248 } else if child.kind() == "identifier" {
249 return Some(ex_text(source, &child));
250 }
251 }
252 }
253 None
254 });
255 let new_caller = func_name
256 .as_ref()
257 .and_then(|n| label_to_id.get(n))
258 .map(|s| s.as_str());
259 let effective_caller = new_caller.or(caller_id);
260 if let Some(do_b) = do_block {
261 for i in 0..do_b.child_count() {
262 if let Some(child) = do_b.child(i) {
263 Self::walk_elixir_calls(
264 child,
265 source,
266 path,
267 file_id,
268 fragment,
269 effective_caller,
270 label_to_id,
271 seen,
272 );
273 }
274 }
275 }
276 return;
277 }
278
279 if kw_str == "defmodule" {
280 if let Some(do_b) = do_block {
281 for i in 0..do_b.child_count() {
282 if let Some(child) = do_b.child(i) {
283 Self::walk_elixir_calls(
284 child,
285 source,
286 path,
287 file_id,
288 fragment,
289 caller_id,
290 label_to_id,
291 seen,
292 );
293 }
294 }
295 }
296 return;
297 }
298
299 if ELIXIR_KEYWORDS.contains(&kw_str) {
301 for i in 0..node.child_count() {
302 if let Some(child) = node.child(i) {
303 Self::walk_elixir_calls(
304 child,
305 source,
306 path,
307 file_id,
308 fragment,
309 caller_id,
310 label_to_id,
311 seen,
312 );
313 }
314 }
315 return;
316 }
317
318 if let Some(caller) = caller_id {
320 if let Some(callee_id) = label_to_id.get(kw_str) {
321 let pair = (caller.to_string(), callee_id.clone());
322 if !seen.contains(&pair) && caller != callee_id.as_str() {
323 seen.insert(pair);
324 fragment.edges.push(Edge {
325 source: caller.to_string(),
326 target: callee_id.clone(),
327 relation: "calls".to_string(),
328 confidence: "EXTRACTED".to_string(),
329 source_file: Some(path.to_string_lossy().to_string()),
330 weight: 1.0,
331 context: Some("call".to_string()),
332 });
333 }
334 }
335 }
336
337 for i in 0..node.child_count() {
338 if let Some(child) = node.child(i) {
339 Self::walk_elixir_calls(
340 child,
341 source,
342 path,
343 file_id,
344 fragment,
345 caller_id,
346 label_to_id,
347 seen,
348 );
349 }
350 }
351 }
352}
353
354impl LanguageExtractor for TsElixirExtractor {
355 fn file_extensions(&self) -> Vec<&'static str> {
356 vec!["ex", "exs"]
357 }
358 fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
359 Self::extract(source, path)
360 }
361 fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
362 vec![]
363 }
364 fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
365}