Skip to main content

codesynapse_core/ts_extract/
elixir.rs

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        // defmodule
51        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        // def/defp functions
76        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        // alias / import / require / use → import edges with context="import"
105        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        // Call extraction: find non-keyword calls inside function bodies
141        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        // Find identifier (keyword)
216        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            // Recurse into do_block with this function as the caller
236            let func_name = args_node.and_then(|args| {
237                // Find call > identifier inside arguments
238                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        // Skip known keywords
300        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        // This is a non-keyword call
319        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}