ck_chunk/

lib.rs

use anyhow::Result;
use ck_core::Span;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Chunk {
    pub span: Span,
    pub text: String,
    pub chunk_type: ChunkType,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ChunkType {
    Text,
    Function,
    Class,
    Method,
    Module,
}

pub fn chunk_text(text: &str, language: Option<&str>) -> Result<Vec<Chunk>> {
    tracing::debug!("Chunking text with language: {:?}, length: {} chars", language, text.len());
    
    let result = match language {
        Some("python") => {
            tracing::debug!("Using Python tree-sitter parser");
            chunk_python(text)
        },
        Some("typescript") | Some("javascript") => {
            tracing::debug!("Using TypeScript/JavaScript tree-sitter parser");
            chunk_typescript(text)
        },
        Some("haskell") => {
            tracing::debug!("Using Haskell tree-sitter parser");
            chunk_haskell(text)
        },
        Some("rust") => {
            tracing::debug!("Using Rust tree-sitter parser");
            chunk_rust(text)
        },
        Some("ruby") => {
            tracing::debug!("Using Ruby tree-sitter parser");
            chunk_ruby(text)
        },
        _ => {
            tracing::debug!("Using generic chunking strategy");
            chunk_generic(text)
        },
    };
    
    match &result {
        Ok(chunks) => tracing::debug!("Successfully created {} chunks", chunks.len()),
        Err(e) => tracing::warn!("Chunking failed: {}", e),
    }
    
    result
}

fn chunk_generic(text: &str) -> Result<Vec<Chunk>> {
    let mut chunks = Vec::new();
    let lines: Vec<&str> = text.lines().collect();
    let chunk_size = 20;
    let overlap = 5;
    
    // Pre-compute cumulative byte offsets for O(1) lookup
    let mut line_byte_offsets = Vec::with_capacity(lines.len() + 1);
    line_byte_offsets.push(0);
    let mut cumulative_offset = 0;
    for line in &lines {
        cumulative_offset += line.len() + 1; // +1 for newline
        line_byte_offsets.push(cumulative_offset);
    }
    
    let mut i = 0;
    while i < lines.len() {
        let end = (i + chunk_size).min(lines.len());
        let chunk_lines = &lines[i..end];
        let chunk_text = chunk_lines.join("\n");
        
        let byte_start = line_byte_offsets[i];
        let byte_end = byte_start + chunk_text.len();
        
        chunks.push(Chunk {
            span: Span {
                byte_start,
                byte_end,
                line_start: i + 1,
                line_end: end,
            },
            text: chunk_text,
            chunk_type: ChunkType::Text,
        });
        
        i += chunk_size - overlap;
        if i >= lines.len() {
            break;
        }
    }
    
    Ok(chunks)
}

fn chunk_python(text: &str) -> Result<Vec<Chunk>> {
    let mut parser = tree_sitter::Parser::new();
    parser.set_language(&tree_sitter_python::language())?;
    
    let tree = parser.parse(text, None).ok_or_else(|| {
        anyhow::anyhow!("Failed to parse Python code")
    })?;
    
    let mut chunks = Vec::new();
    let mut cursor = tree.root_node().walk();
    
    extract_code_chunks(&mut cursor, text, &mut chunks, "python");
    
    if chunks.is_empty() {
        return chunk_generic(text);
    }
    
    Ok(chunks)
}

fn chunk_typescript(text: &str) -> Result<Vec<Chunk>> {
    let mut parser = tree_sitter::Parser::new();
    parser.set_language(&tree_sitter_typescript::language_typescript())?;
    
    let tree = parser.parse(text, None).ok_or_else(|| {
        anyhow::anyhow!("Failed to parse TypeScript code")
    })?;
    
    let mut chunks = Vec::new();
    let mut cursor = tree.root_node().walk();
    
    extract_code_chunks(&mut cursor, text, &mut chunks, "typescript");
    
    if chunks.is_empty() {
        return chunk_generic(text);
    }
    
    Ok(chunks)
}

fn chunk_haskell(text: &str) -> Result<Vec<Chunk>> {
    let mut parser = tree_sitter::Parser::new();
    parser.set_language(&tree_sitter_haskell::language())?;
    
    let tree = parser.parse(text, None).ok_or_else(|| {
        anyhow::anyhow!("Failed to parse Haskell code")
    })?;
    
    let mut chunks = Vec::new();
    let mut cursor = tree.root_node().walk();
    
    extract_code_chunks(&mut cursor, text, &mut chunks, "haskell");
    
    if chunks.is_empty() {
        return chunk_generic(text);
    }
    
    Ok(chunks)
}

fn chunk_rust(text: &str) -> Result<Vec<Chunk>> {
    let mut parser = tree_sitter::Parser::new();
    parser.set_language(&tree_sitter_rust::language())?;
    
    let tree = parser.parse(text, None).ok_or_else(|| {
        anyhow::anyhow!("Failed to parse Rust code")
    })?;
    
    let mut chunks = Vec::new();
    let mut cursor = tree.root_node().walk();
    
    extract_code_chunks(&mut cursor, text, &mut chunks, "rust");
    
    if chunks.is_empty() {
        return chunk_generic(text);
    }
    
    Ok(chunks)
}


fn chunk_ruby(text: &str) -> Result<Vec<Chunk>> {
    let mut parser = tree_sitter::Parser::new();
    parser.set_language(&tree_sitter_ruby::language())?;
    
    let tree = parser.parse(text, None).ok_or_else(|| {
        anyhow::anyhow!("Failed to parse Ruby code")
    })?;
    
    let mut chunks = Vec::new();
    let mut cursor = tree.root_node().walk();
    
    extract_code_chunks(&mut cursor, text, &mut chunks, "ruby");
    
    if chunks.is_empty() {
        return chunk_generic(text);
    }
    
    Ok(chunks)
}


fn extract_code_chunks(
    cursor: &mut tree_sitter::TreeCursor,
    source: &str,
    chunks: &mut Vec<Chunk>,
    language: &str,
) {
    let node = cursor.node();
    let node_kind = node.kind();
    
    
    let is_chunk = match language {
        "python" => matches!(node_kind, "function_definition" | "class_definition"),
        "typescript" | "javascript" => matches!(
            node_kind,
            "function_declaration" | "class_declaration" | "method_definition" | "arrow_function"
        ),
        "haskell" => matches!(
            node_kind,
            "signature" | "data_type" | "newtype" | "type_synomym" | "type_family" | "class" | "instance"
        ),
        "rust" => matches!(
            node_kind,
            "function_item" | "impl_item" | "struct_item" | "enum_item" | "trait_item" | "mod_item"
        ),
        "ruby" => matches!(
            node_kind,
            "method" | "class" | "module" | "singleton_method"
        ),
        _ => false,
    };
    
    if is_chunk {
        let start_byte = node.start_byte();
        let end_byte = node.end_byte();
        let start_pos = node.start_position();
        let end_pos = node.end_position();
        
        let text = &source[start_byte..end_byte];
        
        let chunk_type = match node_kind {
            "function_definition" | "function_declaration" | "arrow_function" | "function" | "signature" | "function_item" | "def" | "defp" | "method" | "singleton_method" | "defn" | "defn-" => ChunkType::Function,
            "class_definition" | "class_declaration" | "instance_declaration" | "class" | "instance" | "struct_item" | "enum_item" | "defstruct" | "defrecord" | "deftype" => ChunkType::Class,
            "method_definition" | "defmacro" => ChunkType::Method,
            "data_type" | "newtype" | "type_synomym" | "type_family" | "impl_item" | "trait_item" | "mod_item" | "defmodule" | "module" | "defprotocol" | "ns" => ChunkType::Module,
            _ => ChunkType::Text,
        };
        
        chunks.push(Chunk {
            span: Span {
                byte_start: start_byte,
                byte_end: end_byte,
                line_start: start_pos.row + 1,
                line_end: end_pos.row + 1,
            },
            text: text.to_string(),
            chunk_type,
        });
    }
    
    if cursor.goto_first_child() {
        loop {
            extract_code_chunks(cursor, source, chunks, language);
            if !cursor.goto_next_sibling() {
                break;
            }
        }
        cursor.goto_parent();
    }
}

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

    #[test]
    fn test_chunk_generic_byte_offsets() {
        // Test that byte offsets are calculated correctly using O(n) algorithm
        let text = "line 1\nline 2\nline 3\nline 4\nline 5";
        let chunks = chunk_generic(text).unwrap();
        
        assert!(!chunks.is_empty());
        
        // First chunk should start at byte 0
        assert_eq!(chunks[0].span.byte_start, 0);
        
        // Each chunk's byte_end should match the actual text length
        for chunk in &chunks {
            let expected_len = chunk.text.len();
            let actual_len = chunk.span.byte_end - chunk.span.byte_start;
            assert_eq!(actual_len, expected_len);
        }
    }

    #[test]
    fn test_chunk_generic_large_file_performance() {
        // Create a large text to ensure O(n) performance
        let lines: Vec<String> = (0..1000).map(|i| format!("Line {}: Some content here", i)).collect();
        let text = lines.join("\n");
        
        let start = std::time::Instant::now();
        let chunks = chunk_generic(&text).unwrap();
        let duration = start.elapsed();
        
        // Should complete quickly even for 1000 lines
        assert!(duration.as_millis() < 100, "Chunking took too long: {:?}", duration);
        assert!(!chunks.is_empty());
        
        // Verify chunks have correct line numbers
        for chunk in &chunks {
            assert!(chunk.span.line_start > 0);
            assert!(chunk.span.line_end >= chunk.span.line_start);
        }
    }

    #[test]
    fn test_chunk_rust() {
        let rust_code = r#"
pub struct Calculator {
    memory: f64,
}

impl Calculator {
    pub fn new() -> Self {
        Calculator { memory: 0.0 }
    }
    
    pub fn add(&mut self, a: f64, b: f64) -> f64 {
        a + b
    }
}

fn main() {
    let calc = Calculator::new();
}

pub mod utils {
    pub fn helper() {}
}
"#;
        
        let chunks = chunk_rust(rust_code).unwrap();
        assert!(!chunks.is_empty());
        
        // Should find struct, impl, functions, and module
        let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
        assert!(chunk_types.contains(&&ChunkType::Class));  // struct
        assert!(chunk_types.contains(&&ChunkType::Module)); // impl and mod
        assert!(chunk_types.contains(&&ChunkType::Function)); // functions
    }

    #[test]
    fn test_chunk_ruby() {
        let ruby_code = r#"
class Calculator
  def initialize
    @memory = 0.0
  end

  def add(a, b)
    a + b
  end

  def self.class_method
    "class method"
  end

  private

  def private_method
    "private"
  end
end

module Utils
  def self.helper
    "helper"
  end
end

def main
  calc = Calculator.new
end
"#;
        
        let chunks = chunk_ruby(ruby_code).unwrap();
        assert!(!chunks.is_empty());
        
        // Should find class, module, and methods
        let chunk_types: Vec<&ChunkType> = chunks.iter().map(|c| &c.chunk_type).collect();
        assert!(chunk_types.contains(&&ChunkType::Class));    // class
        assert!(chunk_types.contains(&&ChunkType::Module));   // module
        assert!(chunk_types.contains(&&ChunkType::Function)); // methods
    }

    #[test]
    fn test_language_detection_fallback() {
        // Test that unknown languages fall back to generic chunking
        let generic_text = "Some text\nwith multiple lines\nto chunk generically";
        
        let chunks_unknown = chunk_text(generic_text, Some("unknown_language")).unwrap();
        let chunks_generic = chunk_generic(generic_text).unwrap();
        
        // Should produce the same result
        assert_eq!(chunks_unknown.len(), chunks_generic.len());
        assert_eq!(chunks_unknown[0].text, chunks_generic[0].text);
    }
}