memlay 0.1.4

//! Structural code parsing via tree-sitter (PRD §11). One generic walker
//! driven by per-language specs; unsupported languages degrade safely to
//! file-level lexical indexing. Extraction results are cached content-
//! addressed in the shared repository cache and reused across worktrees.

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::path::Path;
use tree_sitter::{Language, Node, Parser};

/// Bump when extraction logic changes; part of every cache key (PRD §7.3).
pub const PARSER_VERSION: u32 = 1;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SymbolRec {
    pub name: String,
    pub qualified_name: String,
    pub kind: String,
    pub start_line: u32,
    pub end_line: u32,
    pub signature: String,
    pub is_test: bool,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct Extraction {
    pub symbols: Vec<SymbolRec>,
    /// Import targets: module path / source strings as written.
    pub imports: Vec<String>,
}

/// (node kind, symbol kind) pairs whose `name` field (or first identifier
/// child) names a symbol.
struct LangSpec {
    symbols: &'static [(&'static str, &'static str)],
    /// Node kinds that establish a qualified-name scope for their children.
    containers: &'static [&'static str],
    /// Node kinds representing imports; the first string/dotted child is the
    /// import target.
    imports: &'static [&'static str],
}

const TS_SPEC: LangSpec = LangSpec {
    symbols: &[
        ("function_declaration", "function"),
        ("generator_function_declaration", "function"),
        ("method_definition", "method"),
        ("class_declaration", "class"),
        ("abstract_class_declaration", "class"),
        ("interface_declaration", "interface"),
        ("enum_declaration", "enum"),
        ("type_alias_declaration", "type-alias"),
        ("variable_declarator", "function"), // only kept when value is a function
    ],
    containers: &[
        "class_declaration",
        "abstract_class_declaration",
        "interface_declaration",
        "enum_declaration",
        "internal_module",
    ],
    imports: &["import_statement"],
};

const PY_SPEC: LangSpec = LangSpec {
    symbols: &[
        ("function_definition", "function"),
        ("class_definition", "class"),
    ],
    containers: &["class_definition", "function_definition"],
    imports: &["import_statement", "import_from_statement"],
};

const RUST_SPEC: LangSpec = LangSpec {
    symbols: &[
        ("function_item", "function"),
        ("struct_item", "struct"),
        ("enum_item", "enum"),
        ("trait_item", "trait"),
        ("type_item", "type-alias"),
        ("const_item", "const"),
        ("static_item", "static"),
        ("mod_item", "module"),
    ],
    containers: &["mod_item", "impl_item", "trait_item"],
    imports: &["use_declaration"],
};

const GO_SPEC: LangSpec = LangSpec {
    symbols: &[
        ("function_declaration", "function"),
        ("method_declaration", "method"),
        ("type_spec", "type"),
    ],
    containers: &[],
    imports: &["import_spec"],
};

const LUA_SPEC: LangSpec = LangSpec {
    symbols: &[("function_declaration", "function")],
    containers: &[],
    imports: &[],
};

fn language_for(lang: &str) -> Option<(Language, &'static LangSpec)> {
    match lang {
        "typescript" => Some((tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(), &TS_SPEC)),
        "tsx" => Some((tree_sitter_typescript::LANGUAGE_TSX.into(), &TS_SPEC)),
        "javascript" | "jsx" => Some((tree_sitter_javascript::LANGUAGE.into(), &TS_SPEC)),
        "python" => Some((tree_sitter_python::LANGUAGE.into(), &PY_SPEC)),
        "rust" => Some((tree_sitter_rust::LANGUAGE.into(), &RUST_SPEC)),
        "go" => Some((tree_sitter_go::LANGUAGE.into(), &GO_SPEC)),
        "lua" => Some((tree_sitter_lua::LANGUAGE.into(), &LUA_SPEC)),
        _ => None,
    }
}

pub fn is_supported(lang: &str) -> bool {
    language_for(lang).is_some()
}

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

fn name_of(node: Node, source: &[u8]) -> Option<String> {
    if let Some(n) = node.child_by_field_name("name") {
        return Some(node_text(n, source).to_string());
    }
    // Fallbacks: declarator field (Go type_spec uses name; Rust impl uses type).
    for field in ["declarator", "type"] {
        if let Some(n) = node.child_by_field_name(field) {
            return Some(node_text(n, source).to_string());
        }
    }
    // First identifier-ish child.
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        let k = child.kind();
        if k.contains("identifier") || k == "dot_index_expression" || k == "method_index_expression"
        {
            return Some(node_text(child, source).to_string());
        }
    }
    None
}

fn first_line_signature(node: Node, source: &[u8]) -> String {
    let text = node_text(node, source);
    let line = text.lines().next().unwrap_or("").trim();
    let mut sig: String = line.chars().take(200).collect();
    if sig.len() < line.len() {
        sig.push('…');
    }
    sig
}

fn import_target(node: Node, source: &[u8]) -> Option<String> {
    // Prefer explicit fields used by the grammars we ship.
    for field in ["source", "path", "module_name", "argument"] {
        if let Some(n) = node.child_by_field_name(field) {
            return Some(
                node_text(n, source)
                    .trim_matches(['"', '\'', '`'])
                    .to_string(),
            );
        }
    }
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        match child.kind() {
            "string" | "interpreted_string_literal" | "string_literal" => {
                return Some(
                    node_text(child, source)
                        .trim_matches(['"', '\'', '`'])
                        .to_string(),
                );
            }
            "dotted_name" | "scoped_identifier" | "use_wildcard" | "scoped_use_list"
            | "use_as_clause" | "identifier" => {
                return Some(node_text(child, source).to_string());
            }
            _ => {}
        }
    }
    None
}

fn is_test_symbol(lang: &str, path: &str, name: &str, node: Node, source: &[u8]) -> bool {
    let path_says_test = path.contains(".test.")
        || path.contains(".spec.")
        || path.ends_with("_test.go")
        || path.contains("/tests/")
        || path.starts_with("tests/");
    match lang {
        "python" => name.starts_with("test_") || path_says_test,
        "go" => name.starts_with("Test") && path.ends_with("_test.go"),
        "rust" => {
            // A #[test]-attributed sibling immediately above the function.
            let mut prev = node.prev_sibling();
            while let Some(p) = prev {
                if p.kind() == "attribute_item" {
                    if node_text(p, source).contains("test") {
                        return true;
                    }
                    prev = p.prev_sibling();
                } else {
                    break;
                }
            }
            false
        }
        _ => path_says_test,
    }
}

/// Parse source and extract symbols/imports. Returns None when the language
/// has no structural parser (callers fall back to lexical indexing).
pub fn extract(lang: &str, path: &str, source: &[u8]) -> Option<Result<Extraction>> {
    let (language, spec) = language_for(lang)?;
    Some(extract_with(language, spec, lang, path, source))
}

fn extract_with(
    language: Language,
    spec: &LangSpec,
    lang: &str,
    path: &str,
    source: &[u8],
) -> Result<Extraction> {
    let mut parser = Parser::new();
    parser.set_language(&language)?;
    let tree = parser
        .parse(source, None)
        .ok_or_else(|| anyhow::anyhow!("parser returned no tree"))?;

    let mut out = Extraction::default();
    // Iterative DFS carrying the container-name stack.
    let mut stack: Vec<(Node, Vec<String>)> = vec![(tree.root_node(), Vec::new())];
    while let Some((node, scope)) = stack.pop() {
        let kind = node.kind();

        if spec.imports.contains(&kind) {
            if let Some(target) = import_target(node, source) {
                if !target.is_empty() {
                    out.imports.push(target);
                }
            }
        }

        if let Some((_, sym_kind)) = spec.symbols.iter().find(|(k, _)| *k == kind) {
            let mut keep = true;
            // TS/JS: `const f = () => {}` — only variable declarators whose
            // value is a function count as symbols.
            if kind == "variable_declarator" {
                keep = node
                    .child_by_field_name("value")
                    .map(|v| {
                        matches!(
                            v.kind(),
                            "arrow_function" | "function_expression" | "function"
                        )
                    })
                    .unwrap_or(false);
            }
            if keep {
                if let Some(name) = name_of(node, source) {
                    if !name.is_empty() {
                        let qualified = if scope.is_empty() {
                            name.clone()
                        } else {
                            format!("{}.{}", scope.join("."), name)
                        };
                        out.symbols.push(SymbolRec {
                            is_test: is_test_symbol(lang, path, &name, node, source),
                            name,
                            qualified_name: qualified,
                            kind: sym_kind.to_string(),
                            start_line: node.start_position().row as u32 + 1,
                            end_line: node.end_position().row as u32 + 1,
                            signature: first_line_signature(node, source),
                        });
                    }
                }
            }
        }

        let child_scope = if spec.containers.contains(&kind) {
            let mut s = scope.clone();
            if let Some(name) = name_of(node, source) {
                if !name.is_empty() {
                    s.push(name);
                }
            }
            s
        } else {
            scope
        };

        let mut cursor = node.walk();
        // Push in reverse so iteration order matches document order.
        let children: Vec<Node> = node.children(&mut cursor).collect();
        for child in children.into_iter().rev() {
            stack.push((child, child_scope.clone()));
        }
    }

    // Deterministic order: by line, then qualified name.
    out.symbols
        .sort_by(|a, b| (a.start_line, &a.qualified_name).cmp(&(b.start_line, &b.qualified_name)));
    out.imports.sort();
    out.imports.dedup();
    Ok(out)
}

// ------------------------------------------------------- parse cache ----

/// Content-addressed cache in the shared git-common-dir (PRD §7.3), keyed by
/// parser version, language, and content hash.
pub struct ParseCache {
    dir: std::path::PathBuf,
}

impl ParseCache {
    pub fn new(shared_dir: &Path) -> ParseCache {
        ParseCache {
            dir: shared_dir.join("parse-cache"),
        }
    }

    fn key(&self, lang: &str, content_hash: &str) -> std::path::PathBuf {
        let key = blake3::hash(format!("{PARSER_VERSION}|{lang}|{content_hash}").as_bytes())
            .to_hex()
            .to_string();
        self.dir.join(&key[..2]).join(format!("{key}.json"))
    }

    pub fn get(&self, lang: &str, content_hash: &str) -> Option<Extraction> {
        let path = self.key(lang, content_hash);
        let bytes = std::fs::read(path).ok()?;
        serde_json::from_slice(&bytes).ok()
    }

    pub fn put(&self, lang: &str, content_hash: &str, extraction: &Extraction) {
        let path = self.key(lang, content_hash);
        if let Some(parent) = path.parent() {
            if std::fs::create_dir_all(parent).is_ok() {
                let tmp = path.with_extension("tmp");
                if let Ok(json) = serde_json::to_vec(extraction) {
                    if std::fs::write(&tmp, json).is_ok() {
                        let _ = std::fs::rename(&tmp, &path);
                    }
                }
            }
        }
    }
}

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

    #[test]
    fn typescript_symbols_and_imports() {
        let src = br#"
import { Router } from "express";
export class PaymentWebhookController {
  handle(req: Request): void {}
}
export function helper(): number { return 1; }
const arrowFn = (x: number) => x * 2;
const notAFunction = 42;
"#;
        let e = extract("typescript", "apps/api/webhook.ts", src)
            .unwrap()
            .unwrap();
        let names: Vec<&str> = e
            .symbols
            .iter()
            .map(|s| s.qualified_name.as_str())
            .collect();
        assert!(names.contains(&"PaymentWebhookController"), "{names:?}");
        assert!(
            names.contains(&"PaymentWebhookController.handle"),
            "{names:?}"
        );
        assert!(names.contains(&"helper"), "{names:?}");
        assert!(names.contains(&"arrowFn"), "{names:?}");
        assert!(!names.contains(&"notAFunction"), "{names:?}");
        assert_eq!(e.imports, vec!["express"]);
    }

    #[test]
    fn python_symbols() {
        let src = br#"
import os
from payments import retry

class Consumer:
    def process(self, event):
        pass

def test_processes_event():
    pass
"#;
        let e = extract("python", "worker/consumer.py", src)
            .unwrap()
            .unwrap();
        let names: Vec<&str> = e
            .symbols
            .iter()
            .map(|s| s.qualified_name.as_str())
            .collect();
        assert!(names.contains(&"Consumer"), "{names:?}");
        assert!(names.contains(&"Consumer.process"), "{names:?}");
        let test = e
            .symbols
            .iter()
            .find(|s| s.name == "test_processes_event")
            .unwrap();
        assert!(test.is_test);
        assert!(e.imports.iter().any(|i| i == "os"), "{:?}", e.imports);
        assert!(e.imports.iter().any(|i| i == "payments"), "{:?}", e.imports);
    }

    #[test]
    fn rust_symbols_and_test_attr() {
        let src = br#"
use std::collections::HashMap;

pub struct Engine;

impl Engine {
    pub fn run(&self) {}
}

mod tests {
    #[test]
    fn engine_runs() {}
}
"#;
        let e = extract("rust", "src/engine.rs", src).unwrap().unwrap();
        let names: Vec<&str> = e
            .symbols
            .iter()
            .map(|s| s.qualified_name.as_str())
            .collect();
        assert!(names.contains(&"Engine"), "{names:?}");
        assert!(names.contains(&"Engine.run"), "{names:?}");
        assert!(names.contains(&"tests.engine_runs"), "{names:?}");
        let t = e.symbols.iter().find(|s| s.name == "engine_runs").unwrap();
        assert!(t.is_test);
        assert!(
            e.imports.iter().any(|i| i.contains("HashMap")),
            "{:?}",
            e.imports
        );
    }

    #[test]
    fn go_symbols() {
        let src = br#"
package payments

import "fmt"

type Processor struct{}

func (p *Processor) Handle() {}

func NewProcessor() *Processor { return nil }
"#;
        let e = extract("go", "payments/processor.go", src)
            .unwrap()
            .unwrap();
        let names: Vec<&str> = e.symbols.iter().map(|s| s.name.as_str()).collect();
        assert!(names.contains(&"Processor"), "{names:?}");
        assert!(names.contains(&"Handle"), "{names:?}");
        assert!(names.contains(&"NewProcessor"), "{names:?}");
        assert_eq!(e.imports, vec!["fmt"]);
    }

    #[test]
    fn lua_symbols() {
        let src = br#"
local function helper()
end

function M.process(event)
end
"#;
        let e = extract("lua", "scripts/mod.lua", src).unwrap().unwrap();
        assert!(!e.symbols.is_empty(), "{:?}", e.symbols);
    }

    #[test]
    fn unsupported_language_returns_none() {
        assert!(extract("markdown", "README.md", b"# hi").is_none());
    }

    #[test]
    fn cache_round_trip() {
        let tmp = tempfile::tempdir().unwrap();
        let cache = ParseCache::new(tmp.path());
        let e = Extraction {
            symbols: vec![],
            imports: vec!["x".into()],
        };
        assert!(cache.get("rust", "abc").is_none());
        cache.put("rust", "abc", &e);
        assert_eq!(cache.get("rust", "abc"), Some(e));
    }

    #[test]
    fn broken_source_does_not_panic() {
        let src = b"class {{{{ def )))) import";
        // Tree-sitter always produces a tree; extraction must not fail.
        let e = extract("python", "x.py", src).unwrap();
        assert!(e.is_ok());
    }
}