lean_ctx/core/

chunks_ts.rs

//! Tree-sitter AST-aware code chunking for semantic search.
//!
//! Replaces heuristic line-prefix matching with proper AST parsing.
//! Extracts function bodies, struct definitions, class declarations etc.
//! as complete, self-contained chunks with accurate boundaries.
//!
//! Falls back to heuristic chunking for unsupported languages.

#[cfg(feature = "tree-sitter")]
use tree_sitter::{Node, Parser, Query, QueryCursor, StreamingIterator};

use super::bm25_index::{ChunkKind, CodeChunk};

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_RUST: &str = r"
(function_item name: (identifier) @name) @chunk
(struct_item name: (type_identifier) @name) @chunk
(enum_item name: (type_identifier) @name) @chunk
(trait_item name: (type_identifier) @name) @chunk
(impl_item type: (type_identifier) @name) @chunk
(const_item name: (identifier) @name) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_TYPESCRIPT: &str = r"
(function_declaration name: (identifier) @name) @chunk
(class_declaration name: (type_identifier) @name) @chunk
(abstract_class_declaration name: (type_identifier) @name) @chunk
(interface_declaration name: (type_identifier) @name) @chunk
(type_alias_declaration name: (type_identifier) @name) @chunk
(method_definition name: (property_identifier) @name) @chunk
(variable_declarator name: (identifier) @name value: (arrow_function)) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_JAVASCRIPT: &str = r"
(function_declaration name: (identifier) @name) @chunk
(class_declaration name: (identifier) @name) @chunk
(method_definition name: (property_identifier) @name) @chunk
(variable_declarator name: (identifier) @name value: (arrow_function)) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_PYTHON: &str = r"
(function_definition name: (identifier) @name) @chunk
(class_definition name: (identifier) @name) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_GO: &str = r"
(function_declaration name: (identifier) @name) @chunk
(method_declaration name: (field_identifier) @name) @chunk
(type_spec name: (type_identifier) @name) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_JAVA: &str = r"
(method_declaration name: (identifier) @name) @chunk
(class_declaration name: (identifier) @name) @chunk
(interface_declaration name: (identifier) @name) @chunk
(enum_declaration name: (identifier) @name) @chunk
(constructor_declaration name: (identifier) @name) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_C: &str = r"
(function_definition
  declarator: (function_declarator
    declarator: (identifier) @name)) @chunk
(struct_specifier name: (type_identifier) @name) @chunk
(enum_specifier name: (type_identifier) @name) @chunk
";

#[cfg(feature = "tree-sitter")]
const CHUNK_QUERY_CPP: &str = r"
(function_definition
  declarator: (function_declarator
    declarator: (_) @name)) @chunk
(struct_specifier name: (type_identifier) @name) @chunk
(class_specifier name: (type_identifier) @name) @chunk
(enum_specifier name: (type_identifier) @name) @chunk
(namespace_definition name: (identifier) @name) @chunk
";

/// Extract code chunks from a file using tree-sitter AST parsing.
///
/// Returns `None` if the language is unsupported, allowing callers to fall back
/// to heuristic-based chunking.
#[cfg(feature = "tree-sitter")]
fn get_cached_query(file_ext: &str) -> Option<&'static Query> {
    use std::collections::HashMap;
    use std::sync::OnceLock;

    static QUERY_CACHE: OnceLock<HashMap<&'static str, Query>> = OnceLock::new();

    let cache = QUERY_CACHE.get_or_init(|| {
        let mut map = HashMap::new();
        let exts: &[&str] = &[
            "rs", "ts", "tsx", "js", "jsx", "py", "go", "java", "c", "h", "cpp", "cc", "cxx", "hpp",
        ];
        for &ext in exts {
            if let (Some(lang), Some(src)) = (get_language(ext), get_chunk_query(ext))
                && let Ok(q) = Query::new(&lang, src)
            {
                map.insert(ext, q);
            }
        }
        map
    });

    cache.get(file_ext)
}

/// Visit each structural chunk root (`@chunk` capture) once per tree-sitter query match.
///
/// `start_line` / `end_line` are **1-based** inclusive line numbers (matching [`CodeChunk`]).
/// Returns `None` if the extension is unsupported or parsing fails.
#[cfg(feature = "tree-sitter")]
pub(crate) fn for_each_chunk_node(
    content: &str,
    file_ext: &str,
    mut visitor: impl FnMut(Node, &str, ChunkKind, usize, usize),
) -> Option<()> {
    for_each_chunk_node_while(
        content,
        file_ext,
        || true,
        |node, name, kind, start, end| {
            visitor(node, name, kind, start, end);
            true
        },
    )
}

/// Stop capture traversal as soon as the consumer reaches its work budget.
#[cfg(feature = "tree-sitter")]
pub(crate) fn for_each_chunk_node_while(
    content: &str,
    file_ext: &str,
    mut continue_scan: impl FnMut() -> bool,
    mut visitor: impl FnMut(Node, &str, ChunkKind, usize, usize) -> bool,
) -> Option<()> {
    let language = get_language(file_ext)?;

    thread_local! {
        static PARSER: std::cell::RefCell<Parser> = std::cell::RefCell::new(Parser::new());
    }

    let tree = PARSER.with(|p| {
        let mut parser = p.borrow_mut();
        let _ = parser.set_language(&language);
        parser.parse(content, None)
    })?;

    let query = get_cached_query(file_ext)?;
    let chunk_idx = find_capture_index(query, "chunk")?;
    let name_idx = find_capture_index(query, "name")?;

    let source = content.as_bytes();
    let mut cursor = QueryCursor::new();
    let mut matches = cursor.matches(query, tree.root_node(), source);
    let mut seen_ranges = Vec::new();

    while let Some(m) = matches.next() {
        // Budget admission precedes even containment/deduplication work.
        if !continue_scan() {
            break;
        }
        let mut chunk_node: Option<Node> = None;
        let mut name_text = String::new();

        for cap in m.captures {
            if cap.index == chunk_idx {
                chunk_node = Some(cap.node);
            } else if cap.index == name_idx
                && let Ok(text) = cap.node.utf8_text(source)
            {
                name_text = text.to_string();
            }
        }

        if let Some(node) = chunk_node {
            if name_text.is_empty() {
                continue;
            }

            let start_row0 = node.start_position().row;
            let end_row0 = node.end_position().row;

            let range = (start_row0, end_row0);
            if seen_ranges
                .iter()
                .any(|&(s, e)| s <= start_row0 && end_row0 <= e && range != (s, e))
            {
                continue;
            }
            seen_ranges.push(range);

            let kind = node_kind_to_chunk_kind(node.kind());
            if !visitor(node, name_text.as_str(), kind, start_row0 + 1, end_row0 + 1) {
                break;
            }
        }
    }

    Some(())
}

#[cfg(feature = "tree-sitter")]
pub fn extract_chunks_ts(file_path: &str, content: &str, file_ext: &str) -> Option<Vec<CodeChunk>> {
    let mut chunks = Vec::new();

    for_each_chunk_node(
        content,
        file_ext,
        |node, name_text, kind, start_line, end_line| {
            let block = match node.utf8_text(content.as_bytes()) {
                Ok(text) => text.to_owned(),
                Err(_) => return,
            };
            let token_count = super::bm25_index::tokenize_for_index(&block).len();

            chunks.push(CodeChunk {
                input_origin: None,
                file_path: file_path.to_string(),
                symbol_name: name_text.to_string(),
                kind,
                start_line,
                end_line,
                content: block,
                tokens: Vec::new(),
                token_count,
            });
        },
    )?;

    if chunks.is_empty() {
        return None;
    }

    chunks.sort_by_key(|c| c.start_line);
    Some(chunks)
}

#[cfg(not(feature = "tree-sitter"))]
pub fn extract_chunks_ts(
    _file_path: &str,
    _content: &str,
    _file_ext: &str,
) -> Option<Vec<CodeChunk>> {
    None
}

#[cfg(feature = "tree-sitter")]
use lean_ctx_syntax::chunk_grammar::chunk_language as get_language;

#[cfg(feature = "tree-sitter")]
fn get_chunk_query(ext: &str) -> Option<&'static str> {
    Some(match ext {
        "rs" => CHUNK_QUERY_RUST,
        "ts" | "tsx" => CHUNK_QUERY_TYPESCRIPT,
        "js" | "jsx" => CHUNK_QUERY_JAVASCRIPT,
        "py" => CHUNK_QUERY_PYTHON,
        "go" => CHUNK_QUERY_GO,
        "java" => CHUNK_QUERY_JAVA,
        "c" | "h" => CHUNK_QUERY_C,
        "cpp" | "cc" | "cxx" | "hpp" | "hxx" | "hh" => CHUNK_QUERY_CPP,
        _ => return None,
    })
}

#[cfg(feature = "tree-sitter")]
fn find_capture_index(query: &Query, name: &str) -> Option<u32> {
    query
        .capture_names()
        .iter()
        .position(|n| *n == name)
        .map(|i| i as u32)
}

fn node_kind_to_chunk_kind(kind: &str) -> ChunkKind {
    match kind {
        "function_item"
        | "function_declaration"
        | "function_definition"
        | "method_declaration"
        | "method_definition"
        | "constructor_declaration"
        | "variable_declarator" => ChunkKind::Function,

        "struct_item"
        | "struct_specifier"
        | "struct_declaration"
        | "enum_item"
        | "enum_specifier"
        | "enum_declaration"
        | "trait_item"
        | "interface_declaration"
        | "type_alias_declaration"
        | "type_spec" => ChunkKind::Struct,

        "impl_item" => ChunkKind::Impl,

        "class_declaration"
        | "abstract_class_declaration"
        | "class_specifier"
        | "class_definition" => ChunkKind::Class,

        "namespace_definition" | "namespace_declaration" => ChunkKind::Module,

        _ => ChunkKind::Other,
    }
}

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

    #[test]
    fn extract_rust_chunks() {
        let src = r#"use std::io;

pub fn process(input: &str) -> String {
    input.to_uppercase()
}

pub struct Config {
    pub name: String,
    pub port: u16,
}

impl Config {
    pub fn new() -> Self {
        Self { name: "default".into(), port: 8080 }
    }
}

fn helper() -> bool {
    true
}
"#;
        let chunks = extract_chunks_ts("main.rs", src, "rs").unwrap();
        assert!(
            chunks.len() >= 4,
            "expected >=4 chunks, got {}",
            chunks.len()
        );

        let names: Vec<&str> = chunks.iter().map(|c| c.symbol_name.as_str()).collect();
        assert!(names.contains(&"process"), "got {names:?}");
        assert!(names.contains(&"Config"), "got {names:?}");
        assert!(names.contains(&"helper"), "got {names:?}");

        let process = chunks.iter().find(|c| c.symbol_name == "process").unwrap();
        assert!(matches!(process.kind, ChunkKind::Function));
        assert!(process.content.contains("to_uppercase"));
    }

    #[test]
    fn extract_typescript_chunks() {
        let src = r"
export function greet(name: string): string {
    return `Hello ${name}`;
}

export class UserService {
    findUser(id: number): User {
        return db.find(id);
    }
}

const handler = async (req: Request): Promise<Response> => {
    return new Response();
};
";
        let chunks = extract_chunks_ts("app.ts", src, "ts").unwrap();
        assert!(
            chunks.len() >= 3,
            "expected >=3 chunks, got {}",
            chunks.len()
        );

        let names: Vec<&str> = chunks.iter().map(|c| c.symbol_name.as_str()).collect();
        assert!(names.contains(&"greet"), "got {names:?}");
        assert!(names.contains(&"UserService"), "got {names:?}");
    }

    #[test]
    fn extract_python_chunks() {
        let src = r"
class AuthService:
    def __init__(self, db):
        self.db = db

    def authenticate(self, email: str) -> bool:
        user = self.db.find(email)
        return user is not None

def create_app():
    return Flask(__name__)
";
        let chunks = extract_chunks_ts("app.py", src, "py").unwrap();
        assert!(
            chunks.len() >= 2,
            "expected >=2 chunks, got {}",
            chunks.len()
        );

        let names: Vec<&str> = chunks.iter().map(|c| c.symbol_name.as_str()).collect();
        assert!(names.contains(&"AuthService"), "got {names:?}");
        assert!(names.contains(&"create_app"), "got {names:?}");

        let auth = chunks
            .iter()
            .find(|c| c.symbol_name == "AuthService")
            .unwrap();
        assert!(auth.content.contains("authenticate"));
    }

    #[test]
    fn chunks_contain_full_body() {
        let src = r#"
pub fn complex(x: i32, y: i32) -> Result<String, Error> {
    let sum = x + y;
    let result = format!("Sum: {}", sum);
    if sum > 100 {
        return Err(Error::new("too large"));
    }
    Ok(result)
}
"#;
        let chunks = extract_chunks_ts("math.rs", src, "rs").unwrap();
        let complex = chunks.iter().find(|c| c.symbol_name == "complex").unwrap();
        assert!(complex.content.contains("sum > 100"));
        assert!(complex.content.contains("Ok(result)"));
    }

    #[test]
    fn unsupported_language_returns_none() {
        assert!(extract_chunks_ts("file.xyz", "content", "xyz").is_none());
    }

    #[test]
    fn empty_file_returns_none() {
        assert!(extract_chunks_ts("empty.rs", "", "rs").is_none());
    }

    #[test]
    fn javascript_same_line_functions_use_exact_node_content() {
        let src = "function alpha(){return 1;}function beta(){return 2;}";
        let chunks = extract_chunks_ts("min.js", src, "js").unwrap();

        let alpha = chunks.iter().find(|c| c.symbol_name == "alpha").unwrap();
        let beta = chunks.iter().find(|c| c.symbol_name == "beta").unwrap();

        assert_eq!(alpha.content, "function alpha(){return 1;}");
        assert_eq!(beta.content, "function beta(){return 2;}");
    }

    #[test]
    fn chunks_sorted_by_line() {
        let src = r"
fn b_func() {}
fn a_func() {}
";
        let chunks = extract_chunks_ts("sort.rs", src, "rs").unwrap();
        assert!(chunks[0].start_line <= chunks[1].start_line);
    }
}