fix(diffctx): reduce peak RSS via chunked fragmentation, streaming cache, mimalloc

Author: nikolay-e

diffctx/Cargo.lock

@@ -63,6 +63,7 @@ lang-extra = [
 ]
 
 [dependencies]
+mimalloc = { version = "0.1", default-features = false }
 clap = { version = "4", features = ["derive"] }
 tree-sitter = "0.25"
 tree-sitter-python = { version = "0.23", optional = true }

diffctx/Cargo.toml

@@ -253,37 +253,45 @@ pub fn process_files_for_fragments(
     let max_frags = LIMITS.max_fragments;
     let max_generated = LIMITS.max_generated_fragments;
 
-    let file_contents: Vec<(PathBuf, String)> = files
-        .iter()
-        .filter_map(|file_path| {
-            let content = read_file_content(
-                file_path,
-                root_dir,
-                preferred_revs,
-                batch_reader.as_deref_mut(),
-            )?;
-            Some((file_path.clone(), content))
-        })
-        .collect();
-
-    let parsed: Vec<Vec<Fragment>> = file_contents
-        .par_iter()
-        .map(|(file_path, content)| {
-            let path_arc: Arc<str> = Arc::from(file_path.to_string_lossy().as_ref());
-            let mut raw_frags = fragment_file(path_arc, content);
-
-            let generated = is_generated_file(file_path, content);
-            let cap = if generated { max_generated } else { max_frags };
-            if raw_frags.len() > cap {
-                raw_frags.sort_by(|a, b| b.line_count().cmp(&a.line_count()));
-                raw_frags.truncate(cap);
-            }
-            if generated {
-                raw_frags = truncate_generated_fragments(raw_frags);
-            }
-            raw_frags
-        })
-        .collect();
+    // Process files in chunks: sequential read (CatFileBatch is &mut, !Send) then
+    // parallel parse within each chunk. Peak raw-content memory = chunk_size × max_file_size
+    // instead of N_files × avg_file_size — eliminates OOM on large repos like astropy.
+    let chunk_size = rayon::current_num_threads().max(1);
+    let mut parsed: Vec<Vec<Fragment>> = Vec::with_capacity(files.len());
+    for chunk in files.chunks(chunk_size) {
+        let chunk_contents: Vec<(PathBuf, String)> = chunk
+            .iter()
+            .filter_map(|file_path| {
+                let content = read_file_content(
+                    file_path,
+                    root_dir,
+                    preferred_revs,
+                    batch_reader.as_deref_mut(),
+                )?;
+                Some((file_path.clone(), content))
+            })
+            .collect();
+        parsed.extend(
+            chunk_contents
+                .par_iter()
+                .map(|(file_path, content)| {
+                    let path_arc: Arc<str> = Arc::from(file_path.to_string_lossy().as_ref());
+                    let mut raw_frags = fragment_file(path_arc, content);
+                    let generated = is_generated_file(file_path, content);
+                    let cap = if generated { max_generated } else { max_frags };
+                    if raw_frags.len() > cap {
+                        raw_frags.sort_by(|a, b| b.line_count().cmp(&a.line_count()));
+                        raw_frags.truncate(cap);
+                    }
+                    if generated {
+                        raw_frags = truncate_generated_fragments(raw_frags);
+                    }
+                    raw_frags
+                })
+                .collect::<Vec<_>>(),
+        );
+        // chunk_contents dropped here — raw file text freed before next chunk
+    }
 
     let mut fragments: Vec<Fragment> = Vec::new();
     for file_frags in parsed {

diffctx/src/fragmentation.rs

@@ -1,3 +1,7 @@
+use mimalloc::MiMalloc;
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
 use std::path::PathBuf;
 
 use anyhow::Result;

diffctx/src/main.rs

@@ -552,27 +552,26 @@ fn create_discovery(config: &PipelineConfig) -> Box<dyn DiscoveryStrategy> {
 }
 
 fn build_file_cache(candidate_files: &[PathBuf]) -> FxHashMap<PathBuf, String> {
-    let mut entries: Vec<(PathBuf, String)> = candidate_files
-        .par_iter()
-        .filter_map(|f| {
-            let meta = f.metadata().ok()?;
-            if meta.len() as usize > LIMITS.max_file_size {
-                return None;
-            }
-            let content = std::fs::read_to_string(f).ok()?;
-            Some((f.clone(), content))
-        })
-        .collect();
-    entries.sort_by(|a, b| a.0.cmp(&b.0));
-
+    // Stream files one at a time to avoid materialising all content before the cap.
+    // Previous par_iter().collect() allocated the full eligible corpus into an
+    // intermediate Vec before truncating — on repos with thousands of files this
+    // caused peak memory far above max_cache_bytes.
+    let mut sorted = candidate_files.to_vec();
+    sorted.sort();
     let mut cache: FxHashMap<PathBuf, String> = FxHashMap::default();
     let mut cache_bytes = 0usize;
-    for (path, content) in entries {
+    for path in sorted {
         if cache_bytes > GRAPH_FILTERING.max_cache_bytes {
             break;
         }
-        cache_bytes += content.len();
-        cache.insert(path, content);
+        let Ok(meta) = path.metadata() else { continue };
+        if meta.len() as usize > LIMITS.max_file_size {
+            continue;
+        }
+        if let Ok(content) = std::fs::read_to_string(&path) {
+            cache_bytes += content.len();
+            cache.insert(path, content);
+        }
     }
     cache
 }