Speed up cold indexing by 5–11x with adaptive context

colgrep/src/index/mod.rs

@@ -25,7 +25,9 @@ use serde::{Deserialize, Serialize};
 use crate::acceleration::apply_acceleration_mode;
 use crate::acceleration::{env_acceleration_mode_lossy, AccelerationMode};
 use crate::embed::build_embedding_text;
-use crate::parser::{build_call_graph, detect_language, extract_units, CodeUnit, Language};
+use crate::parser::{
+    build_call_graph, detect_language, extract_units, CodeUnit, Language, UnitType,
+};
 use crate::signal::{is_interrupted, is_interrupted_outside_critical, CriticalSectionGuard};
 
 use paths::{
@@ -56,6 +58,90 @@ const BUILDING_MARKER: &str = ".building";
 /// more often (more resumable) at the cost of more index-append overhead.
 const BUILD_CHECKPOINT_UNITS: usize = 4096;
 
+/// Upper bound for adaptive indexing with the default code model. Most code units use much
+/// less based on their AST-derived shape; complex functions can use the full budget.
+const ADAPTIVE_INDEX_DOCUMENT_LENGTH: usize = 512;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum IndexDocumentLength {
+    Native,
+    Fixed(usize),
+    Adaptive,
+}
+
+fn build_checkpoint_units(binary: bool, index_exists: bool) -> usize {
+    if binary || !index_exists {
+        usize::MAX
+    } else {
+        BUILD_CHECKPOINT_UNITS
+    }
+}
+
+fn index_document_length(model_id: &str) -> IndexDocumentLength {
+    let override_value = std::env::var("COLGREP_INDEX_DOCUMENT_LENGTH").ok();
+    resolve_index_document_length(model_id, override_value.as_deref())
+}
+
+fn resolve_index_document_length(
+    model_id: &str,
+    override_value: Option<&str>,
+) -> IndexDocumentLength {
+    if let Some(value) = override_value {
+        if value == "0" {
+            return IndexDocumentLength::Native;
+        }
+        if let Ok(length) = value.parse::<usize>() {
+            return IndexDocumentLength::Fixed(length.max(2));
+        }
+    }
+    if model_id == crate::model::DEFAULT_MODEL {
+        IndexDocumentLength::Adaptive
+    } else {
+        IndexDocumentLength::Native
+    }
+}
+
+fn model_document_length(policy: IndexDocumentLength) -> Option<usize> {
+    match policy {
+        IndexDocumentLength::Native => None,
+        IndexDocumentLength::Fixed(length) => Some(length),
+        IndexDocumentLength::Adaptive => Some(ADAPTIVE_INDEX_DOCUMENT_LENGTH),
+    }
+}
+
+/// Choose how much source context is worth embedding for one AST-derived code unit.
+///
+/// Identity and signature metadata come first in `build_embedding_text`, so simple symbols
+/// need little context. Branching and high-complexity functions retain progressively more of
+/// their bodies. Non-AST documents use their line span as the only available structure signal.
+fn code_unit_document_length(unit: &CodeUnit, policy: IndexDocumentLength) -> usize {
+    match policy {
+        IndexDocumentLength::Native => usize::MAX,
+        IndexDocumentLength::Fixed(length) => length,
+        IndexDocumentLength::Adaptive => match unit.unit_type {
+            UnitType::Constant | UnitType::Class => 64,
+            UnitType::Function | UnitType::Method => {
+                if unit.complexity >= 10 {
+                    512
+                } else if unit.complexity >= 5 || unit.has_error_handling {
+                    256
+                } else if unit.complexity >= 2 || unit.has_loops || unit.has_branches {
+                    128
+                } else {
+                    64
+                }
+            }
+            UnitType::RawCode => match unit.end_line.saturating_sub(unit.line) {
+                0..=15 => 64,
+                16..=63 => 128,
+                _ => 256,
+            },
+            UnitType::Section => 128,
+            UnitType::Document => 256,
+        },
+    }
+}
+
 /// Test-only counter of expensive `delete_from_index` invocations.
 ///
 /// Issue #116: deleting many files used to call the full-index-rewrite primitive once per
@@ -403,13 +489,15 @@ pub struct UpdatePlan {
 struct SortedUnit {
     unit: Arc<CodeUnit>,
     text: Arc<str>,
+    max_length: usize,
 }
 
 /// After deduplication: unique texts to encode + a map from original positions
 /// back to their unique index, so duplicates share a single GPU encoding pass.
 struct PreparedChunk {
     units: Vec<Arc<CodeUnit>>,
     unique_texts: Vec<Arc<str>>,
+    unique_max_lengths: Vec<usize>,
     original_to_unique: Vec<usize>,
 }
 
@@ -457,12 +545,17 @@ struct ChunkForCoding {
 /// When encoding more than this many units, prompt the user unless auto_confirm is set.
 pub const CONFIRMATION_THRESHOLD: usize = 30_000;
 
-fn prepare_units_for_encoding(units: &[CodeUnit], sample_prefix_size: usize) -> Vec<SortedUnit> {
+fn prepare_units_for_encoding(
+    units: &[CodeUnit],
+    sample_prefix_size: usize,
+    document_length: IndexDocumentLength,
+) -> Vec<SortedUnit> {
     let mut items: Vec<SortedUnit> = units
         .iter()
         .map(|unit| SortedUnit {
             unit: Arc::new(unit.clone()),
             text: Arc::<str>::from(build_embedding_text(unit)),
+            max_length: code_unit_document_length(unit, document_length),
         })
         .collect();
 
@@ -500,17 +593,20 @@ fn prepare_units_for_encoding(units: &[CodeUnit], sample_prefix_size: usize) ->
 /// unit can retrieve its embedding after the GPU pass. On large codebases
 /// this saves ~8% of encoding work (e.g. re-exported types, trait impls).
 fn prepare_deduplicated_chunk(unit_chunk: &[SortedUnit]) -> PreparedChunk {
-    let mut index_by_text: HashMap<&str, usize> = HashMap::new();
+    let mut index_by_text: HashMap<(&str, usize), usize> = HashMap::new();
     let mut unique_texts: Vec<Arc<str>> = Vec::new();
+    let mut unique_max_lengths: Vec<usize> = Vec::new();
     let mut original_to_unique: Vec<usize> = Vec::with_capacity(unit_chunk.len());
 
     for item in unit_chunk.iter() {
-        if let Some(&unique_idx) = index_by_text.get(item.text.as_ref()) {
+        let key = (item.text.as_ref(), item.max_length);
+        if let Some(&unique_idx) = index_by_text.get(&key) {
             original_to_unique.push(unique_idx);
         } else {
             let unique_idx = unique_texts.len();
-            index_by_text.insert(item.text.as_ref(), unique_idx);
+            index_by_text.insert(key, unique_idx);
             unique_texts.push(Arc::clone(&item.text));
+            unique_max_lengths.push(item.max_length);
             original_to_unique.push(unique_idx);
         }
     }
@@ -521,6 +617,7 @@ fn prepare_deduplicated_chunk(unit_chunk: &[SortedUnit]) -> PreparedChunk {
             .map(|item| Arc::clone(&item.unit))
             .collect(),
         unique_texts,
+        unique_max_lengths,
         original_to_unique,
     }
 }
@@ -574,7 +671,8 @@ fn run_tokenize_stage(
             .iter()
             .map(|text| text.as_ref())
             .collect();
-        let prepared_batches = model.tokenize_documents_in_batches(&text_refs)?;
+        let prepared_batches = model
+            .tokenize_documents_in_batches_with_limits(&text_refs, &chunk.unique_max_lengths)?;
 
         sender
             .send(TokenizedChunk {
@@ -1327,20 +1425,24 @@ impl IndexBuilder {
             let batch = self
                 .batch_size
                 .unwrap_or_else(crate::config::get_default_batch_size);
+            let document_length = index_document_length(&self.model_id);
 
             // Suppress stderr during model loading to hide CoreML's harmless
             // "Context leak detected" warnings on macOS.
             // `with_suppressed_stderr` captures any panic message via a temporary
             // panic hook and prints it to the restored stderr before resuming,
             // so panics inside the suppressed region remain visible.
             let model = crate::stderr::with_suppressed_stderr(|| {
-                Colbert::builder(&self.model_path)
+                let mut builder = Colbert::builder(&self.model_path)
                     .with_quantized(self.quantized)
                     .with_parallel(num_sessions)
                     .with_batch_size(batch)
                     .with_dynamic_batch(self.dynamic_batch)
-                    .with_execution_provider(execution_provider)
-                    .build()
+                    .with_execution_provider(execution_provider);
+                if let Some(length) = model_document_length(document_length) {
+                    builder = builder.with_document_length(length);
+                }
+                builder.build()
             })
             .context("Failed to load ColBERT model")?;
 
@@ -1381,15 +1483,19 @@ impl IndexBuilder {
             .parallel_sessions
             .unwrap_or_else(|| default_index_parallel_sessions(ExecutionProvider::Cpu, usize::MAX));
         let batch = crate::config::DEFAULT_BATCH_SIZE_CPU;
+        let document_length = index_document_length(&self.model_id);
 
         let model = crate::stderr::with_suppressed_stderr(|| {
-            Colbert::builder(&self.model_path)
+            let mut builder = Colbert::builder(&self.model_path)
                 .with_quantized(self.quantized)
                 .with_parallel(num_sessions)
                 .with_batch_size(batch)
                 .with_dynamic_batch(false)
-                .with_execution_provider(ExecutionProvider::Cpu)
-                .build()
+                .with_execution_provider(ExecutionProvider::Cpu);
+            if let Some(length) = model_document_length(document_length) {
+                builder = builder.with_document_length(length);
+            }
+            builder.build()
         })
         .context("Failed to load ColBERT model for CPU fallback")?;
 
@@ -2004,7 +2110,11 @@ impl IndexBuilder {
         // the encoding pipeline, which rebuilds their FTS5 entries.
         delete_files_from_index_no_fts_rebuild(index_path, &files_changed)?;
 
-        let sorted_units = prepare_units_for_encoding(&new_units, index_chunk_size);
+        let sorted_units = prepare_units_for_encoding(
+            &new_units,
+            index_chunk_size,
+            index_document_length(&self.model_id),
+        );
         let was_interrupted = self.run_encoding_pipeline(
             &sorted_units,
             index_chunk_size,
@@ -2287,17 +2397,14 @@ impl IndexBuilder {
         // Encode in file-coherent batches of ~BUILD_CHECKPOINT_UNITS units, committing state
         // after each batch so interruptions keep finished work.
         //
-        // Binary builds use a single batch instead: the first committed batch
-        // creates the index and every later batch appends via update_append,
-        // which binary indexes reject. One batch keeps the whole build on the
-        // atomic initial-create path (encoded sign bits are ~4x smaller than
-        // residual codes, so holding them in memory is cheap); the cost is
-        // that an interrupted binary build restarts from zero.
-        let checkpoint_units = if self.binary {
-            usize::MAX
-        } else {
-            BUILD_CHECKPOINT_UNITS
-        };
+        // A fresh build uses a single batch so the initial-create path writes the index once.
+        // Repeated update_append calls rebuild index-wide IVF data, making a new index
+        // quadratic in its number of checkpoints. A resumed build keeps the smaller
+        // checkpoints so an interruption only loses its current batch.
+        //
+        // Binary builds also require one batch because they cannot be appended to.
+        let checkpoint_units =
+            build_checkpoint_units(self.binary, index_dir.join("metadata.json").exists());
         let encode_pb = ProgressBar::new(total_units as u64);
         encode_pb.set_style(
             ProgressStyle::default_bar()
@@ -2410,7 +2517,11 @@ impl IndexBuilder {
         delete_files_from_index_no_fts_rebuild(index_path, batch_files)?;
         self.ensure_model_created(batch_units.len())?;
         let pool_factor = self.resolve_pool_factor(batch_units.len());
-        let sorted_units = prepare_units_for_encoding(batch_units, index_chunk_size);
+        let sorted_units = prepare_units_for_encoding(
+            batch_units,
+            index_chunk_size,
+            index_document_length(&self.model_id),
+        );
         let was_interrupted = self.run_encoding_pipeline(
             &sorted_units,
             index_chunk_size,
@@ -2774,7 +2885,11 @@ impl IndexBuilder {
             // new versions; we do a single rebuild after encoding completes.
             delete_files_from_index_no_fts_rebuild(index_path, &plan.changed)?;
 
-            let sorted_units = prepare_units_for_encoding(&new_units, index_chunk_size);
+            let sorted_units = prepare_units_for_encoding(
+                &new_units,
+                index_chunk_size,
+                index_document_length(&self.model_id),
+            );
             let pipeline_interrupted = self.run_encoding_pipeline(
                 &sorted_units,
                 index_chunk_size,
@@ -3213,7 +3328,11 @@ impl IndexBuilder {
         // Compute effective pool factor based on batch size
         let pool_factor = self.resolve_pool_factor(units.len());
 
-        let sorted_units = prepare_units_for_encoding(units, index_chunk_size);
+        let sorted_units = prepare_units_for_encoding(
+            units,
+            index_chunk_size,
+            index_document_length(&self.model_id),
+        );
         self.ensure_model_created(units.len())?;
         let was_interrupted = self.run_encoding_pipeline(
             &sorted_units,
@@ -4597,6 +4716,75 @@ fn prompt_large_index_confirmation(num_units: usize) -> bool {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_fresh_build_uses_single_write() {
+        assert_eq!(build_checkpoint_units(false, false), usize::MAX);
+    }
+
+    #[test]
+    fn test_resumed_build_keeps_checkpoints() {
+        assert_eq!(build_checkpoint_units(false, true), BUILD_CHECKPOINT_UNITS);
+    }
+
+    #[test]
+    fn test_binary_build_uses_single_write() {
+        assert_eq!(build_checkpoint_units(true, false), usize::MAX);
+        assert_eq!(build_checkpoint_units(true, true), usize::MAX);
+    }
+
+    #[test]
+    fn test_default_model_uses_adaptive_index_document_length() {
+        assert_eq!(
+            resolve_index_document_length(crate::model::DEFAULT_MODEL, None),
+            IndexDocumentLength::Adaptive
+        );
+        assert_eq!(
+            resolve_index_document_length("other/model", None),
+            IndexDocumentLength::Native
+        );
+        assert_eq!(
+            resolve_index_document_length(crate::model::DEFAULT_MODEL, Some("256")),
+            IndexDocumentLength::Fixed(256)
+        );
+        assert_eq!(
+            resolve_index_document_length(crate::model::DEFAULT_MODEL, Some("0")),
+            IndexDocumentLength::Native
+        );
+        assert_eq!(
+            resolve_index_document_length(crate::model::DEFAULT_MODEL, Some("invalid")),
+            IndexDocumentLength::Adaptive
+        );
+    }
+
+    #[test]
+    fn test_adaptive_document_length_uses_structure() {
+        let mut unit = CodeUnit::new(
+            "example".to_string(),
+            PathBuf::from("src/example.rs"),
+            1,
+            3,
+            Language::Rust,
+            UnitType::Function,
+            None,
+        );
+        assert_eq!(
+            code_unit_document_length(&unit, IndexDocumentLength::Adaptive),
+            64
+        );
+
+        unit.has_branches = true;
+        assert_eq!(
+            code_unit_document_length(&unit, IndexDocumentLength::Adaptive),
+            128
+        );
+
+        unit.complexity = 10;
+        assert_eq!(
+            code_unit_document_length(&unit, IndexDocumentLength::Adaptive),
+            512
+        );
+    }
+
     /// The mtime fast path in `compute_update_plan` must skip content hashing
     /// for files whose stored mtime is unchanged. The stored hash is wrong on
     /// purpose: if the plan still reports the file as unchanged, hashing was

colgrep/src/index/state.rs

@@ -12,7 +12,7 @@ use super::paths::get_state_path;
 /// metadata schema). Bump ONLY for incompatible changes: a mismatch discards
 /// the index and re-embeds the entire project on the next run. Routine CLI
 /// releases must NOT bump this.
-pub const INDEX_FORMAT_VERSION: u32 = 2;
+pub const INDEX_FORMAT_VERSION: u32 = 3;
 
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct IndexState {