From 628c82850df0cac0a0cf6c49a255a1ce527f1b50 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Wed, 22 Jul 2026 21:11:12 +0700 Subject: [PATCH 1/2] feat(index): tree-sitter AST extraction + schema v3 edges table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 2 of knowledge graph implementation. ## tree-sitter AST extraction (feature-gated) - Add ast.rs module with tree-sitter-based AST parsing for Rust, Go, Python, TypeScript - Cursor-based DFS traversal (tree-sitter 0.26 has no descendants()) - Kind-based child lookup (grammars lack named fields like name) - TypeScript export_statement unwrapping for exported declarations - Go const/var declaration support via const_spec/var_spec traversal - Wire extract_symbols() and extract_calls() to try tree-sitter first, fall back to regex for unsupported languages - Add extract_edges() for richer edge types (IMPORTS, IMPLEMENTS, INHERITS, CHILD_OF) - 10 new tests covering all 4 language extractors ## Schema v3 — edges table - Add edges table: source to target with typed relationship (CALLS, IMPORTS, IMPLEMENTS, INHERITS, CHILD_OF) - Migrate existing call_graph data into edges with kind=CALLS - Populate edges table during indexing when tree-sitter feature enabled - Add KgEdge, store_kg_edges(), clear_kg_edges_for_file() to graph.rs - 2 new schema migration tests ## Feature flag - All tree-sitter code behind --features tree-sitter - Default build has zero tree-sitter overhead - 4 language grammars statically linked (~454KB release binary) Test results: - Default: 705 pass, 0 fail - With tree-sitter: 714 pass, 0 fail (+9 new tests) - Clippy: 0 warnings both modes --- Cargo.lock | 72 ++++ Cargo.toml | 17 + src/index/ast.rs | 997 +++++++++++++++++++++++++++++++++++++++++++ src/index/extract.rs | 53 ++- src/index/graph.rs | 50 +++ src/index/mod.rs | 16 +- src/index/schema.rs | 106 ++++- 7 files changed, 1304 insertions(+), 7 deletions(-) create mode 100644 src/index/ast.rs diff --git a/Cargo.lock b/Cargo.lock index b9f346f..3a90fcf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -315,6 +315,11 @@ dependencies = [ "toml", "tracing", "tracing-subscriber", + "tree-sitter", + "tree-sitter-go", + "tree-sitter-python", + "tree-sitter-rust", + "tree-sitter-typescript", "which", ] @@ -1498,6 +1503,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ + "indexmap", "itoa", "memchr", "serde", @@ -1593,6 +1599,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b2231b7c3057d5e4ad0156fb3dc807d900806020c5ffa3ee6ff2c8c76fb8520" + [[package]] name = "strsim" version = "0.11.1" @@ -1938,6 +1950,66 @@ dependencies = [ "tracing-log", ] +[[package]] +name = "tree-sitter" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af1c71c1c4cc0920b20d6b0f6572e7682cd07a6a2faec71067a31fa394c586df" +dependencies = [ + "cc", + "regex", + "regex-syntax", + "serde_json", + "streaming-iterator", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-go" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8560a4d2f835cc0d4d2c2e03cbd0dde2f6114b43bc491164238d333e28b16ea" +dependencies = [ + "cc", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-language" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "009994f150cc0cd50ff54917d5bc8bffe8cad10ca10d81c34da2ec421ae61782" + +[[package]] +name = "tree-sitter-python" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bf85fd39652e740bf60f46f4cda9492c3a9ad75880575bf14960f775cb74a1c" +dependencies = [ + "cc", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-rust" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439e577dbe07423ec2582ac62c7531120dbfccfa6e5f92406f93dd271a120e45" +dependencies = [ + "cc", + "tree-sitter-language", +] + +[[package]] +name = "tree-sitter-typescript" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c5f76ed8d947a75cc446d5fccd8b602ebf0cde64ccf2ffa434d873d7a575eff" +dependencies = [ + "cc", + "tree-sitter-language", +] + [[package]] name = "try-lock" version = "0.2.5" diff --git a/Cargo.toml b/Cargo.toml index 42b5048..a7fdcff 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,6 +59,23 @@ chrono = { version = "0.4.44", features = ["serde"] } # Symbol index (v0.6 — Code Intelligence) rusqlite = { version = "0.31", features = ["bundled"] } +# Tree-sitter AST parsing (optional — enable with --features tree-sitter) +tree-sitter = { version = "0.26", optional = true } +tree-sitter-rust = { version = "0.24", optional = true } +tree-sitter-go = { version = "0.25", optional = true } +tree-sitter-python = { version = "0.25", optional = true } +tree-sitter-typescript = { version = "0.23", optional = true } + +[features] +default = [] +tree-sitter = [ + "dep:tree-sitter", + "dep:tree-sitter-rust", + "dep:tree-sitter-go", + "dep:tree-sitter-python", + "dep:tree-sitter-typescript", +] + [dev-dependencies] assert_cmd = "2" predicates = "3" diff --git a/src/index/ast.rs b/src/index/ast.rs new file mode 100644 index 0000000..29cfdcf --- /dev/null +++ b/src/index/ast.rs @@ -0,0 +1,997 @@ +// ! Tree-sitter AST extraction. +// ! +// ! Provides proper AST-based symbol and edge extraction using tree-sitter. +// ! Only compiled when the `tree-sitter` feature is enabled. + +#![cfg(feature = "tree-sitter")] +#![allow(dead_code, unused)] + +use crate::index::extract::CallSite; +use crate::index::extract::ExtractedDef; +use crate::index::symbols::SymbolKind; + +/// A node in the code graph — a definition extracted from source. +#[derive(Debug, Clone)] +pub struct AstNode { + /// Symbol name. + pub name: String, + /// Kind of definition. + pub kind: SymbolKind, + /// Source file path. + pub file: String, + /// 1-based line number. + pub line: u32, + /// Full signature text. + pub signature: String, + /// Parent scope name (struct/class for methods). + pub parent: Option, +} + +impl From<&AstNode> for ExtractedDef { + fn from(n: &AstNode) -> Self { + Self { + name: n.name.clone(), + kind: n.kind.clone(), + file: n.file.clone(), + line: n.line, + signature: n.signature.clone(), + } + } +} + +impl From for ExtractedDef { + fn from(n: AstNode) -> Self { + ExtractedDef::from(&n) + } +} + +/// Edge type in the knowledge graph. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum EdgeKind { + Calls, + Imports, + Implements, + Inherits, + ChildOf, +} + +impl EdgeKind { + pub fn as_str(&self) -> &'static str { + match self { + Self::Calls => "CALLS", + Self::Imports => "IMPORTS", + Self::Implements => "IMPLEMENTS", + Self::Inherits => "INHERITS", + Self::ChildOf => "CHILD_OF", + } + } +} + +/// An edge in the code graph. +#[derive(Debug, Clone)] +pub struct AstEdge { + pub source: String, + pub kind: EdgeKind, + pub target: String, + pub file: String, + pub line: u32, +} + +impl From<&AstEdge> for CallSite { + fn from(e: &AstEdge) -> Self { + Self { + caller: e.source.clone(), + callee: e.target.clone(), + file: e.file.clone(), + line: e.line, + } + } +} + +impl From for CallSite { + fn from(e: AstEdge) -> Self { + CallSite::from(&e) + } +} + +// ─── Public API ───────────────────────────────────────────────────────── + +/// Get the tree-sitter Language for a file extension. +/// Returns `None` for unsupported languages. +pub fn get_language(ext: &str) -> Option { + match ext { + "rs" => Some(tree_sitter_rust::LANGUAGE.into()), + "go" => Some(tree_sitter_go::LANGUAGE.into()), + "py" | "pyi" => Some(tree_sitter_python::LANGUAGE.into()), + "ts" => Some(tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into()), + "tsx" => Some(tree_sitter_typescript::LANGUAGE_TSX.into()), + _ => None, + } +} + +/// Parse source code into a tree-sitter tree. +pub fn parse(source: &[u8], language: tree_sitter::Language) -> Option { + let mut parser = tree_sitter::Parser::new(); + parser.set_language(&language).ok()?; + parser.parse(source, None) +} + +/// Extract definitions and edges using tree-sitter. +/// +/// Returns `(nodes, edges)`. Empty vectors for unsupported languages +/// or parse failures. +pub fn extract(content: &str, language: &str, file_path: &str) -> (Vec, Vec) { + let lang = match get_language(language) { + Some(l) => l, + None => return (Vec::new(), Vec::new()), + }; + + let tree = match parse(content.as_bytes(), lang) { + Some(t) => t, + None => return (Vec::new(), Vec::new()), + }; + + match language { + "rs" => extract_rust(&tree.root_node(), content, file_path), + "go" => extract_go(&tree.root_node(), content, file_path), + "py" | "pyi" => extract_python(&tree.root_node(), content, file_path), + "ts" | "tsx" => extract_typescript(&tree.root_node(), content, file_path), + _ => (Vec::new(), Vec::new()), + } +} +// ─── Helpers ─────────────────────────────────────────────────────────── + +fn node_text(node: &tree_sitter::Node, source: &str) -> String { + node.utf8_text(source.as_bytes()).unwrap_or("").to_string() +} + +fn signature_for_node(node: &tree_sitter::Node, source: &str) -> String { + let lines: Vec<&str> = source.lines().collect(); + let start_row = node.start_position().row; + let end_row = node.end_position().row; + let start_col = node.start_position().column; + let end_col = node.end_position().column; + if start_row >= lines.len() { + return String::new(); + } + let first = &lines[start_row][start_col..]; + if start_row == end_row { + first[..std::cmp::min(end_col, first.len())].to_string() + } else if end_row < lines.len() { + let last = &lines[end_row][..std::cmp::min(end_col, lines[end_row].len())]; + format!("{first}\n...\n{last}") + } else { + first.to_string() + } +} + +/// Find the first child node matching one of the given kinds. +fn find_child_by_kind<'a>( + parent: &tree_sitter::Node<'a>, + kinds: &[&str], +) -> Option> { + let mut c = parent.walk(); + if c.goto_first_child() { + loop { + if kinds.contains(&c.node().kind()) { + return Some(c.node()); + } + if !c.goto_next_sibling() { + break; + } + } + } + None +} + +/// Get the "name" of a definition node — tries field name first, then kind-based fallback. +fn node_name(node: &tree_sitter::Node, source: &str) -> String { + if let Some(n) = node.child_by_field_name("name") { + return node_text(&n, source); + } + if let Some(n) = + find_child_by_kind(node, &["identifier", "type_identifier", "field_identifier"]) + { + return node_text(&n, source); + } + String::new() +} + +/// Get the type/trait target from a node (for impl, type_spec, etc.). +fn node_type(node: &tree_sitter::Node, source: &str) -> String { + if let Some(n) = node.child_by_field_name("type") { + return node_text(&n, source); + } + if let Some(n) = find_child_by_kind(node, &["type_identifier"]) { + return node_text(&n, source); + } + String::new() +} + +/// Get the trait/interface name from an impl node. +fn node_trait(node: &tree_sitter::Node, source: &str) -> String { + if let Some(n) = node.child_by_field_name("trait") { + return node_text(&n, source); + } + let mut c = node.walk(); + if c.goto_first_child() { + while c.goto_next_sibling() { + let n = c.node(); + if n.kind() == "type_identifier" { + return node_text(&n, source); + } + if n.kind() == "declaration_list" || n.kind() == "for" { + break; + } + } + } + String::new() +} + +// ─── Rust ────────────────────────────────────────────────────────────── + +fn extract_rust( + root: &tree_sitter::Node, + source: &str, + file_path: &str, +) -> (Vec, Vec) { + let mut nodes = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = root.walk(); + + for child in root.children(&mut cursor) { + match child.kind() { + "function_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name: name.clone(), + kind: SymbolKind::Function, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + extract_calls_from_node(&child, source, file_path, &name, &mut edges); + } + } + "struct_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Struct, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + } + } + "enum_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Enum, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + } + } + "trait_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Trait, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + } + } + "type_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::TypeAlias, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + } + } + "const_item" | "static_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Constant, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + } + } + "mod_item" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Module, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + } + } + "impl_item" => { + let trait_n = node_trait(&child, source); + let type_n = node_type(&child, source); + if !trait_n.is_empty() && !type_n.is_empty() { + edges.push(AstEdge { + source: type_n.clone(), + kind: EdgeKind::Implements, + target: trait_n, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + }); + } + // Extract methods inside impl block + let mut mc = child.walk(); + if mc.goto_first_child() { + loop { + if mc.node().kind() == "declaration_list" { + let mut dc = mc.node().walk(); + if dc.goto_first_child() { + loop { + let gc = dc.node(); + if gc.kind() == "function_item" { + let name = node_name(&gc, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Function, + file: file_path.to_string(), + line: (gc.start_position().row + 1) as u32, + signature: signature_for_node(&gc, source), + parent: Some(type_n.clone()), + }); + } + } + if !dc.goto_next_sibling() { + break; + } + } + } + } + if !mc.goto_next_sibling() { + break; + } + } + } + } + "use_declaration" => { + if let Some(arg) = child.child_by_field_name("argument") { + edges.push(AstEdge { + source: file_path.to_string(), + kind: EdgeKind::Imports, + target: node_text(&arg, source), + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + }); + } + } + _ => {} + } + } + (nodes, edges) +} + +/// Walk function body for call expressions using cursor-based DFS. +fn extract_calls_from_node( + node: &tree_sitter::Node, + source: &str, + file_path: &str, + caller: &str, + edges: &mut Vec, +) { + fn walk_calls( + node: &tree_sitter::Node, + source: &str, + file_path: &str, + caller: &str, + edges: &mut Vec, + ) { + if node.kind() == "call_expression" { + if let Some(fn_node) = node.child_by_field_name("function") { + let callee = node_text(&fn_node, source); + if !callee.is_empty() { + edges.push(AstEdge { + source: caller.to_string(), + kind: EdgeKind::Calls, + target: callee, + file: file_path.to_string(), + line: (node.start_position().row + 1) as u32, + }); + } + } + } + let mut cursor = node.walk(); + if cursor.goto_first_child() { + loop { + walk_calls(&cursor.node(), source, file_path, caller, edges); + if !cursor.goto_next_sibling() { + break; + } + } + } + } + walk_calls(node, source, file_path, caller, edges); +} + +// ─── Go ─────────────────────────────────────────────────────────────── + +fn extract_go( + root: &tree_sitter::Node, + source: &str, + file_path: &str, +) -> (Vec, Vec) { + let mut nodes = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = root.walk(); + + for child in root.children(&mut cursor) { + match child.kind() { + "function_declaration" | "method_declaration" => { + let name = node_name(&child, source); + if !name.is_empty() { + let kind = if child.kind() == "method_declaration" { + SymbolKind::Method + } else { + SymbolKind::Function + }; + nodes.push(AstNode { + name: name.clone(), + kind, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + extract_calls_from_node(&child, source, file_path, &name, &mut edges); + } + } + "type_declaration" => { + let mut tc = child.walk(); + if tc.goto_first_child() { + loop { + if tc.node().kind() == "type_spec" { + let name = node_name(&tc.node(), source); + let type_node = tc.node().child_by_field_name("type"); + let kind = type_node + .as_ref() + .map(|t| match t.kind() { + "struct_type" => SymbolKind::Struct, + "interface_type" => SymbolKind::Interface, + _ => SymbolKind::TypeAlias, + }) + .unwrap_or(SymbolKind::TypeAlias); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind, + file: file_path.to_string(), + line: (tc.node().start_position().row + 1) as u32, + signature: signature_for_node(&tc.node(), source), + parent: None, + }); + } + } + if !tc.goto_next_sibling() { + break; + } + } + } + } + "import_declaration" => { + if let Some(path_node) = child.child_by_field_name("path") { + edges.push(AstEdge { + source: file_path.to_string(), + kind: EdgeKind::Imports, + target: node_text(&path_node, source), + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + }); + } + } + "const_declaration" => { + let mut cc = child.walk(); + if cc.goto_first_child() { + loop { + if cc.node().kind() == "const_spec" { + let name = node_name(&cc.node(), source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Constant, + file: file_path.to_string(), + line: (cc.node().start_position().row + 1) as u32, + signature: signature_for_node(&cc.node(), source), + parent: None, + }); + } + } + if !cc.goto_next_sibling() { break; } + } + } + } + "var_declaration" => { + let mut vc = child.walk(); + if vc.goto_first_child() { + loop { + if vc.node().kind() == "var_spec" { + let name = node_name(&vc.node(), source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Constant, + file: file_path.to_string(), + line: (vc.node().start_position().row + 1) as u32, + signature: signature_for_node(&vc.node(), source), + parent: None, + }); + } + } + if !vc.goto_next_sibling() { break; } + } + } + } + _ => {} + } + } + (nodes, edges) +} + +// ─── Python ───────────────────────────────────────────────────────── + +fn extract_python( + root: &tree_sitter::Node, + source: &str, + file_path: &str, +) -> (Vec, Vec) { + let mut nodes = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = root.walk(); + + for child in root.children(&mut cursor) { + match child.kind() { + "function_definition" => { + let name = node_name(&child, source); + if !name.is_empty() { + nodes.push(AstNode { + name: name.clone(), + kind: SymbolKind::Function, + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + signature: signature_for_node(&child, source), + parent: None, + }); + extract_calls_from_node(&child, source, file_path, &name, &mut edges); + } + } + "class_definition" => { + let name = node_name(&child, source); + if !name.is_empty() { + let line = (child.start_position().row + 1) as u32; + // Check for parent classes (inheritance) + if let Some(bases) = child.child_by_field_name("superclasses") { + let mut bc = bases.walk(); + if bc.goto_first_child() { + loop { + let base_name = node_text(&bc.node(), source); + if !base_name.is_empty() { + edges.push(AstEdge { + source: name.clone(), + kind: EdgeKind::Inherits, + target: base_name, + file: file_path.to_string(), + line, + }); + } + if !bc.goto_next_sibling() { + break; + } + } + } + } + nodes.push(AstNode { + name: name.clone(), + kind: SymbolKind::Class, + file: file_path.to_string(), + line, + signature: signature_for_node(&child, source), + parent: None, + }); + // Extract methods inside class body + if let Some(body) = find_child_by_kind(&child, &["block"]) { + let mut mc = body.walk(); + if mc.goto_first_child() { + loop { + if mc.node().kind() == "function_definition" { + let mname = node_name(&mc.node(), source); + if !mname.is_empty() { + nodes.push(AstNode { + name: mname, + kind: SymbolKind::Method, + file: file_path.to_string(), + line: (mc.node().start_position().row + 1) as u32, + signature: signature_for_node(&mc.node(), source), + parent: Some(name.clone()), + }); + } + } + if !mc.goto_next_sibling() { + break; + } + } + } + } + } + } + "import_statement" | "import_from_statement" => { + edges.push(AstEdge { + source: file_path.to_string(), + kind: EdgeKind::Imports, + target: signature_for_node(&child, source), + file: file_path.to_string(), + line: (child.start_position().row + 1) as u32, + }); + } + _ => {} + } + } + (nodes, edges) +} + +// ─── TypeScript ───────────────────────────────────────────────────── + +// ─── TypeScript ───────────────────────────────────────────────────── + +fn extract_typescript( + root: &tree_sitter::Node, + source: &str, + file_path: &str, +) -> (Vec, Vec) { + let mut nodes = Vec::new(); + let mut edges = Vec::new(); + + fn process_ts_node( + node: &tree_sitter::Node, + source: &str, + file_path: &str, + nodes: &mut Vec, + edges: &mut Vec, + ) { + match node.kind() { + "function_declaration" => { + let name = node_name(node, source); + if !name.is_empty() { + nodes.push(AstNode { + name: name.clone(), + kind: SymbolKind::Function, + file: file_path.to_string(), + line: (node.start_position().row + 1) as u32, + signature: signature_for_node(node, source), + parent: None, + }); + extract_calls_from_node(node, source, file_path, &name, edges); + } + } + "class_declaration" => { + let name = node_name(node, source); + if !name.is_empty() { + let line = (node.start_position().row + 1) as u32; + nodes.push(AstNode { + name: name.clone(), + kind: SymbolKind::Class, + file: file_path.to_string(), + line, + signature: signature_for_node(node, source), + parent: None, + }); + if let Some(body) = find_child_by_kind(node, &["class_body"]) { + let mut mc = body.walk(); + if mc.goto_first_child() { + loop { + if mc.node().kind() == "method_definition" { + let mname = node_name(&mc.node(), source); + if !mname.is_empty() { + nodes.push(AstNode { + name: mname, + kind: SymbolKind::Method, + file: file_path.to_string(), + line: (mc.node().start_position().row + 1) as u32, + signature: signature_for_node(&mc.node(), source), + parent: Some(name.clone()), + }); + } + } + if !mc.goto_next_sibling() { + break; + } + } + } + } + } + } + "interface_declaration" => { + let name = node_name(node, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::Interface, + file: file_path.to_string(), + line: (node.start_position().row + 1) as u32, + signature: signature_for_node(node, source), + parent: None, + }); + } + } + "type_alias_declaration" => { + let name = node_name(node, source); + if !name.is_empty() { + nodes.push(AstNode { + name, + kind: SymbolKind::TypeAlias, + file: file_path.to_string(), + line: (node.start_position().row + 1) as u32, + signature: signature_for_node(node, source), + parent: None, + }); + } + } + "lexical_declaration" => { + // const/let declarations + let mut vc = node.walk(); + if vc.goto_first_child() { + loop { + if vc.node().kind() == "variable_declarator" { + let name = node_name(&vc.node(), source); + // Only extract as constant if ALL_CAPS + if !name.is_empty() && name == name.to_uppercase() && name.len() > 1 { + nodes.push(AstNode { + name, + kind: SymbolKind::Constant, + file: file_path.to_string(), + line: (vc.node().start_position().row + 1) as u32, + signature: signature_for_node(&vc.node(), source), + parent: None, + }); + } + } + if !vc.goto_next_sibling() { + break; + } + } + } + } + "variable_declaration" => { + let mut vc = node.walk(); + if vc.goto_first_child() { + loop { + if vc.node().kind() == "variable_declarator" { + let name = node_name(&vc.node(), source); + if !name.is_empty() && name == name.to_uppercase() && name.len() > 1 { + nodes.push(AstNode { + name, + kind: SymbolKind::Constant, + file: file_path.to_string(), + line: (vc.node().start_position().row + 1) as u32, + signature: signature_for_node(&vc.node(), source), + parent: None, + }); + } + } + if !vc.goto_next_sibling() { + break; + } + } + } + } + "export_statement" => { + // Unwrap: recurse into the exported declaration + let mut ec = node.walk(); + if ec.goto_first_child() { + loop { + let inner = ec.node(); + if inner.kind() != "export" && inner.kind() != ";" { + process_ts_node(&inner, source, file_path, nodes, edges); + } + if !ec.goto_next_sibling() { + break; + } + } + } + } + "import_statement" => { + edges.push(AstEdge { + source: file_path.to_string(), + kind: EdgeKind::Imports, + target: signature_for_node(node, source), + file: file_path.to_string(), + line: (node.start_position().row + 1) as u32, + }); + } + _ => {} + } + } + + let mut cursor = root.walk(); + for child in root.children(&mut cursor) { + process_ts_node(&child, source, file_path, &mut nodes, &mut edges); + } + + (nodes, edges) +} +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_get_language_supported() { + assert!(get_language("rs").is_some()); + assert!(get_language("go").is_some()); + assert!(get_language("py").is_some()); + assert!(get_language("ts").is_some()); + assert!(get_language("tsx").is_some()); + assert!(get_language("pyi").is_some()); + } + + #[test] + fn test_get_language_unsupported() { + assert!(get_language("rb").is_none()); + assert!(get_language("java").is_none()); + assert!(get_language("").is_none()); + } + + #[test] + fn test_extract_rust_functions() { + let code = r#"pub fn hello() -> String { + "hello".to_string() +} + +fn add(a: i32, b: i32) -> i32 { + a + b +}"#; + let (nodes, edges) = extract(code, "rs", "test.rs"); + let names: Vec<&str> = nodes.iter().map(|n| n.name.as_str()).collect(); + assert!(names.contains(&"hello")); + assert!(names.contains(&"add")); + assert!(edges.iter().any(|e| e.kind == EdgeKind::Calls)); + } + + #[test] + fn test_extract_rust_struct_and_impl() { + let code = r#"pub struct Cache { + data: HashMap, +} + +impl Cache { + pub fn new() -> Self { + Self { data: HashMap::new() } + } + + pub fn get(&self, key: &str) -> Option<&String> { + self.data.get(key) + } +} + +pub trait Store { + fn save(&self, data: &str); +} + +impl Store for Cache { + fn save(&self, data: &str) { + todo!() + } +}"#; + let (nodes, edges) = extract(code, "rs", "cache.rs"); + let names: Vec<&str> = nodes.iter().map(|n| n.name.as_str()).collect(); + assert!(names.contains(&"Cache")); + assert!(names.contains(&"Store")); + assert!(names.contains(&"new")); + assert!(names.contains(&"get")); + assert!(names.contains(&"save")); + assert!( + edges.iter().any(|e| e.kind == EdgeKind::Implements + && e.source == "Cache" + && e.target == "Store") + ); + } + + #[test] + fn test_extract_go_functions() { + let code = r#"package main + +func NewServer(port int) *Server { + return &Server{port: port} +} + +func (s *Server) Start() error { + return nil +}"#; + let (nodes, _edges) = extract(code, "go", "server.go"); + let names: Vec<&str> = nodes.iter().map(|n| n.name.as_str()).collect(); + assert!(names.contains(&"NewServer")); + assert!(names.contains(&"Start")); + } + + #[test] + fn test_extract_python_class() { + let code = r#"class AuthService: + def __init__(self): + self.secret = "" + + async def validate(self, token: str) -> bool: + return False + + def refresh(self): + pass"#; + let (nodes, _edges) = extract(code, "py", "auth.py"); + let classes: Vec<&str> = nodes + .iter() + .filter(|n| matches!(n.kind, SymbolKind::Class)) + .map(|n| n.name.as_str()) + .collect(); + let methods: Vec<&str> = nodes + .iter() + .filter(|n| matches!(n.kind, SymbolKind::Method)) + .map(|n| n.name.as_str()) + .collect(); + assert!(classes.contains(&"AuthService")); + assert!(methods.contains(&"validate")); + assert!(methods.contains(&"refresh")); + } + + #[test] + fn test_extract_typescript_interface() { + let code = r#"export interface User { + id: string; + name: string; +} + +export const DEFAULT_TIMEOUT = 5000; + +export function getUser(id: string): User { + return null as any; +}"#; + let (nodes, _edges) = extract(code, "ts", "user.ts"); + let names: Vec<&str> = nodes.iter().map(|n| n.name.as_str()).collect(); + assert!(names.contains(&"User")); + assert!(names.contains(&"DEFAULT_TIMEOUT")); + assert!(names.contains(&"getUser")); + } + + #[test] + fn test_extract_unsupported_fallback() { + let (nodes, edges) = extract("fn test() {}", "rb", "test.rb"); + assert!(nodes.is_empty()); + assert!(edges.is_empty()); + } + + #[test] + fn test_extract_empty() { + let (nodes, edges) = extract("", "rs", "empty.rs"); + assert!(nodes.is_empty()); + assert!(edges.is_empty()); + } +} diff --git a/src/index/extract.rs b/src/index/extract.rs index 0b013f2..30896e3 100644 --- a/src/index/extract.rs +++ b/src/index/extract.rs @@ -146,8 +146,20 @@ static RE_ZIG_CONST: LazyLock = /// Extract symbols from source code. /// -/// Returns a list of `IndexedSymbol` entries (without id, which is assigned by the database). +/// When the `tree-sitter` feature is enabled, uses AST-based extraction for +/// supported languages (Rust, Go, Python, TypeScript/TSX). Falls back to +/// regex for all other languages. pub fn extract_symbols(content: &str, language: &str, file_path: &str) -> Vec { + #[cfg(feature = "tree-sitter")] + { + use super::ast; + if ast::get_language(language).is_some() { + let (nodes, _edges) = ast::extract(content, language, file_path); + return nodes.into_iter().map(Into::into).collect(); + } + } + + // Regex fallback for unsupported languages or when feature is disabled let mut symbols = Vec::new(); for (line_num, line) in content.lines().enumerate() { @@ -350,11 +362,25 @@ fn extract_zig(line: &str, line_no: u32, file: &str, raw: &str, out: &mut Vec Vec { + #[cfg(feature = "tree-sitter")] + { + use super::ast; + if ast::get_language(language).is_some() { + let (_nodes, edges) = ast::extract(content, language, file_path); + return edges + .into_iter() + .filter(|e| e.kind == ast::EdgeKind::Calls) + .map(Into::into) + .collect(); + } + } + + // Regex fallback use crate::engine::context::extraction as ctx_extract; use crate::engine::context::types::SymbolKind as CtxSymbolKind; @@ -413,6 +439,24 @@ pub fn extract_calls(content: &str, language: &str, file_path: &str) -> Vec Vec { + use super::ast; + if ast::get_language(language).is_some() { + let (_nodes, edges) = ast::extract(content, language, file_path); + return edges; + } + Vec::new() +} + /// Detect function entry from a line (returns function name). fn detect_function_entry(line: &str, language: &str) -> Option { let trimmed = line.trim(); @@ -575,6 +619,7 @@ const MAX_SIZE: usize = 100; "#; let symbols = extract_symbols(code, "rs", "src/cache.rs"); let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect(); + eprintln!("GO symbols: {:?}", symbols.iter().map(|s| (&s.kind, &s.name)).collect::>()); assert!(names.contains(&"Cache")); assert!(names.contains(&"new")); diff --git a/src/index/graph.rs b/src/index/graph.rs index 45a4bbb..490488b 100644 --- a/src/index/graph.rs +++ b/src/index/graph.rs @@ -42,6 +42,56 @@ pub fn store_edges( Ok(count) } + + +/// A typed edge in the knowledge graph. +#[cfg(feature = "tree-sitter")] +#[derive(Debug, Clone)] +pub struct KgEdge { + /// Source symbol. + pub source: String, + /// Edge kind: CALLS, IMPORTS, IMPLEMENTS, INHERITS, CHILD_OF. + pub kind: String, + /// Target symbol. + pub target: String, + /// File where the relationship is defined. + pub file: String, + /// Line number. + pub line: u32, +} + +/// Store knowledge graph edges in the `edges` table. +#[cfg(feature = "tree-sitter")] +#[allow(dead_code)] +pub fn store_kg_edges( + conn: &Connection, + edges: &[KgEdge], + project_id: i64, +) -> anyhow::Result { + let tx = conn.unchecked_transaction()?; + let mut count = 0; + for edge in edges { + tx.execute( + "INSERT INTO edges (source, kind, target, file, line, project_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6)", + rusqlite::params![edge.source, edge.kind, edge.target, edge.file, edge.line as i64, project_id], + )?; + count += 1; + } + tx.commit()?; + Ok(count) +} + +/// Clear knowledge graph edges for a specific file. +#[cfg(feature = "tree-sitter")] +#[allow(dead_code)] +pub fn clear_kg_edges_for_file(conn: &Connection, file: &str, project_id: i64) -> anyhow::Result<()> { + conn.execute( + "DELETE FROM edges WHERE file = ?1 AND project_id = ?2", + rusqlite::params![file, project_id], + )?; + Ok(()) +} + /// Clear call graph edges for a specific file (before re-indexing), scoped to project. #[allow(dead_code)] pub fn clear_edges_for_file(conn: &Connection, file: &str, project_id: i64) -> anyhow::Result<()> { diff --git a/src/index/mod.rs b/src/index/mod.rs index 829cb05..8d8b1af 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -4,6 +4,8 @@ //! Uses regex-based extraction (same approach as `engine/context/extraction.rs`) //! stored in SQLite with FTS5 for fast full-text search. +#[cfg(feature = "tree-sitter")] +mod ast; mod extract; pub mod graph; pub mod schema; @@ -113,10 +115,22 @@ pub fn index_file( )?; } + #[cfg(feature = "tree-sitter")] + { + graph::clear_kg_edges_for_file(&tx, file_path, project_id)?; + let kg_edges = extract::extract_edges(content, language, file_path); + for e in &kg_edges { + tx.execute( + "INSERT INTO edges (source, kind, target, file, line, project_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6)", + rusqlite::params![e.source, e.kind.as_str(), e.target, e.file, e.line as i64, project_id], + )?; + } + } + tx.commit()?; debug!( - "Indexed {file_path}: {count} symbols, {} edges ({language})", + "Indexed {file_path}: {count} symbols, {} call edges ({language})", call_sites.len() ); Ok(count) diff --git a/src/index/schema.rs b/src/index/schema.rs index 3c7548a..a27b11e 100644 --- a/src/index/schema.rs +++ b/src/index/schema.rs @@ -4,7 +4,7 @@ use rusqlite::Connection; /// Current schema version. #[allow(dead_code)] -const SCHEMA_VERSION: i32 = 2; +const SCHEMA_VERSION: i32 = 3; /// Run database migrations (creates tables if not exist). pub fn run_migrations(conn: &Connection) -> anyhow::Result<()> { @@ -28,10 +28,14 @@ pub fn run_migrations(conn: &Connection) -> anyhow::Result<()> { if current < 2 { migrate_v2(conn)?; } + if current < 3 { + migrate_v3(conn)?; + } Ok(()) } + /// Migration v1: Initial schema — symbols, files, FTS5 index. fn migrate_v1(conn: &Connection) -> anyhow::Result<()> { conn.execute_batch( @@ -155,6 +159,42 @@ fn migrate_v2(conn: &Connection) -> anyhow::Result<()> { Ok(()) } +/// Migration v3: Knowledge graph edges. +/// +/// Adds `edges` table — a richer version of `call_graph` that supports +/// multiple edge types (CALLS, IMPORTS, IMPLEMENTS, INHERITS, CHILD_OF). +/// Existing `call_graph` rows are migrated into `edges` with kind='CALLS'. +fn migrate_v3(conn: &Connection) -> anyhow::Result<()> { + conn.execute_batch( + " + -- Knowledge graph edges: source → target with typed relationship + CREATE TABLE IF NOT EXISTS edges ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + source TEXT NOT NULL, + kind TEXT NOT NULL, + target TEXT NOT NULL, + file TEXT NOT NULL, + line INTEGER NOT NULL, + project_id INTEGER REFERENCES projects(id) ON DELETE CASCADE + ); + + CREATE INDEX IF NOT EXISTS idx_edges_source ON edges(source); + CREATE INDEX IF NOT EXISTS idx_edges_target ON edges(target); + CREATE INDEX IF NOT EXISTS idx_edges_kind ON edges(kind); + CREATE INDEX IF NOT EXISTS idx_edges_project ON edges(project_id); + CREATE INDEX IF NOT EXISTS idx_edges_file ON edges(file); + + -- Migrate existing call_graph data into edges + INSERT OR IGNORE INTO edges (source, kind, target, file, line, project_id) + SELECT caller, 'CALLS', callee, file, line, project_id FROM call_graph; + ", + )?; + + conn.execute("INSERT INTO schema_version (version) VALUES (3)", [])?; + + Ok(()) +} + /// Get or create a project entry by root path. /// /// Returns the project ID. @@ -340,4 +380,66 @@ mod tests { .unwrap(); assert_eq!(proj_count, 0); } -} + + #[test] + fn test_v3_edges_table() { + let conn = mem_conn(); + let pid = get_or_create_project(&conn, "/test/proj").unwrap(); + + // Check edges table exists + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM edges", [], |row| row.get(0)) + .unwrap(); + assert_eq!(count, 0); + + // Insert an edge + conn.execute( + "INSERT INTO edges (source, kind, target, file, line, project_id) VALUES (?1, ?2, ?3, ?4, ?5, ?6)", + rusqlite::params!["Cache", "IMPLEMENTS", "Store", "cache.rs", 10, pid], + ) + .unwrap(); + + // Query it back + let (source, kind, target): (String, String, String) = conn + .query_row( + "SELECT source, kind, target FROM edges WHERE kind = 'IMPLEMENTS'", + [], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + ) + .unwrap(); + assert_eq!(source, "Cache"); + assert_eq!(kind, "IMPLEMENTS"); + assert_eq!(target, "Store"); + } + + #[test] + fn test_v3_migrate_call_graph() { + let conn = mem_conn(); + let pid = get_or_create_project(&conn, "/test/proj").unwrap(); + + // Insert into old call_graph + conn.execute( + "INSERT INTO call_graph (caller, callee, file, line, project_id) VALUES (?1, ?2, ?3, ?4, ?5)", + rusqlite::params!["main", "helper", "main.rs", 5, pid], + ) + .unwrap(); + + // Migration already ran, so edges should be empty + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM edges WHERE kind = 'CALLS'", [], |row| row.get(0)) + .unwrap(); + assert_eq!(count, 0); + + // Manually verify the migration SQL works + conn.execute( + "INSERT OR IGNORE INTO edges (source, kind, target, file, line, project_id) SELECT caller, 'CALLS', callee, file, line, project_id FROM call_graph", + [], + ) + .unwrap(); + + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM edges WHERE kind = 'CALLS'", [], |row| row.get(0)) + .unwrap(); + assert_eq!(count, 1); + } +} \ No newline at end of file From 45d4535b41deab13cc307ffcd9fdeeb39fcc662f Mon Sep 17 00:00:00 2001 From: ajianaz Date: Wed, 22 Jul 2026 21:21:59 +0700 Subject: [PATCH 2/2] style: cargo fmt --all --- src/index/ast.rs | 8 ++++++-- src/index/extract.rs | 8 +++++++- src/index/graph.rs | 8 +++++--- src/index/schema.rs | 15 +++++++++++---- 4 files changed, 29 insertions(+), 10 deletions(-) diff --git a/src/index/ast.rs b/src/index/ast.rs index 29cfdcf..e97c04a 100644 --- a/src/index/ast.rs +++ b/src/index/ast.rs @@ -530,7 +530,9 @@ fn extract_go( }); } } - if !cc.goto_next_sibling() { break; } + if !cc.goto_next_sibling() { + break; + } } } } @@ -551,7 +553,9 @@ fn extract_go( }); } } - if !vc.goto_next_sibling() { break; } + if !vc.goto_next_sibling() { + break; + } } } } diff --git a/src/index/extract.rs b/src/index/extract.rs index 30896e3..ff78506 100644 --- a/src/index/extract.rs +++ b/src/index/extract.rs @@ -619,7 +619,13 @@ const MAX_SIZE: usize = 100; "#; let symbols = extract_symbols(code, "rs", "src/cache.rs"); let names: Vec<&str> = symbols.iter().map(|s| s.name.as_str()).collect(); - eprintln!("GO symbols: {:?}", symbols.iter().map(|s| (&s.kind, &s.name)).collect::>()); + eprintln!( + "GO symbols: {:?}", + symbols + .iter() + .map(|s| (&s.kind, &s.name)) + .collect::>() + ); assert!(names.contains(&"Cache")); assert!(names.contains(&"new")); diff --git a/src/index/graph.rs b/src/index/graph.rs index 490488b..35f7e27 100644 --- a/src/index/graph.rs +++ b/src/index/graph.rs @@ -42,8 +42,6 @@ pub fn store_edges( Ok(count) } - - /// A typed edge in the knowledge graph. #[cfg(feature = "tree-sitter")] #[derive(Debug, Clone)] @@ -84,7 +82,11 @@ pub fn store_kg_edges( /// Clear knowledge graph edges for a specific file. #[cfg(feature = "tree-sitter")] #[allow(dead_code)] -pub fn clear_kg_edges_for_file(conn: &Connection, file: &str, project_id: i64) -> anyhow::Result<()> { +pub fn clear_kg_edges_for_file( + conn: &Connection, + file: &str, + project_id: i64, +) -> anyhow::Result<()> { conn.execute( "DELETE FROM edges WHERE file = ?1 AND project_id = ?2", rusqlite::params![file, project_id], diff --git a/src/index/schema.rs b/src/index/schema.rs index a27b11e..5a69ae9 100644 --- a/src/index/schema.rs +++ b/src/index/schema.rs @@ -35,7 +35,6 @@ pub fn run_migrations(conn: &Connection) -> anyhow::Result<()> { Ok(()) } - /// Migration v1: Initial schema — symbols, files, FTS5 index. fn migrate_v1(conn: &Connection) -> anyhow::Result<()> { conn.execute_batch( @@ -426,7 +425,11 @@ mod tests { // Migration already ran, so edges should be empty let count: i64 = conn - .query_row("SELECT COUNT(*) FROM edges WHERE kind = 'CALLS'", [], |row| row.get(0)) + .query_row( + "SELECT COUNT(*) FROM edges WHERE kind = 'CALLS'", + [], + |row| row.get(0), + ) .unwrap(); assert_eq!(count, 0); @@ -438,8 +441,12 @@ mod tests { .unwrap(); let count: i64 = conn - .query_row("SELECT COUNT(*) FROM edges WHERE kind = 'CALLS'", [], |row| row.get(0)) + .query_row( + "SELECT COUNT(*) FROM edges WHERE kind = 'CALLS'", + [], + |row| row.get(0), + ) .unwrap(); assert_eq!(count, 1); } -} \ No newline at end of file +}