aurora_semantic/chunker/

mod.rs

//! Code chunking module for extracting meaningful code segments.
//!
//! This module provides functionality to parse source code and extract
//! semantically meaningful chunks (functions, classes, etc.).

mod strategies;

use crate::config::ChunkingConfig;
use crate::error::Result;
use crate::types::{Chunk, ChunkId, ChunkType, Document, Language};

/// Trait for code chunking implementations.
pub trait Chunker: Send + Sync {
    /// Extract chunks from source code content.
    fn chunk(&self, document: &Document, content: &str) -> Result<Vec<Chunk>>;

    /// Get the name of this chunker.
    fn name(&self) -> &'static str;
}

/// Default chunker using regex-based parsing.
pub struct DefaultChunker {
    config: ChunkingConfig,
}

impl DefaultChunker {
    /// Create a new default chunker with the given configuration.
    pub fn new(config: ChunkingConfig) -> Result<Self> {
        Ok(Self { config })
    }
}

impl Chunker for DefaultChunker {
    fn chunk(&self, document: &Document, content: &str) -> Result<Vec<Chunk>> {
        // Skip if content is too small
        if content.len() < self.config.min_chunk_size {
            return Ok(vec![]);
        }

        // Use semantic chunking based on language patterns
        let chunks = extract_chunks_by_language(document, content, &self.config);

        // If we got no chunks, fall back to line-based chunking
        if chunks.is_empty() && content.len() >= self.config.min_chunk_size {
            return Ok(vec![Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: content.to_string(),
                chunk_type: ChunkType::Block,
                start_line: 1,
                end_line: content.lines().count() as u32,
                start_byte: 0,
                end_byte: content.len(),
                symbol_name: None,
                parent_symbol: None,
            }]);
        }

        Ok(chunks)
    }

    fn name(&self) -> &'static str {
        "default"
    }
}

/// Extract chunks based on language-specific patterns.
fn extract_chunks_by_language(
    document: &Document,
    content: &str,
    config: &ChunkingConfig,
) -> Vec<Chunk> {
    match document.language {
        Language::Rust => extract_rust_chunks(document, content, config),
        Language::Python => extract_python_chunks(document, content, config),
        Language::JavaScript | Language::TypeScript => {
            extract_js_chunks(document, content, config)
        }
        Language::Go => extract_go_chunks(document, content, config),
        Language::Java => extract_java_chunks(document, content, config),
        _ => extract_generic_chunks(document, content, config),
    }
}

/// Extract Rust functions/structs/impls.
fn extract_rust_chunks(document: &Document, content: &str, config: &ChunkingConfig) -> Vec<Chunk> {
    let mut chunks = Vec::new();
    let lines: Vec<&str> = content.lines().collect();

    let mut i = 0;
    while i < lines.len() {
        let line = lines[i].trim();

        // Detect function, struct, enum, impl, trait, mod
        let (chunk_type, name) = if line.starts_with("pub fn ")
            || line.starts_with("fn ")
            || line.starts_with("pub async fn ")
            || line.starts_with("async fn ")
        {
            (ChunkType::Function, extract_name_after(line, "fn "))
        } else if line.starts_with("pub struct ") || line.starts_with("struct ") {
            (ChunkType::Struct, extract_name_after(line, "struct "))
        } else if line.starts_with("pub enum ") || line.starts_with("enum ") {
            (ChunkType::Enum, extract_name_after(line, "enum "))
        } else if line.starts_with("impl ") {
            (ChunkType::Implementation, extract_impl_name(line))
        } else if line.starts_with("pub trait ") || line.starts_with("trait ") {
            (ChunkType::Interface, extract_name_after(line, "trait "))
        } else if line.starts_with("mod ") || line.starts_with("pub mod ") {
            (ChunkType::Module, extract_name_after(line, "mod "))
        } else {
            i += 1;
            continue;
        };

        // Find the end of the block
        let start_line = i;
        let end_line = find_block_end(&lines, i);

        let chunk_content: String = lines[start_line..=end_line].join("\n");

        if chunk_content.len() >= config.min_chunk_size
            && chunk_content.len() <= config.max_chunk_size
        {
            let start_byte = lines[..start_line].iter().map(|l| l.len() + 1).sum();
            let end_byte = start_byte + chunk_content.len();

            chunks.push(Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: chunk_content,
                chunk_type,
                start_line: start_line as u32 + 1,
                end_line: end_line as u32 + 1,
                start_byte,
                end_byte,
                symbol_name: name,
                parent_symbol: None,
            });
        }

        i = end_line + 1;
    }

    chunks
}

/// Extract Python functions/classes.
fn extract_python_chunks(
    document: &Document,
    content: &str,
    config: &ChunkingConfig,
) -> Vec<Chunk> {
    let mut chunks = Vec::new();
    let lines: Vec<&str> = content.lines().collect();

    let mut i = 0;
    while i < lines.len() {
        let line = lines[i];
        let trimmed = line.trim();

        let (chunk_type, name) = if trimmed.starts_with("def ")
            || trimmed.starts_with("async def ")
        {
            (ChunkType::Function, extract_name_after(trimmed, "def "))
        } else if trimmed.starts_with("class ") {
            (ChunkType::Class, extract_name_after(trimmed, "class "))
        } else {
            i += 1;
            continue;
        };

        // Find indentation level
        let indent = line.len() - line.trim_start().len();
        let start_line = i;
        let end_line = find_python_block_end(&lines, i, indent);

        let chunk_content: String = lines[start_line..=end_line].join("\n");

        if chunk_content.len() >= config.min_chunk_size
            && chunk_content.len() <= config.max_chunk_size
        {
            let start_byte = lines[..start_line].iter().map(|l| l.len() + 1).sum();
            let end_byte = start_byte + chunk_content.len();

            chunks.push(Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: chunk_content,
                chunk_type,
                start_line: start_line as u32 + 1,
                end_line: end_line as u32 + 1,
                start_byte,
                end_byte,
                symbol_name: name,
                parent_symbol: None,
            });
        }

        i = end_line + 1;
    }

    chunks
}

/// Extract JavaScript/TypeScript functions/classes.
fn extract_js_chunks(document: &Document, content: &str, config: &ChunkingConfig) -> Vec<Chunk> {
    let mut chunks = Vec::new();
    let lines: Vec<&str> = content.lines().collect();

    let mut i = 0;
    while i < lines.len() {
        let line = lines[i].trim();

        let (chunk_type, name) = if line.starts_with("function ")
            || line.starts_with("async function ")
            || line.starts_with("export function ")
            || line.starts_with("export async function ")
        {
            (ChunkType::Function, extract_name_after(line, "function "))
        } else if line.starts_with("class ") || line.starts_with("export class ") {
            (ChunkType::Class, extract_name_after(line, "class "))
        } else if line.starts_with("interface ") || line.starts_with("export interface ") {
            (ChunkType::Interface, extract_name_after(line, "interface "))
        } else if line.contains("=>") && (line.starts_with("const ") || line.starts_with("export const ")) {
            (ChunkType::Function, extract_name_after(line, "const "))
        } else {
            i += 1;
            continue;
        };

        let start_line = i;
        let end_line = find_block_end(&lines, i);

        let chunk_content: String = lines[start_line..=end_line].join("\n");

        if chunk_content.len() >= config.min_chunk_size
            && chunk_content.len() <= config.max_chunk_size
        {
            let start_byte = lines[..start_line].iter().map(|l| l.len() + 1).sum();
            let end_byte = start_byte + chunk_content.len();

            chunks.push(Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: chunk_content,
                chunk_type,
                start_line: start_line as u32 + 1,
                end_line: end_line as u32 + 1,
                start_byte,
                end_byte,
                symbol_name: name,
                parent_symbol: None,
            });
        }

        i = end_line + 1;
    }

    chunks
}

/// Extract Go functions/structs.
fn extract_go_chunks(document: &Document, content: &str, config: &ChunkingConfig) -> Vec<Chunk> {
    let mut chunks = Vec::new();
    let lines: Vec<&str> = content.lines().collect();

    let mut i = 0;
    while i < lines.len() {
        let line = lines[i].trim();

        let (chunk_type, name) = if line.starts_with("func ") {
            (ChunkType::Function, extract_go_func_name(line))
        } else if line.starts_with("type ") && line.contains("struct") {
            (ChunkType::Struct, extract_name_after(line, "type "))
        } else if line.starts_with("type ") && line.contains("interface") {
            (ChunkType::Interface, extract_name_after(line, "type "))
        } else {
            i += 1;
            continue;
        };

        let start_line = i;
        let end_line = find_block_end(&lines, i);

        let chunk_content: String = lines[start_line..=end_line].join("\n");

        if chunk_content.len() >= config.min_chunk_size
            && chunk_content.len() <= config.max_chunk_size
        {
            let start_byte = lines[..start_line].iter().map(|l| l.len() + 1).sum();
            let end_byte = start_byte + chunk_content.len();

            chunks.push(Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: chunk_content,
                chunk_type,
                start_line: start_line as u32 + 1,
                end_line: end_line as u32 + 1,
                start_byte,
                end_byte,
                symbol_name: name,
                parent_symbol: None,
            });
        }

        i = end_line + 1;
    }

    chunks
}

/// Extract Java methods/classes.
fn extract_java_chunks(document: &Document, content: &str, config: &ChunkingConfig) -> Vec<Chunk> {
    let mut chunks = Vec::new();
    let lines: Vec<&str> = content.lines().collect();

    let mut i = 0;
    while i < lines.len() {
        let line = lines[i].trim();

        let (chunk_type, name) =
            if line.contains("class ") && (line.starts_with("public") || line.starts_with("class"))
            {
                (ChunkType::Class, extract_name_after(line, "class "))
            } else if line.contains("interface ")
                && (line.starts_with("public") || line.starts_with("interface"))
            {
                (ChunkType::Interface, extract_name_after(line, "interface "))
            } else if (line.starts_with("public ")
                || line.starts_with("private ")
                || line.starts_with("protected "))
                && line.contains("(")
                && !line.contains("class")
            {
                (ChunkType::Function, extract_java_method_name(line))
            } else {
                i += 1;
                continue;
            };

        let start_line = i;
        let end_line = find_block_end(&lines, i);

        let chunk_content: String = lines[start_line..=end_line].join("\n");

        if chunk_content.len() >= config.min_chunk_size
            && chunk_content.len() <= config.max_chunk_size
        {
            let start_byte = lines[..start_line].iter().map(|l| l.len() + 1).sum();
            let end_byte = start_byte + chunk_content.len();

            chunks.push(Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: chunk_content,
                chunk_type,
                start_line: start_line as u32 + 1,
                end_line: end_line as u32 + 1,
                start_byte,
                end_byte,
                symbol_name: name,
                parent_symbol: None,
            });
        }

        i = end_line + 1;
    }

    chunks
}

/// Generic chunking for unknown languages.
fn extract_generic_chunks(
    document: &Document,
    content: &str,
    config: &ChunkingConfig,
) -> Vec<Chunk> {
    let lines: Vec<&str> = content.lines().collect();
    let mut chunks = Vec::new();

    // Split into chunks by blank lines or size limits
    let mut start = 0;
    let mut current_chunk = String::new();

    for (i, line) in lines.iter().enumerate() {
        current_chunk.push_str(line);
        current_chunk.push('\n');

        // Check if we should create a chunk
        let should_split = current_chunk.len() >= config.max_chunk_size
            || (line.is_empty() && current_chunk.len() >= config.min_chunk_size);

        if should_split {
            let start_byte: usize = lines[..start].iter().map(|l| l.len() + 1).sum();

            chunks.push(Chunk {
                id: ChunkId::new(),
                document_id: document.id.clone(),
                content: current_chunk.trim().to_string(),
                chunk_type: ChunkType::Block,
                start_line: start as u32 + 1,
                end_line: i as u32 + 1,
                start_byte,
                end_byte: start_byte + current_chunk.len(),
                symbol_name: None,
                parent_symbol: None,
            });

            current_chunk.clear();
            start = i + 1;
        }
    }

    // Don't forget the last chunk
    if current_chunk.len() >= config.min_chunk_size {
        let start_byte: usize = lines[..start].iter().map(|l| l.len() + 1).sum();

        chunks.push(Chunk {
            id: ChunkId::new(),
            document_id: document.id.clone(),
            content: current_chunk.trim().to_string(),
            chunk_type: ChunkType::Block,
            start_line: start as u32 + 1,
            end_line: lines.len() as u32,
            start_byte,
            end_byte: start_byte + current_chunk.len(),
            symbol_name: None,
            parent_symbol: None,
        });
    }

    chunks
}

// Helper functions

fn extract_name_after(line: &str, keyword: &str) -> Option<String> {
    line.find(keyword).and_then(|idx| {
        let rest = &line[idx + keyword.len()..];
        let end = rest
            .find(|c: char| !c.is_alphanumeric() && c != '_')
            .unwrap_or(rest.len());
        if end > 0 {
            Some(rest[..end].to_string())
        } else {
            None
        }
    })
}

fn extract_impl_name(line: &str) -> Option<String> {
    // impl Type or impl Trait for Type
    let rest = line.strip_prefix("impl")?.trim();
    let end = rest
        .find(|c: char| !c.is_alphanumeric() && c != '_' && c != '<')
        .unwrap_or(rest.len());
    if end > 0 {
        Some(rest[..end].to_string())
    } else {
        None
    }
}

fn extract_go_func_name(line: &str) -> Option<String> {
    // func Name(...) or func (r Receiver) Name(...)
    let rest = line.strip_prefix("func")?.trim();
    if rest.starts_with('(') {
        // Method with receiver
        let after_receiver = rest.find(')')? + 1;
        let name_part = rest[after_receiver..].trim();
        extract_name_after(name_part, "")
    } else {
        extract_name_after(rest, "")
    }
}

fn extract_java_method_name(line: &str) -> Option<String> {
    // public void methodName(...)
    let paren_idx = line.find('(')?;
    let before_paren = &line[..paren_idx];
    let words: Vec<&str> = before_paren.split_whitespace().collect();
    words.last().map(|s| s.to_string())
}

fn find_block_end(lines: &[&str], start: usize) -> usize {
    let mut brace_count = 0;
    let mut found_open = false;

    for (i, line) in lines.iter().enumerate().skip(start) {
        for c in line.chars() {
            if c == '{' {
                brace_count += 1;
                found_open = true;
            } else if c == '}' {
                brace_count -= 1;
            }
        }

        if found_open && brace_count == 0 {
            return i;
        }
    }

    lines.len().saturating_sub(1)
}

fn find_python_block_end(lines: &[&str], start: usize, base_indent: usize) -> usize {
    for (i, line) in lines.iter().enumerate().skip(start + 1) {
        if line.trim().is_empty() {
            continue;
        }
        let indent = line.len() - line.trim_start().len();
        if indent <= base_indent {
            return i.saturating_sub(1);
        }
    }
    lines.len().saturating_sub(1)
}

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

    fn make_document(language: Language) -> Document {
        Document {
            id: crate::types::DocumentId::new(),
            relative_path: PathBuf::from("test.rs"),
            absolute_path: PathBuf::from("/test/test.rs"),
            language,
            content_hash: "test".to_string(),
            size_bytes: 0,
            modified_at: chrono::Utc::now(),
        }
    }

    #[test]
    fn test_rust_chunking() {
        let doc = make_document(Language::Rust);
        let content = r#"
fn hello() {
    println!("Hello");
}

pub struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn new() -> Self {
        Self { x: 0, y: 0 }
    }
}
"#;
        let config = ChunkingConfig::default();
        let chunks = extract_rust_chunks(&doc, content, &config);

        assert!(chunks.len() >= 2);
        assert!(chunks.iter().any(|c| c.symbol_name == Some("hello".to_string())));
    }

    #[test]
    fn test_python_chunking() {
        let doc = make_document(Language::Python);
        let content = r#"
def hello():
    print("Hello")

class MyClass:
    def __init__(self):
        pass
"#;
        let config = ChunkingConfig::default();
        let chunks = extract_python_chunks(&doc, content, &config);

        assert!(!chunks.is_empty());
    }
}