use super::languages::SupportedLanguage;
use crate::db::CodeChunk;
use anyhow::Result;
use chrono::Utc;
use tree_sitter::{Parser, Tree};
use uuid::Uuid;
#[derive(Clone)]
pub struct ChunkConfig {
pub min_chunk_lines: usize,
pub max_chunk_lines: usize,
pub overlap_lines: usize,
pub include_file_chunks: bool,
}
impl Default for ChunkConfig {
fn default() -> Self {
Self {
min_chunk_lines: 5,
max_chunk_lines: 100,
overlap_lines: 2,
include_file_chunks: true,
}
}
}
fn derive_module_path(file_path: &str) -> String {
let normalised = file_path.replace('\\', "/");
let stripped = normalised
.strip_prefix("src/")
.or_else(|| normalised.strip_prefix("lib/"))
.unwrap_or(&normalised);
let no_ext = if let Some(pos) = stripped.rfind('.') {
let after_last_slash = stripped.rfind('/').map(|i| i + 1).unwrap_or(0);
if pos > after_last_slash {
&stripped[..pos]
} else {
stripped
}
} else {
stripped
};
let parts: Vec<&str> = no_ext.split('/').collect();
let trim_last = parts
.last()
.map(|last| *last == "mod" || *last == "__init__" || *last == "index")
.unwrap_or(false);
let final_parts = if trim_last {
&parts[..parts.len().saturating_sub(1)]
} else {
parts.as_slice()
};
let result = final_parts.join(".");
if result.is_empty() {
no_ext.replace('/', ".")
} else {
result
}
}
pub fn chunk_file(
content: &str,
file_path: &str,
language: SupportedLanguage,
config: &ChunkConfig,
) -> Result<(Vec<CodeChunk>, Tree)> {
let mut chunks = Vec::new();
let lines: Vec<&str> = content.lines().collect();
let timestamp = Utc::now().to_rfc3339();
let module_path = Some(derive_module_path(file_path));
let mut parser = Parser::new();
parser
.set_language(&language.tree_sitter_language())
.map_err(|e| anyhow::anyhow!("Failed to set language: {}", e))?;
let tree = parser
.parse(content, None)
.ok_or_else(|| anyhow::anyhow!("Failed to parse file"))?;
extract_semantic_chunks(
&tree,
content,
file_path,
language,
config.min_chunk_lines,
&mut chunks,
×tamp,
module_path.as_deref(),
);
if chunks.is_empty() || (config.include_file_chunks && lines.len() <= config.max_chunk_lines) {
let file_chunk = CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: content.to_string(),
start_line: 1,
end_line: lines.len() as u32,
chunk_type: "file".to_string(),
language: language.name().to_string(),
symbol_name: None,
content_hash: compute_content_hash(content),
indexed_at: timestamp.clone(),
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: module_path.clone(),
};
chunks.push(file_chunk);
}
if chunks.is_empty() && lines.len() > config.max_chunk_lines {
let mut start = 0;
while start < lines.len() {
let end = (start + config.max_chunk_lines).min(lines.len());
let chunk_content = lines[start..end].join("\n");
let chunk = CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: chunk_content.clone(),
start_line: (start + 1) as u32,
end_line: end as u32,
chunk_type: "block".to_string(),
language: language.name().to_string(),
symbol_name: None,
content_hash: compute_content_hash(&chunk_content),
indexed_at: timestamp.clone(),
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: module_path.clone(),
};
chunks.push(chunk);
if end >= lines.len() {
break;
}
start = end - config.overlap_lines;
}
}
Ok((chunks, tree))
}
pub fn chunk_plain_text(
content: &str,
file_path: &str,
language_label: &str,
config: &ChunkConfig,
) -> Result<Vec<CodeChunk>> {
let lines: Vec<&str> = content.lines().collect();
let timestamp = Utc::now().to_rfc3339();
let module_path = Some(derive_module_path(file_path));
if lines.is_empty() {
return Ok(Vec::new());
}
let chunks = if language_label == "markdown" {
chunk_markdown(
&lines,
file_path,
language_label,
config,
×tamp,
module_path.as_deref(),
)
} else {
chunk_lines_overlap(
&lines,
file_path,
language_label,
config,
×tamp,
module_path.as_deref(),
)
};
Ok(chunks)
}
fn chunk_markdown(
lines: &[&str],
file_path: &str,
language_label: &str,
config: &ChunkConfig,
timestamp: &str,
module_path: Option<&str>,
) -> Vec<CodeChunk> {
let mut chunks = Vec::new();
let mut section_start = 0;
let mut section_heading: Option<String> = None;
let mut found_heading = false;
for (i, line) in lines.iter().enumerate() {
let is_heading = line.starts_with('#');
let is_last = i == lines.len() - 1;
if is_heading || is_last {
let section_end = if is_last && !is_heading { i + 1 } else { i };
if section_end > section_start {
let section_content = lines[section_start..section_end].join("\n");
if !section_content.trim().is_empty() {
if section_end - section_start > config.max_chunk_lines {
let sub_lines = &lines[section_start..section_end];
let sub_chunks = chunk_lines_overlap_with_offset(
sub_lines,
file_path,
language_label,
config,
timestamp,
module_path,
section_start,
section_heading.as_deref(),
);
chunks.extend(sub_chunks);
} else {
chunks.push(CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: section_content.clone(),
start_line: (section_start + 1) as u32,
end_line: section_end as u32,
chunk_type: "section".to_string(),
language: language_label.to_string(),
symbol_name: section_heading.clone(),
content_hash: compute_content_hash(§ion_content),
indexed_at: timestamp.to_string(),
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: module_path.map(|s| s.to_string()),
});
}
}
}
if is_heading {
found_heading = true;
section_start = i;
section_heading = Some(line.trim_start_matches('#').trim().to_string());
if is_last {
let section_content = lines[i..=i].join("\n");
chunks.push(CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: section_content.clone(),
start_line: (i + 1) as u32,
end_line: (i + 1) as u32,
chunk_type: "section".to_string(),
language: language_label.to_string(),
symbol_name: section_heading.clone(),
content_hash: compute_content_hash(§ion_content),
indexed_at: timestamp.to_string(),
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: module_path.map(|s| s.to_string()),
});
}
}
}
}
if !found_heading {
return chunk_lines_overlap(
lines,
file_path,
language_label,
config,
timestamp,
module_path,
);
}
chunks
}
fn chunk_lines_overlap(
lines: &[&str],
file_path: &str,
language_label: &str,
config: &ChunkConfig,
timestamp: &str,
module_path: Option<&str>,
) -> Vec<CodeChunk> {
chunk_lines_overlap_with_offset(
lines,
file_path,
language_label,
config,
timestamp,
module_path,
0,
None,
)
}
fn chunk_lines_overlap_with_offset(
lines: &[&str],
file_path: &str,
language_label: &str,
config: &ChunkConfig,
timestamp: &str,
module_path: Option<&str>,
line_offset: usize,
symbol_name: Option<&str>,
) -> Vec<CodeChunk> {
let mut chunks = Vec::new();
if lines.len() <= config.max_chunk_lines {
let content = lines.join("\n");
if !content.trim().is_empty() {
chunks.push(CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: content.clone(),
start_line: (line_offset + 1) as u32,
end_line: (line_offset + lines.len()) as u32,
chunk_type: "block".to_string(),
language: language_label.to_string(),
symbol_name: symbol_name.map(|s| s.to_string()),
content_hash: compute_content_hash(&content),
indexed_at: timestamp.to_string(),
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: module_path.map(|s| s.to_string()),
});
}
return chunks;
}
let mut start = 0;
while start < lines.len() {
let end = (start + config.max_chunk_lines).min(lines.len());
let chunk_content = lines[start..end].join("\n");
chunks.push(CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: chunk_content.clone(),
start_line: (line_offset + start + 1) as u32,
end_line: (line_offset + end) as u32,
chunk_type: "block".to_string(),
language: language_label.to_string(),
symbol_name: symbol_name.map(|s| s.to_string()),
content_hash: compute_content_hash(&chunk_content),
indexed_at: timestamp.to_string(),
parent_symbol: None,
signature: None,
doc_comment: None,
module_path: module_path.map(|s| s.to_string()),
});
if end >= lines.len() {
break;
}
start = end - config.overlap_lines;
}
chunks
}
fn extract_semantic_chunks(
tree: &Tree,
content: &str,
file_path: &str,
language: SupportedLanguage,
min_chunk_lines: usize,
chunks: &mut Vec<CodeChunk>,
timestamp: &str,
module_path: Option<&str>,
) {
let root = tree.root_node();
let mut cursor = root.walk();
let lines: Vec<&str> = content.lines().collect();
let function_types = language.function_node_types();
let class_types = language.class_node_types();
fn visit_node(
cursor: &mut tree_sitter::TreeCursor,
content: &str,
lines: &[&str],
file_path: &str,
language: SupportedLanguage,
min_chunk_lines: usize,
function_types: &[&str],
class_types: &[&str],
chunks: &mut Vec<CodeChunk>,
timestamp: &str,
module_path: Option<&str>,
) {
let node = cursor.node();
let node_type = node.kind();
let chunk_type = if function_types.contains(&node_type) {
Some("function")
} else if class_types.contains(&node_type) {
Some("class")
} else {
None
};
if let Some(chunk_type) = chunk_type {
let start_line = node.start_position().row;
let end_line = node.end_position().row;
let line_count = end_line.saturating_sub(start_line) + 1;
if line_count >= min_chunk_lines {
let chunk_content = lines[start_line..=end_line.min(lines.len() - 1)].join("\n");
let symbol_name = node
.child_by_field_name(language.name_field())
.map(|n| n.utf8_text(content.as_bytes()).unwrap_or("").to_string())
.filter(|s| !s.is_empty());
let parent_symbol = extract_parent_symbol(&node, content, class_types, language);
let signature = if chunk_type == "function" {
extract_signature(&node, content)
} else {
None
};
let doc_comment = extract_doc_comment(&node, content, language);
let chunk = CodeChunk {
id: Uuid::new_v4().to_string(),
file_path: file_path.to_string(),
content: chunk_content.clone(),
start_line: (start_line + 1) as u32,
end_line: (end_line + 1) as u32,
chunk_type: chunk_type.to_string(),
language: language.name().to_string(),
symbol_name,
content_hash: compute_content_hash(&chunk_content),
indexed_at: timestamp.to_string(),
parent_symbol,
signature,
doc_comment,
module_path: module_path.map(|s| s.to_string()),
};
chunks.push(chunk);
}
}
if cursor.goto_first_child() {
loop {
visit_node(
cursor,
content,
lines,
file_path,
language,
min_chunk_lines,
function_types,
class_types,
chunks,
timestamp,
module_path,
);
if !cursor.goto_next_sibling() {
break;
}
}
cursor.goto_parent();
}
}
visit_node(
&mut cursor,
content,
&lines,
file_path,
language,
min_chunk_lines,
function_types,
class_types,
chunks,
timestamp,
module_path,
);
}
fn extract_parent_symbol(
node: &tree_sitter::Node,
content: &str,
class_types: &[&str],
language: SupportedLanguage,
) -> Option<String> {
let mut current = node.parent()?;
loop {
if class_types.contains(¤t.kind()) {
let name = current
.child_by_field_name(language.name_field())
.and_then(|n| n.utf8_text(content.as_bytes()).ok())
.map(|s| s.to_string())
.filter(|s| !s.is_empty());
if name.is_some() {
return name;
}
}
match current.parent() {
Some(p) => current = p,
None => break,
}
}
None
}
fn extract_signature(node: &tree_sitter::Node, content: &str) -> Option<String> {
let content_bytes = content.as_bytes();
let body_start = node
.child_by_field_name("body")
.map(|b| b.start_byte())
.or_else(|| {
(0..node.child_count())
.filter_map(|i| node.child(i))
.find(|child| {
matches!(
child.kind(),
"body"
| "block"
| "declaration_list"
| "field_declaration_list"
| "class_body"
| "statement_block"
| "compound_statement"
| "function_body"
)
})
.map(|body_node| body_node.start_byte())
});
let sig_bytes = if let Some(body_start_byte) = body_start {
let node_start = node.start_byte();
if body_start_byte > node_start {
&content_bytes[node_start..body_start_byte]
} else {
return None;
}
} else {
let start = node.start_byte();
let end = node.end_byte().min(content_bytes.len());
let text = std::str::from_utf8(&content_bytes[start..end]).unwrap_or("");
let first_line = text.lines().next().unwrap_or("").trim();
return if first_line.is_empty() {
None
} else {
Some(first_line.to_string())
};
};
let sig_str = std::str::from_utf8(sig_bytes).unwrap_or("").trim();
if sig_str.is_empty() {
None
} else {
Some(sig_str.to_string())
}
}
fn extract_doc_comment(
node: &tree_sitter::Node,
content: &str,
language: SupportedLanguage,
) -> Option<String> {
let content_bytes = content.as_bytes();
if let Some(prev) = node.prev_sibling() {
let kind = prev.kind();
if kind.contains("comment") || kind.contains("doc") {
let text = prev
.utf8_text(content_bytes)
.unwrap_or("")
.trim()
.to_string();
if !text.is_empty() {
return Some(text);
}
}
if matches!(language, SupportedLanguage::Python)
&& (kind == "expression_statement" || kind == "string")
{
let string_node = if kind == "expression_statement" {
(0..prev.child_count())
.filter_map(|i| prev.child(i))
.find(|c| c.kind() == "string")
} else {
Some(prev)
};
if let Some(s) = string_node {
let text = s.utf8_text(content_bytes).unwrap_or("").trim().to_string();
if !text.is_empty() {
return Some(text);
}
}
}
}
if matches!(language, SupportedLanguage::Python) {
if let Some(body) = node.child_by_field_name("body") {
for i in 0..body.child_count() {
if let Some(child) = body.child(i) {
if !child.is_named() {
continue;
}
let kind = child.kind();
if kind == "string" {
let text = child
.utf8_text(content_bytes)
.unwrap_or("")
.trim()
.to_string();
if !text.is_empty() {
return Some(text);
}
} else if kind == "expression_statement" {
let string_node = (0..child.child_count())
.filter_map(|i| child.child(i))
.find(|c| c.kind() == "string");
if let Some(s) = string_node {
let text = s.utf8_text(content_bytes).unwrap_or("").trim().to_string();
if !text.is_empty() {
return Some(text);
}
}
}
break; }
}
}
}
None
}
fn compute_content_hash(content: &str) -> String {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(content.as_bytes());
hex::encode(hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_derive_module_path_rust() {
assert_eq!(derive_module_path("src/search/bm25.rs"), "search.bm25");
assert_eq!(derive_module_path("src/indexer/mod.rs"), "indexer");
assert_eq!(derive_module_path("src/lib.rs"), "lib");
assert_eq!(derive_module_path("src/db/mod.rs"), "db");
assert_eq!(derive_module_path("src/cli/search.rs"), "cli.search");
}
#[test]
fn test_derive_module_path_python() {
assert_eq!(derive_module_path("app/utils/__init__.py"), "app.utils");
assert_eq!(
derive_module_path("app/utils/helpers.py"),
"app.utils.helpers"
);
}
#[test]
fn test_derive_module_path_no_prefix() {
assert_eq!(derive_module_path("main.rs"), "main");
assert_eq!(derive_module_path("utils/helpers.ts"), "utils.helpers");
}
#[test]
fn test_derive_module_path_windows_separators() {
assert_eq!(derive_module_path("src\\search\\bm25.rs"), "search.bm25");
}
fn default_config() -> ChunkConfig {
ChunkConfig {
min_chunk_lines: 1,
max_chunk_lines: 10,
overlap_lines: 2,
include_file_chunks: true,
}
}
#[test]
fn test_chunk_plain_text_empty_file() {
let cfg = default_config();
let result = chunk_plain_text("", "README.md", "markdown", &cfg).unwrap();
assert!(result.is_empty(), "empty file should produce no chunks");
}
#[test]
fn test_chunk_plain_text_whitespace_only() {
let cfg = default_config();
let result = chunk_plain_text(" \n\n ", "README.md", "markdown", &cfg).unwrap();
assert!(
result.is_empty(),
"whitespace-only file should produce no chunks"
);
}
#[test]
fn test_chunk_markdown_sections() {
let cfg = default_config();
let content = "# Introduction\nHello world.\n\n# Usage\nRun `cargo build`.\n";
let chunks = chunk_plain_text(content, "README.md", "markdown", &cfg).unwrap();
assert_eq!(chunks.len(), 2);
let intro = &chunks[0];
assert_eq!(intro.chunk_type, "section");
assert_eq!(intro.symbol_name.as_deref(), Some("Introduction"));
assert_eq!(intro.language, "markdown");
let usage = &chunks[1];
assert_eq!(usage.chunk_type, "section");
assert_eq!(usage.symbol_name.as_deref(), Some("Usage"));
}
#[test]
fn test_chunk_markdown_no_headings_falls_back_to_blocks() {
let cfg = default_config();
let content = "Just some prose\nwith multiple lines\nbut no headings at all.\n";
let chunks = chunk_plain_text(content, "notes.md", "markdown", &cfg).unwrap();
assert!(!chunks.is_empty(), "should produce at least one chunk");
for chunk in &chunks {
assert_eq!(chunk.chunk_type, "block");
assert_eq!(chunk.language, "markdown");
}
}
#[test]
fn test_chunk_plain_text_yaml_produces_blocks() {
let cfg = default_config();
let content = "name: my-app\nversion: 1.0\nenv: production\n";
let chunks = chunk_plain_text(content, ".github/workflows/ci.yml", "yaml", &cfg).unwrap();
assert!(!chunks.is_empty());
for chunk in &chunks {
assert_eq!(chunk.chunk_type, "block");
assert_eq!(chunk.language, "yaml");
}
}
#[test]
fn test_chunk_plain_text_overlapping_large_file() {
let cfg = ChunkConfig {
max_chunk_lines: 5,
overlap_lines: 2,
..default_config()
};
let lines: Vec<String> = (1..=12).map(|i| format!("line {i}")).collect();
let content = lines.join("\n");
let chunks = chunk_plain_text(&content, "big.sql", "sql", &cfg).unwrap();
assert!(
chunks.len() > 1,
"large file should split into multiple chunks"
);
for chunk in &chunks {
assert_eq!(chunk.chunk_type, "block");
}
assert_eq!(chunks[0].start_line, 1);
let last = chunks.last().unwrap();
assert_eq!(last.end_line, 12);
}
#[test]
fn test_chunk_markdown_heading_only_last_line() {
let cfg = default_config();
let content = "Some intro text.\n# Trailing Heading";
let chunks = chunk_plain_text(content, "README.md", "markdown", &cfg).unwrap();
let trailing = chunks
.iter()
.find(|c| c.symbol_name.as_deref() == Some("Trailing Heading"));
assert!(
trailing.is_some(),
"heading on last line should create a section chunk"
);
}
}