use crate::types::*;
use super::{node_text, field_text, extract_doc_comment, extract_signature};
pub fn extract(tree: &tree_sitter::Tree, source: &[u8]) -> (Vec<Symbol>, Vec<Import>) {
let root = tree.root_node();
let mut symbols = Vec::new();
let mut imports = Vec::new();
let mut cursor = root.walk();
for child in root.children(&mut cursor) {
extract_top_level(&child, source, &mut symbols, &mut imports);
}
(symbols, imports)
}
fn extract_top_level(
node: &tree_sitter::Node,
source: &[u8],
symbols: &mut Vec<Symbol>,
imports: &mut Vec<Import>,
) {
match node.kind() {
"function_declaration" => {
if let Some(sym) = extract_function(node, source) {
symbols.push(sym);
}
}
"local_function" | "local_function_definition_statement" => {
if let Some(sym) = extract_local_function(node, source) {
symbols.push(sym);
}
}
"local_variable_declaration" | "variable_declaration" => {
extract_variable(node, source, symbols, imports);
}
"assignment_statement" | "assignment" => {
extract_assignment(node, source, symbols, imports);
}
"function_call" | "function_call_statement" => {
if is_require_call(node, source) {
let text = node_text(node, source);
if let Some(path) = extract_require_path(text) {
imports.push(Import {
path,
alias: None,
span: Span::from_node(node),
});
}
}
}
"expression_statement" => {
let mut inner = node.walk();
for child in node.children(&mut inner) {
if child.kind() == "function_call" || child.kind() == "function_call_statement" {
if is_require_call(&child, source) {
let text = node_text(&child, source);
if let Some(path) = extract_require_path(text) {
imports.push(Import {
path,
alias: None,
span: Span::from_node(&child),
});
}
}
}
}
}
"return_statement" => {}
_ => {}
}
}
fn extract_function(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
let name = extract_function_name(node, source)?;
let kind = if name.contains(':') {
SymbolKind::Method
} else if name.contains('.') {
SymbolKind::Function
} else {
SymbolKind::Function
};
let parent = if name.contains(':') || name.contains('.') {
Some(name.split(&[':', '.'][..]).next().unwrap_or("").to_string())
} else {
None
};
let short_name = name
.rsplit(&[':', '.'][..])
.next()
.unwrap_or(&name)
.to_string();
Some(Symbol {
name: short_name,
kind,
span: Span::from_node(node),
signature: extract_signature(node, "body", source),
doc_comment: extract_doc_comment(node, source),
parent,
children: Vec::new(),
})
}
fn extract_local_function(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
let name = extract_function_name(node, source)?;
Some(Symbol {
name,
kind: SymbolKind::Function,
span: Span::from_node(node),
signature: extract_signature(node, "body", source),
doc_comment: extract_doc_comment(node, source),
parent: None,
children: Vec::new(),
})
}
fn extract_function_name(node: &tree_sitter::Node, source: &[u8]) -> Option<String> {
if let Some(name) = field_text(node, "name", source) {
if !name.is_empty() {
return Some(name.to_string());
}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
match child.kind() {
"identifier" | "IDENTIFIER" => {
let t = node_text(&child, source);
if !t.is_empty() && t != "function" && t != "local" {
return Some(t.to_string());
}
}
"dot_index_expression" | "method_index_expression" => {
return Some(node_text(&child, source).to_string());
}
"function_name" | "function_name_field" => {
return Some(node_text(&child, source).to_string());
}
_ => {}
}
}
let text = node_text(node, source);
extract_name_from_function_text(text)
}
fn extract_name_from_function_text(text: &str) -> Option<String> {
let text = text.trim();
let text = text.strip_prefix("local ").unwrap_or(text);
let text = text.strip_prefix("function ").unwrap_or(text);
let paren = text.find('(')?;
let name = text[..paren].trim();
if name.is_empty() {
None
} else {
Some(name.to_string())
}
}
fn extract_variable(
node: &tree_sitter::Node,
source: &[u8],
symbols: &mut Vec<Symbol>,
imports: &mut Vec<Import>,
) {
let text = node_text(node, source);
if text.contains("require") {
if let Some(path) = extract_require_path(text) {
let alias = extract_local_var_name(node, source)
.or_else(|| extract_var_name_from_text(text));
imports.push(Import {
path,
alias,
span: Span::from_node(node),
});
return;
}
}
if text.contains("function") && text.contains("end") {
if let Some(name) = extract_local_var_name(node, source)
.or_else(|| extract_var_name_from_text(text))
{
symbols.push(Symbol {
name,
kind: SymbolKind::Function,
span: Span::from_node(node),
signature: first_line(text).to_string(),
doc_comment: extract_doc_comment(node, source),
parent: None,
children: Vec::new(),
});
return;
}
}
if let Some(name) = extract_local_var_name(node, source)
.or_else(|| extract_var_name_from_text(text))
{
symbols.push(Symbol {
name,
kind: SymbolKind::Const,
span: Span::from_node(node),
signature: first_line(text).to_string(),
doc_comment: extract_doc_comment(node, source),
parent: None,
children: Vec::new(),
});
}
}
fn extract_assignment(
node: &tree_sitter::Node,
source: &[u8],
symbols: &mut Vec<Symbol>,
imports: &mut Vec<Import>,
) {
let text = node_text(node, source);
if text.contains("require") {
if let Some(path) = extract_require_path(text) {
let alias = extract_var_name_from_text(text);
imports.push(Import {
path,
alias,
span: Span::from_node(node),
});
return;
}
}
if text.contains("function") {
if let Some(name) = extract_var_name_from_text(text) {
symbols.push(Symbol {
name,
kind: SymbolKind::Function,
span: Span::from_node(node),
signature: first_line(text).to_string(),
doc_comment: extract_doc_comment(node, source),
parent: None,
children: Vec::new(),
});
return;
}
}
if let Some(name) = extract_var_name_from_text(text) {
symbols.push(Symbol {
name,
kind: SymbolKind::Const,
span: Span::from_node(node),
signature: first_line(text).to_string(),
doc_comment: extract_doc_comment(node, source),
parent: None,
children: Vec::new(),
});
}
}
fn extract_local_var_name(node: &tree_sitter::Node, source: &[u8]) -> Option<String> {
if let Some(name) = field_text(node, "name", source) {
if !name.is_empty() {
return Some(name.to_string());
}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "identifier" || child.kind() == "variable_list" {
let t = node_text(&child, source).trim().to_string();
if !t.is_empty() && t != "local" {
return Some(t);
}
}
}
None
}
fn extract_var_name_from_text(text: &str) -> Option<String> {
let text = text.trim();
let text = text.strip_prefix("local ").unwrap_or(text);
let eq = text.find('=')?;
let name = text[..eq].trim();
if name.is_empty() {
None
} else {
Some(name.to_string())
}
}
fn extract_require_path(text: &str) -> Option<String> {
let start_dq = text.find("require");
if let Some(req_pos) = start_dq {
let after = &text[req_pos + 7..];
if let Some(q_start) = after.find(|c: char| c == '"' || c == '\'') {
let quote_char = after.as_bytes()[q_start] as char;
let inner = &after[q_start + 1..];
if let Some(q_end) = inner.find(quote_char) {
return Some(inner[..q_end].to_string());
}
}
}
None
}
fn is_require_call(node: &tree_sitter::Node, source: &[u8]) -> bool {
let text = node_text(node, source);
text.trim().starts_with("require")
}
fn first_line(text: &str) -> &str {
text.lines().next().unwrap_or(text)
}
#[cfg(test)]
mod tests {
use crate::{parse, Language, SymbolKind};
#[test]
fn test_lua_extract() {
let source = r#"local json = require("cjson")
local utils = require('utils')
local MAX_SIZE = 1024
local function helper(x)
return x * 2
end
function greet(name)
print("Hello, " .. name)
end
function M:method(arg)
self.value = arg
end
local tbl = {}
"#;
let parsed = parse(source, Language::Lua);
if let Some(parsed) = parsed {
let names: Vec<&str> = parsed.symbols.iter().map(|s| s.name.as_str()).collect();
let kinds: Vec<(&str, SymbolKind)> = parsed
.symbols
.iter()
.map(|s| (s.name.as_str(), s.kind))
.collect();
let import_paths: Vec<&str> = parsed.imports.iter().map(|i| i.path.as_str()).collect();
assert!(
import_paths.iter().any(|p| *p == "cjson"),
"Should find cjson import, got: {:?}",
import_paths
);
assert!(
import_paths.iter().any(|p| *p == "utils"),
"Should find utils import, got: {:?}",
import_paths
);
assert!(
kinds.iter().any(|(n, k)| *n == "helper" && *k == SymbolKind::Function),
"Should find 'helper' function, got: {:?}",
kinds
);
assert!(
kinds.iter().any(|(n, k)| *n == "greet" && *k == SymbolKind::Function),
"Should find 'greet' function, got: {:?}",
kinds
);
assert!(
kinds.iter().any(|(n, k)| *n == "method" && *k == SymbolKind::Method),
"Should find 'method' method, got: {:?}",
kinds
);
assert!(
kinds.iter().any(|(n, k)| *n == "MAX_SIZE" && *k == SymbolKind::Const),
"Should find 'MAX_SIZE' const, got: {:?}",
kinds
);
eprintln!("Lua symbols: {:?}", names);
eprintln!("Lua imports: {:?}", import_paths);
} else {
eprintln!("Lua parser not available or parse failed — skipping assertions");
}
}
}