fn chunk_go_file(source: &str) -> Result<Vec<CodeChunk>, String> {
let tree = parse_go(source)?;
let root = tree.root_node();
let mut chunks = Vec::new();
extract_go_items(root, source, &mut chunks);
Ok(chunks)
}
#[cfg(feature = "go-ast")]
fn parse_go(source: &str) -> Result<Tree, String> {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_go::LANGUAGE.into())
.map_err(|e| format!("Failed to set Go language: {e}"))?;
parser
.parse(source, None)
.ok_or_else(|| "Failed to parse Go source".to_string())
}
#[cfg(not(feature = "go-ast"))]
fn parse_go(_source: &str) -> Result<Tree, String> {
Err("go-ast feature is disabled".to_string())
}
fn extract_go_type_name(node: Node, source: &str) -> Option<String> {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "type_spec" {
if let Some(name_node) = child.child_by_field_name("name") {
return Some(source[name_node.byte_range()].to_string());
}
}
}
None
}
fn extract_go_items(node: Node, source: &str, chunks: &mut Vec<CodeChunk>) {
match node.kind() {
"function_declaration" => {
if let Some(name_node) = node.child_by_field_name("name") {
let name = source[name_node.byte_range()].to_string();
let start_byte = find_doc_comment_start(node, source);
let content = source
.get(start_byte..node.end_byte())
.unwrap_or_default()
.to_string();
push_chunk(chunks, ChunkType::Function, name, "go", node, content);
}
return; }
"method_declaration" => {
if let Some(name_node) = node.child_by_field_name("name") {
let name = source[name_node.byte_range()].to_string();
let content = source[node.byte_range()].to_string();
push_chunk(chunks, ChunkType::Function, name, "go", node, content);
}
return; }
"type_declaration" => {
if let Some(name) = extract_go_type_name(node, source) {
let start_byte = find_doc_comment_start(node, source);
let content = source
.get(start_byte..node.end_byte())
.unwrap_or_default()
.to_string();
push_chunk(chunks, ChunkType::Class, name, "go", node, content);
}
return; }
_ => {}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
extract_go_items(child, source, chunks);
}
}
fn chunk_lua_file(source: &str) -> Result<Vec<CodeChunk>, String> {
let tree = parse_lua(source)?;
let root = tree.root_node();
let mut chunks = Vec::new();
extract_lua_items(root, source, &mut chunks);
Ok(chunks)
}
#[cfg(feature = "lua-ast")]
fn parse_lua(source: &str) -> Result<Tree, String> {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_lua::LANGUAGE.into())
.map_err(|e| format!("Failed to set Lua language: {e}"))?;
parser
.parse(source, None)
.ok_or_else(|| "Failed to parse Lua source".to_string())
}
#[cfg(not(feature = "lua-ast"))]
fn parse_lua(_source: &str) -> Result<Tree, String> {
Err("lua-ast feature is disabled".to_string())
}
fn extract_lua_items(node: Node, source: &str, chunks: &mut Vec<CodeChunk>) {
if node.kind() == "function_declaration" {
if let Some(name) = extract_lua_function_name(&node, source) {
let start_byte = find_doc_comment_start(node, source);
let content = source
.get(start_byte..node.end_byte())
.unwrap_or_default()
.to_string();
chunks.push(CodeChunk {
file_path: String::new(),
chunk_type: ChunkType::Function,
chunk_name: name,
language: "lua".to_string(),
start_line: node.start_position().row + 1,
end_line: node.end_position().row + 1,
content: content.clone(),
content_checksum: compute_checksum(&content),
});
}
return;
}
else if node.kind() == "local_function_declaration" {
if let Some(name_node) = node.child_by_field_name("name") {
let name = source[name_node.byte_range()].to_string();
let start_byte = find_doc_comment_start(node, source);
let content = source
.get(start_byte..node.end_byte())
.unwrap_or_default()
.to_string();
chunks.push(CodeChunk {
file_path: String::new(),
chunk_type: ChunkType::Function,
chunk_name: name,
language: "lua".to_string(),
start_line: node.start_position().row + 1,
end_line: node.end_position().row + 1,
content: content.clone(),
content_checksum: compute_checksum(&content),
});
}
return;
}
else if node.kind() == "variable_declaration" {
extract_lua_variable_function(node, source, chunks);
return;
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
extract_lua_items(child, source, chunks);
}
}
fn extract_lua_function_name(node: &Node, source: &str) -> Option<String> {
if let Some(name_node) = node.child_by_field_name("name") {
return Some(source[name_node.byte_range()].to_string());
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
match child.kind() {
"identifier" | "dot_index_expression" | "method_index_expression" => {
return Some(source[child.byte_range()].to_string());
}
_ => {}
}
}
None
}
fn find_lua_var_name(node: Node, source: &str) -> Option<String> {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "identifier" {
return Some(source[child.byte_range()].to_string());
}
}
None
}
fn has_function_definition(node: Node) -> bool {
if node.kind() == "function_definition" {
return true;
}
let mut cursor = node.walk();
let found = node
.children(&mut cursor)
.any(|c| c.kind() == "function_definition");
found
}
fn extract_lua_variable_function(node: Node, source: &str, chunks: &mut Vec<CodeChunk>) {
let mut cursor = node.walk();
let mut var_name = None;
for child in node.children(&mut cursor) {
if child.kind() == "assignment_statement" || child.kind() == "variable_list" {
var_name = var_name.or_else(|| find_lua_var_name(child, source));
}
if !has_function_definition(child) {
continue;
}
if let Some(name) = &var_name {
let content = source[node.byte_range()].to_string();
push_chunk(
chunks,
ChunkType::Function,
name.clone(),
"lua",
node,
content,
);
return; }
}
}