Skip to main content

codesynapse_core/ts_extract/
zig.rs

1use super::{add_contains_edge, add_node_if_missing, make_file_node, node_text, run_query_named};
2use crate::error::{CodeSynapseError, 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;
8
9fn contains_struct_decl(node: TsNode) -> bool {
10    if node.kind() == "struct_declaration" {
11        return true;
12    }
13    let mut cursor = node.walk();
14    for child in node.children(&mut cursor) {
15        if contains_struct_decl(child) {
16            return true;
17        }
18    }
19    false
20}
21
22fn tree_walk_structs(
23    node: TsNode,
24    source: &[u8],
25    file_id: &str,
26    fragment: &mut ExtractionFragment,
27    path: &Path,
28) {
29    if node.kind() == "variable_declaration" {
30        let mut name = None;
31        let mut cursor = node.walk();
32        for child in node.children(&mut cursor) {
33            if child.kind() == "identifier" {
34                if let Ok(text) = node_text(source, &child) {
35                    name = Some(text.to_string());
36                }
37            }
38        }
39        if let Some(n) = name {
40            if contains_struct_decl(node) {
41                let id = make_id(&[file_id, &n]);
42                fragment.nodes.push(Node {
43                    id: id.clone(),
44                    label: n,
45                    file_type: "struct".to_string(),
46                    source_file: path.to_string_lossy().to_string(),
47                    source_location: None,
48                    community: None,
49                    rationale: None,
50                    docstring: None,
51                    metadata: HashMap::new(),
52                });
53                add_contains_edge(fragment, &file_id.to_string(), id, path);
54            }
55        }
56    }
57    let mut cursor = node.walk();
58    for child in node.children(&mut cursor) {
59        tree_walk_structs(child, source, file_id, fragment, path);
60    }
61}
62
63pub struct TsZigExtractor;
64
65impl TsZigExtractor {
66    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
67        let (file_id, _, file_node) = make_file_node(path);
68        let mut fragment = ExtractionFragment {
69            nodes: vec![file_node],
70            edges: vec![],
71        };
72
73        let lang = tree_sitter_zig::LANGUAGE.into();
74
75        let mut parser = tree_sitter::Parser::new();
76        parser
77            .set_language(&lang)
78            .map_err(|e| CodeSynapseError::Parse(format!("failed to set zig language: {}", e)))?;
79        if let Some(tree) = parser.parse(source, None) {
80            let root = tree.root_node();
81            tree_walk_structs(root, source, &file_id, &mut fragment, path);
82        }
83
84        let func_query = r#"
85            (function_declaration
86                name: (identifier) @func.name
87            )
88        "#;
89        if let Ok(mut captures) = run_query_named(source, &lang, func_query) {
90            for name in captures.remove("func.name").unwrap_or_default() {
91                let label = format!("{}()", name);
92                let id = make_id(&[&file_id, &name]);
93                fragment.nodes.push(Node {
94                    id: id.clone(),
95                    label,
96                    file_type: "function".to_string(),
97                    source_file: path.to_string_lossy().to_string(),
98                    source_location: None,
99                    community: None,
100                    rationale: None,
101                    docstring: None,
102                    metadata: HashMap::new(),
103                });
104                add_contains_edge(&mut fragment, &file_id, id, path);
105            }
106        }
107
108        let import_query = r#"
109            ((builtin_function
110                (builtin_identifier) @builtin.name
111                (arguments (string) @import.path))
112             (#eq? @builtin.name "@import"))
113        "#;
114        if let Ok(mut captures) = run_query_named(source, &lang, import_query) {
115            for path_str in captures.remove("import.path").unwrap_or_default() {
116                let module = path_str.trim_matches('"').trim();
117                let mod_id = make_id(&[&file_id, module]);
118                let mod_node = Node {
119                    id: mod_id.clone(),
120                    label: module.to_string(),
121                    file_type: "module".to_string(),
122                    source_file: path.to_string_lossy().to_string(),
123                    source_location: None,
124                    community: None,
125                    rationale: None,
126                    docstring: None,
127                    metadata: HashMap::new(),
128                };
129                add_node_if_missing(&mut fragment, mod_node);
130                fragment.edges.push(Edge {
131                    source: file_id.clone(),
132                    target: mod_id,
133                    relation: "imports".to_string(),
134                    confidence: "EXTRACTED".to_string(),
135                    source_file: Some(path.to_string_lossy().to_string()),
136                    weight: 1.0,
137                    context: None,
138                });
139            }
140        }
141
142        Ok(fragment)
143    }
144}
145
146impl LanguageExtractor for TsZigExtractor {
147    fn file_extensions(&self) -> Vec<&'static str> {
148        vec!["zig"]
149    }
150    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
151        Self::extract(source, path)
152    }
153    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
154        vec![]
155    }
156    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
157}