amql_engine/extractor/
go_structure.rs

1//! General-purpose Go structure extractor.
2//!
3//! Extracts functions, methods, structs, interfaces, type aliases, constants,
4//! and variables as annotations. Mirrors the Go resolver logic but produces
5//! `Annotation` metadata instead of `CodeElement` trees, discarding
6//! implementation bodies.
7
8use super::BuiltinExtractor;
9use crate::store::Annotation;
10use crate::types::{AttrName, Binding, RelativePath, TagName};
11use rustc_hash::FxHashMap;
12use serde_json::Value as JsonValue;
13use std::cell::RefCell;
14
15/// General-purpose Go structure extractor.
16pub struct GoStructureExtractor;
17
18impl BuiltinExtractor for GoStructureExtractor {
19    fn name(&self) -> &str {
20        "go-structure"
21    }
22
23    fn extensions(&self) -> &[&str] {
24        &[".go"]
25    }
26
27    fn extract(&self, source: &str, file: &RelativePath) -> Vec<Annotation> {
28        let tree = match parse_go(source) {
29            Some(t) => t,
30            None => return vec![],
31        };
32        let mut annotations = Vec::new();
33        let root = tree.root_node();
34        let src = source.as_bytes();
35        let mut cursor = root.walk();
36        for child in root.named_children(&mut cursor) {
37            extract_elements(&child, src, file, &mut annotations);
38        }
39        annotations
40    }
41}
42
43// ---------------------------------------------------------------------------
44// Tree-sitter helpers
45// ---------------------------------------------------------------------------
46
47fn node_text<'a>(node: &tree_sitter::Node, src: &'a [u8]) -> &'a str {
48    node.utf8_text(src).unwrap_or("")
49}
50
51fn get_name(node: &tree_sitter::Node, src: &[u8]) -> String {
52    node.child_by_field_name("name")
53        .map(|n| node_text(&n, src).to_string())
54        .unwrap_or_default()
55}
56
57fn make_annotation(
58    tag: &str,
59    binding: String,
60    attrs: FxHashMap<AttrName, JsonValue>,
61    file: &RelativePath,
62    children: Vec<Annotation>,
63) -> Annotation {
64    Annotation {
65        tag: TagName::from(tag),
66        attrs,
67        binding: Binding::from(binding),
68        file: file.clone(),
69        children,
70    }
71}
72
73/// Check if a Go identifier is exported (starts with uppercase).
74fn is_exported(name: &str) -> bool {
75    name.starts_with(|c: char| c.is_uppercase())
76}
77
78// ---------------------------------------------------------------------------
79// Element extraction
80// ---------------------------------------------------------------------------
81
82fn extract_elements(
83    node: &tree_sitter::Node,
84    src: &[u8],
85    file: &RelativePath,
86    result: &mut Vec<Annotation>,
87) {
88    match node.kind() {
89        "function_declaration" => {
90            if let Some(ann) = extract_function(node, src, file) {
91                result.push(ann);
92            }
93        }
94        "method_declaration" => {
95            if let Some(ann) = extract_method(node, src, file) {
96                result.push(ann);
97            }
98        }
99        "type_declaration" => {
100            let mut cursor = node.walk();
101            for child in node.named_children(&mut cursor) {
102                if child.kind() == "type_spec" {
103                    if let Some(ann) = extract_type_spec(&child, src, file) {
104                        result.push(ann);
105                    }
106                }
107            }
108        }
109        "const_declaration" => {
110            let mut cursor = node.walk();
111            for child in node.named_children(&mut cursor) {
112                if child.kind() == "const_spec" {
113                    if let Some(ann) = extract_const_spec(&child, src, file) {
114                        result.push(ann);
115                    }
116                }
117            }
118        }
119        "var_declaration" => {
120            let mut cursor = node.walk();
121            for child in node.named_children(&mut cursor) {
122                if child.kind() == "var_spec" {
123                    if let Some(ann) = extract_var_spec(&child, src, file) {
124                        result.push(ann);
125                    }
126                }
127            }
128        }
129        _ => {}
130    }
131}
132
133fn extract_function(
134    node: &tree_sitter::Node,
135    src: &[u8],
136    file: &RelativePath,
137) -> Option<Annotation> {
138    let name = get_name(node, src);
139    if name.is_empty() {
140        return None;
141    }
142    let mut attrs = FxHashMap::default();
143    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
144    if is_exported(&name) {
145        attrs.insert(AttrName::from("exported"), JsonValue::Bool(true));
146    }
147    Some(make_annotation("function", name, attrs, file, vec![]))
148}
149
150fn extract_method(node: &tree_sitter::Node, src: &[u8], file: &RelativePath) -> Option<Annotation> {
151    let name = get_name(node, src);
152    if name.is_empty() {
153        return None;
154    }
155    let mut attrs = FxHashMap::default();
156    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
157
158    // Extract receiver type
159    if let Some(receiver) = node.child_by_field_name("receiver") {
160        let receiver_text = node_text(&receiver, src);
161        let cleaned = receiver_text
162            .trim_matches(|c: char| c == '(' || c == ')')
163            .trim();
164        let type_name = cleaned
165            .split_whitespace()
166            .last()
167            .unwrap_or(cleaned)
168            .trim_start_matches('*');
169        attrs.insert(
170            AttrName::from("receiver"),
171            JsonValue::String(type_name.to_string()),
172        );
173    }
174
175    if is_exported(&name) {
176        attrs.insert(AttrName::from("exported"), JsonValue::Bool(true));
177    }
178    Some(make_annotation("method", name, attrs, file, vec![]))
179}
180
181fn extract_type_spec(
182    node: &tree_sitter::Node,
183    src: &[u8],
184    file: &RelativePath,
185) -> Option<Annotation> {
186    let name = get_name(node, src);
187    if name.is_empty() {
188        return None;
189    }
190    let type_node = node.child_by_field_name("type")?;
191    let mut attrs = FxHashMap::default();
192    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
193    if is_exported(&name) {
194        attrs.insert(AttrName::from("exported"), JsonValue::Bool(true));
195    }
196
197    let (tag, children) = match type_node.kind() {
198        "struct_type" => {
199            let children = extract_struct_fields(&type_node, src, file);
200            ("struct", children)
201        }
202        "interface_type" => {
203            let children = extract_interface_methods(&type_node, src, file);
204            ("interface", children)
205        }
206        _ => ("type", vec![]),
207    };
208
209    Some(make_annotation(tag, name, attrs, file, children))
210}
211
212/// Extract struct field declarations as children.
213fn extract_struct_fields(
214    node: &tree_sitter::Node,
215    src: &[u8],
216    file: &RelativePath,
217) -> Vec<Annotation> {
218    let mut fields = Vec::new();
219    let mut cursor = node.walk();
220    for child in node.named_children(&mut cursor) {
221        if child.kind() == "field_declaration_list" {
222            let mut inner_cursor = child.walk();
223            for field in child.named_children(&mut inner_cursor) {
224                if field.kind() == "field_declaration" {
225                    let field_name = field
226                        .child_by_field_name("name")
227                        .map(|n| node_text(&n, src).to_string())
228                        .unwrap_or_default();
229                    if field_name.is_empty() {
230                        continue; // embedded type
231                    }
232                    let mut field_attrs = FxHashMap::default();
233                    field_attrs.insert(
234                        AttrName::from("name"),
235                        JsonValue::String(field_name.clone()),
236                    );
237                    if let Some(type_node) = field.child_by_field_name("type") {
238                        field_attrs.insert(
239                            AttrName::from("fieldType"),
240                            JsonValue::String(node_text(&type_node, src).to_string()),
241                        );
242                    }
243                    fields.push(make_annotation(
244                        "field",
245                        field_name,
246                        field_attrs,
247                        file,
248                        vec![],
249                    ));
250                }
251            }
252        }
253    }
254    fields
255}
256
257/// Extract interface method signatures as children.
258fn extract_interface_methods(
259    node: &tree_sitter::Node,
260    src: &[u8],
261    file: &RelativePath,
262) -> Vec<Annotation> {
263    let mut methods = Vec::new();
264    let mut cursor = node.walk();
265    for child in node.named_children(&mut cursor) {
266        if child.kind() == "method_elem" {
267            let method_name = child
268                .child_by_field_name("name")
269                .map(|n| node_text(&n, src).to_string())
270                .unwrap_or_default();
271            if method_name.is_empty() {
272                continue;
273            }
274            let mut method_attrs = FxHashMap::default();
275            method_attrs.insert(
276                AttrName::from("name"),
277                JsonValue::String(method_name.clone()),
278            );
279            methods.push(make_annotation(
280                "method",
281                method_name,
282                method_attrs,
283                file,
284                vec![],
285            ));
286        }
287    }
288    methods
289}
290
291fn extract_const_spec(
292    node: &tree_sitter::Node,
293    src: &[u8],
294    file: &RelativePath,
295) -> Option<Annotation> {
296    let name = get_name(node, src);
297    if name.is_empty() {
298        return None;
299    }
300    let mut attrs = FxHashMap::default();
301    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
302    if is_exported(&name) {
303        attrs.insert(AttrName::from("exported"), JsonValue::Bool(true));
304    }
305    Some(make_annotation("const", name, attrs, file, vec![]))
306}
307
308fn extract_var_spec(
309    node: &tree_sitter::Node,
310    src: &[u8],
311    file: &RelativePath,
312) -> Option<Annotation> {
313    let name = get_name(node, src);
314    if name.is_empty() {
315        return None;
316    }
317    let mut attrs = FxHashMap::default();
318    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
319    if is_exported(&name) {
320        attrs.insert(AttrName::from("exported"), JsonValue::Bool(true));
321    }
322    Some(make_annotation("var", name, attrs, file, vec![]))
323}
324
325// ---------------------------------------------------------------------------
326// Parser cache (thread-local)
327// ---------------------------------------------------------------------------
328
329thread_local! {
330    static GO_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
331}
332
333fn parse_go(source: &str) -> Option<tree_sitter::Tree> {
334    GO_PARSER.with(|cell| {
335        let mut opt = cell.borrow_mut();
336        let parser = opt.get_or_insert_with(|| {
337            let mut p = tree_sitter::Parser::new();
338            p.set_language(&tree_sitter_go::LANGUAGE.into())
339                .expect("Failed to set Go language");
340            p
341        });
342        parser.parse(source, None)
343    })
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349
350    fn run(source: &str) -> Vec<Annotation> {
351        let file = RelativePath::from("main.go");
352        GoStructureExtractor.extract(source, &file)
353    }
354
355    #[test]
356    fn extracts_functions() {
357        // Arrange
358        let source = "package main\n\nfunc hello() {}\nfunc World() string { return \"\" }";
359
360        // Act
361        let anns = run(source);
362
363        // Assert
364        assert_eq!(anns.len(), 2, "should find 2 functions");
365        assert_eq!(anns[0].tag.as_ref(), "function", "should be function");
366        assert_eq!(anns[0].binding.as_ref(), "hello", "function name");
367        assert_eq!(
368            anns[0].attrs.get(&AttrName::from("exported")),
369            None,
370            "hello should not be exported"
371        );
372        assert_eq!(
373            anns[1].attrs.get(&AttrName::from("exported")),
374            Some(&JsonValue::Bool(true)),
375            "World should be exported"
376        );
377    }
378
379    #[test]
380    fn extracts_methods() {
381        // Arrange
382        let source = "package main\n\ntype Server struct{}\n\nfunc (s *Server) Handle() {}";
383
384        // Act
385        let anns = run(source);
386        let method = anns.iter().find(|a| a.tag.as_ref() == "method");
387
388        // Assert
389        assert!(method.is_some(), "should find a method");
390        let method = method.unwrap();
391        assert_eq!(method.binding.as_ref(), "Handle", "method name");
392        assert_eq!(
393            method.attrs.get(&AttrName::from("receiver")),
394            Some(&JsonValue::String("Server".to_string())),
395            "receiver type"
396        );
397    }
398
399    #[test]
400    fn extracts_structs_and_interfaces() {
401        // Arrange
402        let source = "package main\n\ntype Config struct {\n\tHost string\n\tPort int\n}\n\ntype Handler interface {\n\tServeHTTP()\n}";
403
404        // Act
405        let anns = run(source);
406
407        // Assert
408        assert_eq!(anns.len(), 2, "should find struct + interface");
409        assert_eq!(anns[0].tag.as_ref(), "struct", "should be struct");
410        assert_eq!(anns[0].binding.as_ref(), "Config", "struct name");
411        assert_eq!(anns[0].children.len(), 2, "struct should have 2 fields");
412        assert_eq!(anns[1].tag.as_ref(), "interface", "should be interface");
413        assert_eq!(anns[1].binding.as_ref(), "Handler", "interface name");
414        assert_eq!(anns[1].children.len(), 1, "interface should have 1 method");
415    }
416
417    #[test]
418    fn extracts_consts_and_vars() {
419        // Arrange
420        let source = "package main\n\nconst MaxRetries = 3\n\nvar defaultTimeout = 30";
421
422        // Act
423        let anns = run(source);
424
425        // Assert
426        assert_eq!(anns.len(), 2, "should find const + var");
427        assert_eq!(anns[0].tag.as_ref(), "const", "should be const");
428        assert_eq!(anns[0].binding.as_ref(), "MaxRetries", "const name");
429        assert_eq!(
430            anns[0].attrs.get(&AttrName::from("exported")),
431            Some(&JsonValue::Bool(true)),
432            "MaxRetries should be exported"
433        );
434        assert_eq!(anns[1].tag.as_ref(), "var", "should be var");
435        assert_eq!(anns[1].binding.as_ref(), "defaultTimeout", "var name");
436        assert_eq!(
437            anns[1].attrs.get(&AttrName::from("exported")),
438            None,
439            "defaultTimeout should not be exported"
440        );
441    }
442}
amql_engine/extractor/go_structure.rs

amql_engine/extractor/
go_structure.rs
