codemem_engine/index/

chunker.rs

//! CST-aware code chunking.
//!
//! Splits source files into semantically meaningful chunks using the
//! concrete syntax tree (CST) produced by ast-grep/tree-sitter. The algorithm:
//!
//! 1. If a CST node fits within `max_chunk_size` (non-whitespace chars) -> emit it as a chunk.
//! 2. If too large -> recurse into named children.
//! 3. Adjacent small siblings are merged greedily until the merged size would exceed `max_chunk_size`.
//!
//! Each chunk records its parent symbol (resolved by line-range containment).

use crate::index::symbol::Symbol;
use ast_grep_core::{Doc, Node};
use serde::{Deserialize, Serialize};

/// A code chunk produced by the CST-aware chunker.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeChunk {
    /// 0-based index of this chunk within the file.
    pub index: usize,
    /// The source text of this chunk.
    pub text: String,
    /// The tree-sitter node kind (e.g., "function_item", "impl_item").
    pub node_kind: String,
    /// 0-based starting line.
    pub line_start: usize,
    /// 0-based ending line.
    pub line_end: usize,
    /// Byte offset start.
    pub byte_start: usize,
    /// Byte offset end.
    pub byte_end: usize,
    /// Count of non-whitespace characters.
    pub non_ws_chars: usize,
    /// Qualified name of the innermost containing symbol, if any.
    pub parent_symbol: Option<String>,
    /// Path of the source file.
    pub file_path: String,
}

/// Configuration for the chunker.
#[derive(Debug, Clone)]
pub struct ChunkConfig {
    /// Maximum chunk size in non-whitespace characters.
    pub max_chunk_size: usize,
    /// Minimum chunk size in non-whitespace characters.
    pub min_chunk_size: usize,
    /// Number of lines to overlap between adjacent chunks (0 = no overlap).
    pub overlap_lines: usize,
}

impl Default for ChunkConfig {
    fn default() -> Self {
        Self {
            max_chunk_size: 1500,
            min_chunk_size: 50,
            overlap_lines: 0,
        }
    }
}

/// Count non-whitespace characters in a string.
fn count_non_ws(s: &str) -> usize {
    s.chars().filter(|c| !c.is_whitespace()).count()
}

/// Intermediate chunk before index assignment and parent resolution.
struct RawChunk {
    text: String,
    node_kind: String,
    line_start: usize,
    line_end: usize,
    byte_start: usize,
    byte_end: usize,
    non_ws_chars: usize,
}

/// Chunk a file using its CST tree.
///
/// - `root`: the ast-grep parsed root (AstGrep instance).
/// - `source`: the raw source string.
/// - `file_path`: path to the source file (stored in each chunk).
/// - `symbols`: symbols extracted from this file (used for parent resolution).
/// - `config`: chunking parameters.
pub fn chunk_file<D: Doc>(
    root: &ast_grep_core::AstGrep<D>,
    source: &str,
    file_path: &str,
    symbols: &[Symbol],
    config: &ChunkConfig,
) -> Vec<CodeChunk>
where
    D::Lang: ast_grep_core::Language,
{
    if source.trim().is_empty() {
        return Vec::new();
    }

    let root_node = root.root();
    let mut raw_chunks = Vec::new();
    collect_chunks(&root_node, config, &mut raw_chunks);

    // Merge adjacent small chunks greedily
    let merged = merge_small_chunks(raw_chunks, source, config);

    // C3: Apply overlap — prepend trailing lines from previous chunk
    let merged = if config.overlap_lines > 0 {
        apply_overlap(merged, source, config.overlap_lines)
    } else {
        merged
    };

    // C2: Build interval index once for O(log n) parent resolution per chunk
    let interval_index = SymbolIntervalIndex::build(symbols);

    // Assign indices and resolve parent symbols
    merged
        .into_iter()
        .enumerate()
        .map(|(idx, raw)| {
            let parent_symbol = interval_index
                .resolve(raw.line_start, raw.line_end)
                .map(|s| s.qualified_name.clone());
            CodeChunk {
                index: idx,
                text: raw.text,
                node_kind: raw.node_kind,
                line_start: raw.line_start,
                line_end: raw.line_end,
                byte_start: raw.byte_start,
                byte_end: raw.byte_end,
                non_ws_chars: raw.non_ws_chars,
                parent_symbol,
                file_path: file_path.to_string(),
            }
        })
        .collect()
}

/// Recursively collect chunks from a CST node.
fn collect_chunks<D: Doc>(node: &Node<'_, D>, config: &ChunkConfig, out: &mut Vec<RawChunk>)
where
    D::Lang: ast_grep_core::Language,
{
    let text = node.text();
    let nws = count_non_ws(&text);

    // If the node fits, emit it as a single chunk
    if nws <= config.max_chunk_size {
        let range = node.range();
        out.push(RawChunk {
            text: text.to_string(),
            node_kind: node.kind().to_string(),
            line_start: node.start_pos().line(),
            line_end: node.end_pos().line(),
            byte_start: range.start,
            byte_end: range.end,
            non_ws_chars: nws,
        });
        return;
    }

    // Too large: recurse into named children
    let named_children: Vec<_> = node.children().filter(|c| c.is_named()).collect();
    if named_children.is_empty() {
        // No named children (e.g., a very large string literal), emit as-is
        let range = node.range();
        out.push(RawChunk {
            text: text.to_string(),
            node_kind: node.kind().to_string(),
            line_start: node.start_pos().line(),
            line_end: node.end_pos().line(),
            byte_start: range.start,
            byte_end: range.end,
            non_ws_chars: nws,
        });
    } else {
        for child in &named_children {
            collect_chunks(child, config, out);
        }
    }
}

/// Merge adjacent small chunks greedily.
fn merge_small_chunks(chunks: Vec<RawChunk>, source: &str, config: &ChunkConfig) -> Vec<RawChunk> {
    if chunks.is_empty() {
        return Vec::new();
    }

    let mut result: Vec<RawChunk> = Vec::new();

    for chunk in chunks {
        if let Some(last) = result.last_mut() {
            // If both the current accumulator and new chunk are small, try merging
            if last.non_ws_chars < config.min_chunk_size
                || chunk.non_ws_chars < config.min_chunk_size
            {
                // Compute actual merged non-whitespace count before deciding to merge
                let merged_start = last.byte_start;
                let merged_end = chunk.byte_end;
                let merged_text = if merged_end <= source.len() {
                    source[merged_start..merged_end].to_string()
                } else {
                    format!("{}\n{}", last.text, chunk.text)
                };
                let merged_nws = count_non_ws(&merged_text);

                if merged_nws <= config.max_chunk_size {
                    last.text = merged_text;
                    // C4: Preserve individual node_kinds as comma-separated
                    if last.node_kind.contains(&chunk.node_kind) {
                        // Already contains this kind, no-op
                    } else {
                        last.node_kind = format!("{},{}", last.node_kind, chunk.node_kind);
                    }
                    last.line_end = chunk.line_end;
                    last.byte_end = merged_end;
                    last.non_ws_chars = merged_nws;
                    continue;
                }
            }
        }
        result.push(chunk);
    }

    result
}

/// C3: Apply overlap between adjacent chunks by prepending trailing lines
/// from the previous chunk to the current one.
fn apply_overlap(chunks: Vec<RawChunk>, source: &str, overlap_lines: usize) -> Vec<RawChunk> {
    if chunks.len() <= 1 || overlap_lines == 0 {
        return chunks;
    }

    let source_lines: Vec<&str> = source.lines().collect();
    let mut result = Vec::with_capacity(chunks.len());

    for (i, mut chunk) in chunks.into_iter().enumerate() {
        if i > 0 && chunk.line_start > 0 {
            // Prepend `overlap_lines` lines from before this chunk's start
            let overlap_start = chunk.line_start.saturating_sub(overlap_lines);
            if overlap_start < chunk.line_start && overlap_start < source_lines.len() {
                let end = chunk.line_start.min(source_lines.len());
                let prefix: String = source_lines[overlap_start..end].join("\n");
                chunk.text = format!("{}\n{}", prefix, chunk.text);
                chunk.line_start = overlap_start;
                chunk.non_ws_chars = count_non_ws(&chunk.text);
            }
        }
        result.push(chunk);
    }

    result
}

/// Pre-sorted symbol index for O(log n) parent resolution via binary search.
struct SymbolIntervalIndex<'a> {
    /// Symbols sorted by (line_start ASC, line_end DESC) — outermost first at each start.
    sorted: Vec<&'a Symbol>,
}

impl<'a> SymbolIntervalIndex<'a> {
    fn build(symbols: &'a [Symbol]) -> Self {
        let mut sorted: Vec<&Symbol> = symbols.iter().collect();
        sorted.sort_by(|a, b| {
            a.line_start
                .cmp(&b.line_start)
                .then_with(|| b.line_end.cmp(&a.line_end))
        });
        Self { sorted }
    }

    /// Find the innermost symbol containing [line_start, line_end].
    /// Uses binary search to find candidates starting at or before line_start,
    /// then scans forward for the tightest containment.
    fn resolve(&self, line_start: usize, line_end: usize) -> Option<&'a Symbol> {
        if self.sorted.is_empty() {
            return None;
        }

        // Binary search: find the rightmost symbol whose line_start <= line_start
        let idx = match self
            .sorted
            .binary_search_by(|s| s.line_start.cmp(&line_start))
        {
            Ok(i) => i,
            Err(i) => {
                if i == 0 {
                    return None;
                }
                i - 1
            }
        };

        let mut best: Option<&Symbol> = None;
        let mut best_span = usize::MAX;

        // Scan backwards from idx (all candidates have line_start <= line_start)
        for &sym in self.sorted[..=idx].iter().rev() {
            if sym.line_start > line_start {
                continue;
            }
            // Once we pass symbols that start too early and are too short, stop
            if best.is_some() && sym.line_end < line_end {
                // Symbols are sorted with largest span first at each start position,
                // so once we see one that doesn't contain us and we already have
                // a best, earlier symbols with the same start won't either.
                // But symbols with smaller line_start may still contain us.
                continue;
            }
            if sym.line_end >= line_end {
                let span = sym.line_end - sym.line_start;
                if span < best_span {
                    best_span = span;
                    best = Some(sym);
                }
            }
        }

        best
    }
}

#[cfg(test)]
#[path = "tests/chunker_tests.rs"]
mod tests;