amql_engine/resolver/

go.rs

//! Go source file resolver using tree-sitter.
//!
//! Parses `.go` files into `CodeElement` trees, extracting functions, methods,
//! structs, interfaces, type aliases, constants, and variables.

use super::{CodeElement, SourceLocation};
use crate::error::AqlError;
use crate::types::{AttrName, CodeElementName, RelativePath, TagName};
use rustc_hash::FxHashMap;
use std::cell::RefCell;
use std::path::Path;

/// Go source file resolver using tree-sitter.
pub struct GoResolver;

impl super::CodeResolver for GoResolver {
    fn resolve(&self, file_path: &Path) -> Result<CodeElement, AqlError> {
        let source =
            std::fs::read_to_string(file_path).map_err(|e| format!("Failed to read file: {e}"))?;
        let root = parse_go_source(&source, file_path)?;
        Ok(root)
    }

    fn extensions(&self) -> &[&str] {
        &[".go"]
    }

    fn code_tags(&self) -> &[&str] {
        &[
            "function",
            "method",
            "struct",
            "interface",
            "type",
            "const",
            "var",
            "module",
        ]
    }
}

// Thread-local cached tree-sitter parser to avoid re-creating on each call.
thread_local! {
    static GO_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
}

fn with_go_parser<F, R>(f: F) -> Result<R, String>
where
    F: FnOnce(&mut tree_sitter::Parser) -> Result<R, String>,
{
    GO_PARSER.with(|cell| {
        let mut opt = cell.borrow_mut();
        let parser = opt.get_or_insert_with(|| {
            let mut p = tree_sitter::Parser::new();
            p.set_language(&tree_sitter_go::LANGUAGE.into())
                .expect("Failed to set Go language for tree-sitter");
            p
        });
        f(parser)
    })
}

/// Parse a Go source string into a CodeElement tree.
fn parse_go_source(source: &str, file_path: &Path) -> Result<CodeElement, String> {
    let tree = with_go_parser(|parser| {
        parser
            .parse(source, None)
            .ok_or_else(|| "Failed to parse source".to_string())
    })?;

    let root_node = tree.root_node();
    let src = source.as_bytes();
    let file_str = file_path.to_string_lossy().to_string();

    let mut children = Vec::new();
    let mut cursor = root_node.walk();
    for child in root_node.named_children(&mut cursor) {
        extract_elements(&child, src, &file_str, &mut children);
    }

    let filename = file_path
        .file_name()
        .map(|f| f.to_string_lossy().to_string())
        .unwrap_or_else(|| file_str.clone());

    Ok(CodeElement {
        tag: TagName::from("module"),
        name: CodeElementName::from(filename),
        attrs: FxHashMap::default(),
        children,
        source: SourceLocation {
            file: RelativePath::from(file_str),
            line: 1,
            column: 0,
            end_line: Some(root_node.end_position().row + 1),
            end_column: Some(root_node.end_position().column),
            start_byte: root_node.start_byte(),
            end_byte: root_node.end_byte(),
        },
    })
}

/// Extract CodeElements from a tree-sitter node and its children.
fn extract_elements(
    node: &tree_sitter::Node,
    src: &[u8],
    file: &str,
    result: &mut Vec<CodeElement>,
) {
    match node.kind() {
        "function_declaration" => {
            result.push(extract_function(node, src, file));
        }
        "method_declaration" => {
            result.push(extract_method(node, src, file));
        }
        "type_declaration" => {
            // type_declaration wraps one or more type_spec children
            let mut cursor = node.walk();
            for child in node.named_children(&mut cursor) {
                if child.kind() == "type_spec" {
                    if let Some(el) = extract_type_spec(&child, src, file) {
                        result.push(el);
                    }
                }
            }
        }
        "const_declaration" => {
            let mut cursor = node.walk();
            for child in node.named_children(&mut cursor) {
                if child.kind() == "const_spec" {
                    result.push(extract_const_spec(&child, src, file));
                }
            }
        }
        "var_declaration" => {
            let mut cursor = node.walk();
            for child in node.named_children(&mut cursor) {
                if child.kind() == "var_spec" {
                    result.push(extract_var_spec(&child, src, file));
                }
            }
        }
        _ => {}
    }
}

fn node_text<'a>(node: &tree_sitter::Node, src: &'a [u8]) -> &'a str {
    node.utf8_text(src).unwrap_or("")
}

fn get_name(node: &tree_sitter::Node, src: &[u8]) -> CodeElementName {
    CodeElementName::from(
        node.child_by_field_name("name")
            .map(|n| node_text(&n, src).to_string())
            .unwrap_or_default(),
    )
}

fn make_source_location(node: &tree_sitter::Node, file: &str) -> SourceLocation {
    let start = node.start_position();
    let end = node.end_position();
    SourceLocation {
        file: RelativePath::from(file),
        line: start.row + 1,
        column: start.column,
        end_line: Some(end.row + 1),
        end_column: Some(end.column),
        start_byte: node.start_byte(),
        end_byte: node.end_byte(),
    }
}

/// Extract a function declaration.
fn extract_function(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    // Check if exported (starts with uppercase)
    let name_str = name.to_string();
    if name_str.starts_with(|c: char| c.is_uppercase()) {
        attrs.insert(AttrName::from("export"), serde_json::Value::Bool(true));
    }

    // Extract return type
    if let Some(result_node) = node.child_by_field_name("result") {
        attrs.insert(
            AttrName::from("returnType"),
            serde_json::Value::String(node_text(&result_node, src).to_string()),
        );
    }

    CodeElement {
        tag: TagName::from("function"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

/// Extract a method declaration (function with receiver).
fn extract_method(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    // Extract receiver type
    if let Some(receiver) = node.child_by_field_name("receiver") {
        let receiver_text = node_text(&receiver, src);
        // Strip parens and pointer: (s *Server) -> Server
        let cleaned = receiver_text
            .trim_matches(|c: char| c == '(' || c == ')')
            .trim();
        let type_name = cleaned
            .split_whitespace()
            .last()
            .unwrap_or(cleaned)
            .trim_start_matches('*');
        attrs.insert(
            AttrName::from("receiver"),
            serde_json::Value::String(type_name.to_string()),
        );

        // Check for pointer receiver
        if cleaned.contains('*') {
            attrs.insert(AttrName::from("pointer"), serde_json::Value::Bool(true));
        }
    }

    // Check if exported
    let name_str = name.to_string();
    if name_str.starts_with(|c: char| c.is_uppercase()) {
        attrs.insert(AttrName::from("export"), serde_json::Value::Bool(true));
    }

    // Extract return type
    if let Some(result_node) = node.child_by_field_name("result") {
        attrs.insert(
            AttrName::from("returnType"),
            serde_json::Value::String(node_text(&result_node, src).to_string()),
        );
    }

    CodeElement {
        tag: TagName::from("method"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

/// Extract a type_spec — could be struct, interface, or type alias.
fn extract_type_spec(node: &tree_sitter::Node, src: &[u8], file: &str) -> Option<CodeElement> {
    let name = get_name(node, src);
    let type_node = node.child_by_field_name("type")?;
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    // Check if exported
    let name_str = name.to_string();
    if name_str.starts_with(|c: char| c.is_uppercase()) {
        attrs.insert(AttrName::from("export"), serde_json::Value::Bool(true));
    }

    let (tag, children) = match type_node.kind() {
        "struct_type" => {
            let children = extract_struct_fields(&type_node, src, file);
            ("struct", children)
        }
        "interface_type" => {
            let children = extract_interface_methods(&type_node, src, file);
            ("interface", children)
        }
        _ => ("type", vec![]),
    };

    Some(CodeElement {
        tag: TagName::from(tag),
        name,
        attrs,
        children,
        source: make_source_location(node, file),
    })
}

/// Extract struct field declarations as children.
fn extract_struct_fields(node: &tree_sitter::Node, src: &[u8], file: &str) -> Vec<CodeElement> {
    let mut fields = Vec::new();
    // field_declaration_list is a direct child of struct_type
    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        if child.kind() == "field_declaration_list" {
            let mut inner_cursor = child.walk();
            for field in child.named_children(&mut inner_cursor) {
                if field.kind() == "field_declaration" {
                    // Field name is the first identifier child
                    let field_name = field
                        .child_by_field_name("name")
                        .map(|n| node_text(&n, src).to_string())
                        .unwrap_or_default();
                    if field_name.is_empty() {
                        continue; // embedded type
                    }
                    let mut attrs = FxHashMap::default();
                    attrs.insert(
                        AttrName::from("name"),
                        serde_json::Value::String(field_name.clone()),
                    );
                    if let Some(type_node) = field.child_by_field_name("type") {
                        attrs.insert(
                            AttrName::from("fieldType"),
                            serde_json::Value::String(node_text(&type_node, src).to_string()),
                        );
                    }
                    fields.push(CodeElement {
                        tag: TagName::from("field"),
                        name: CodeElementName::from(field_name),
                        attrs,
                        children: vec![],
                        source: make_source_location(&field, file),
                    });
                }
            }
        }
    }
    fields
}

/// Extract interface method signatures as children.
fn extract_interface_methods(node: &tree_sitter::Node, src: &[u8], file: &str) -> Vec<CodeElement> {
    let mut methods = Vec::new();
    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        if child.kind() == "method_elem" {
            let method_name = child
                .child_by_field_name("name")
                .map(|n| node_text(&n, src).to_string())
                .unwrap_or_default();
            if method_name.is_empty() {
                continue;
            }
            let mut attrs = FxHashMap::default();
            attrs.insert(
                AttrName::from("name"),
                serde_json::Value::String(method_name.clone()),
            );
            methods.push(CodeElement {
                tag: TagName::from("method"),
                name: CodeElementName::from(method_name),
                attrs,
                children: vec![],
                source: make_source_location(&child, file),
            });
        }
    }
    methods
}

/// Extract a const_spec.
fn extract_const_spec(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );
    let name_str = name.to_string();
    if name_str.starts_with(|c: char| c.is_uppercase()) {
        attrs.insert(AttrName::from("export"), serde_json::Value::Bool(true));
    }
    CodeElement {
        tag: TagName::from("const"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

/// Extract a var_spec.
fn extract_var_spec(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );
    let name_str = name.to_string();
    if name_str.starts_with(|c: char| c.is_uppercase()) {
        attrs.insert(AttrName::from("export"), serde_json::Value::Bool(true));
    }
    CodeElement {
        tag: TagName::from("var"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn parse_snippet(source: &str) -> CodeElement {
        parse_go_source(source, Path::new("test.go")).unwrap()
    }

    #[test]
    fn parses_functions() {
        // Arrange
        let root =
            parse_snippet("package main\n\nfunc hello() {}\nfunc World() string { return \"\" }");

        // Act
        let hello = &root.children[0];
        let world = &root.children[1];

        // Assert
        assert_eq!(hello.tag, "function", "hello tag");
        assert_eq!(hello.name, "hello", "hello name");
        assert_eq!(hello.attrs.get("export"), None, "hello not exported");

        assert_eq!(world.tag, "function", "World tag");
        assert_eq!(world.name, "World", "World name");
        assert_eq!(
            world.attrs.get("export"),
            Some(&serde_json::Value::Bool(true)),
            "World is exported"
        );
    }

    #[test]
    fn parses_methods() {
        // Arrange
        let root =
            parse_snippet("package main\n\ntype Server struct{}\n\nfunc (s *Server) Handle() {}");

        // Act — type_declaration + method_declaration
        let method = root
            .children
            .iter()
            .find(|c| c.tag.as_ref() == "method")
            .unwrap();

        // Assert
        assert_eq!(method.name, "Handle", "method name");
        assert_eq!(
            method.attrs.get("receiver"),
            Some(&serde_json::Value::String("Server".to_string())),
            "receiver type"
        );
        assert_eq!(
            method.attrs.get("pointer"),
            Some(&serde_json::Value::Bool(true)),
            "pointer receiver"
        );
    }

    #[test]
    fn parses_structs_and_interfaces() {
        // Arrange
        let root = parse_snippet(
            "package main\n\ntype Config struct {\n\tHost string\n\tPort int\n}\n\ntype Handler interface {\n\tServeHTTP()\n}",
        );

        // Act
        let config = &root.children[0];
        let handler = &root.children[1];

        // Assert
        assert_eq!(config.tag, "struct", "Config tag");
        assert_eq!(config.name, "Config", "Config name");
        assert_eq!(config.children.len(), 2, "Config has 2 fields");
        assert_eq!(config.children[0].name, "Host", "first field");

        assert_eq!(handler.tag, "interface", "Handler tag");
        assert_eq!(handler.name, "Handler", "Handler name");
        assert_eq!(handler.children.len(), 1, "Handler has 1 method");
        assert_eq!(handler.children[0].name, "ServeHTTP", "interface method");
    }

    #[test]
    fn parses_consts_and_vars() {
        // Arrange
        let root = parse_snippet("package main\n\nconst MaxRetries = 3\n\nvar defaultTimeout = 30");

        // Act
        let c = &root.children[0];
        let v = &root.children[1];

        // Assert
        assert_eq!(c.tag, "const", "const tag");
        assert_eq!(c.name, "MaxRetries", "const name");
        assert_eq!(
            c.attrs.get("export"),
            Some(&serde_json::Value::Bool(true)),
            "const exported"
        );

        assert_eq!(v.tag, "var", "var tag");
        assert_eq!(v.name, "defaultTimeout", "var name");
        assert_eq!(v.attrs.get("export"), None, "var not exported");
    }
}