perf(index): optimize initial indexing from 28.5s to 13s (2.2x faster)

Four targeted optimizations addressing measured bottlenecks in the
extraction pipeline for large Swift projects (~2000 files):

- Thread-local tree-sitter parsers: eliminate per-file Parser::new() +
  set_language() allocation in both Swift and Rust extractors
- Pre-canonicalize module directories: ModuleMap.with_fallback() now
  canonicalizes dirs once at construction, removing O(M×N) redundant
  syscalls from the parallel extraction loop
- Outward-expanding snippet search with radius limit: replace O(N×L)
  full-source scan per node with outward search from expected line,
  capped at 200 lines. The worst straggler (Strings.generated.swift,
  12K nodes, 28K lines) dropped from 18.5s to 0.4s
- Eager index store warmup: buildFileIndex() now runs during
  prepare_project instead of lazily blocking all rayon threads on
  first extraction

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ahavili

grapha-core/src/module.rs

@@ -26,18 +26,27 @@ impl ModuleMap {
                 .to_string();
             self.modules.insert(name, vec![root.to_path_buf()]);
         }
+        self.canonicalize_dirs();
         self
     }
 
+    fn canonicalize_dirs(&mut self) {
+        for dirs in self.modules.values_mut() {
+            for dir in dirs.iter_mut() {
+                if let Ok(canonical) = dir.canonicalize() {
+                    *dir = canonical;
+                }
+            }
+        }
+    }
+
     pub fn module_for_file(&self, file: &Path) -> Option<String> {
         let canonical_file = normalize_path(file);
         let mut best_match: Option<(&str, usize)> = None;
 
         for (name, dirs) in &self.modules {
             for dir in dirs {
-                let canonical_dir = normalize_path(dir);
-
-                if let Ok(suffix) = canonical_file.strip_prefix(&canonical_dir) {
+                if let Ok(suffix) = canonical_file.strip_prefix(dir) {
                     let depth = suffix.components().count();
                     match best_match {
                         Some((_, best_depth)) if depth < best_depth => {
@@ -52,7 +61,7 @@ impl ModuleMap {
 
                 if best_match.is_none()
                     && file.is_relative()
-                    && let Some(dir_name) = canonical_dir.file_name().and_then(|name| name.to_str())
+                    && let Some(dir_name) = dir.file_name().and_then(|name| name.to_str())
                 {
                     let file_str = file.to_string_lossy();
                     if file_str.starts_with(dir_name)

grapha-core/src/pipeline.rs

@@ -213,6 +213,7 @@ mod tests {
         };
         let mut modules = ModuleMap::new();
         modules.modules.insert("core".to_string(), vec![src_dir]);
+        modules = modules.with_fallback(dir.path());
         let file_context = file_context(&project, &modules, &file);
 
         let result = extract_with_registry(&registry, b"fn main() {}", &file_context).unwrap();

grapha-swift/src/bridge.rs

@@ -32,6 +32,7 @@ type IndexStoreOpenFn = unsafe extern "C" fn(*const i8, *mut i32) -> *mut std::f
 type IndexStoreCloseFn = unsafe extern "C" fn(*mut std::ffi::c_void);
 type IndexStoreExtractFn =
     unsafe extern "C" fn(*mut std::ffi::c_void, *const i8, *mut u32, *mut i32) -> *const u8;
+type IndexStoreWarmupFn = unsafe extern "C" fn(*mut std::ffi::c_void);
 type SwiftSyntaxExtractFn = unsafe extern "C" fn(*const i8, usize, *const i8) -> *const i8;
 type FreeStringFn = unsafe extern "C" fn(*mut i8);
 type FreeBufferFn = unsafe extern "C" fn(*mut u8);
@@ -41,6 +42,7 @@ pub struct SwiftBridge {
     pub indexstore_open: IndexStoreOpenFn,
     pub indexstore_close: IndexStoreCloseFn,
     pub indexstore_extract: IndexStoreExtractFn,
+    pub indexstore_warmup: Option<IndexStoreWarmupFn>,
     pub swiftsyntax_extract: SwiftSyntaxExtractFn,
     pub free_string: FreeStringFn,
     pub free_buffer: FreeBufferFn,
@@ -69,6 +71,10 @@ impl SwiftBridge {
             let indexstore_extract = *lib
                 .get::<IndexStoreExtractFn>(b"grapha_indexstore_extract")
                 .ok()?;
+            let indexstore_warmup = lib
+                .get::<IndexStoreWarmupFn>(b"grapha_indexstore_warmup")
+                .ok()
+                .map(|s| *s);
             let swiftsyntax_extract = *lib
                 .get::<SwiftSyntaxExtractFn>(b"grapha_swiftsyntax_extract")
                 .ok()?;
@@ -80,6 +86,7 @@ impl SwiftBridge {
                 indexstore_open,
                 indexstore_close,
                 indexstore_extract,
+                indexstore_warmup,
                 swiftsyntax_extract,
                 free_string,
                 free_buffer,

grapha-swift/src/indexstore.rs

@@ -138,6 +138,19 @@ fn get_or_open_store(index_store_path: &Path) -> Option<Arc<StoreHandle>> {
     })
 }
 
+pub fn warmup_indexstore(index_store_path: &Path) {
+    let Some(bridge) = bridge::bridge() else {
+        return;
+    };
+    let Some(warmup_fn) = bridge.indexstore_warmup else {
+        return;
+    };
+    let Some(handle) = get_or_open_store(index_store_path) else {
+        return;
+    };
+    unsafe { warmup_fn(handle.ptr) };
+}
+
 pub fn extract_from_indexstore(
     file_path: &Path,
     index_store_path: &Path,

grapha-swift/src/lib.rs

@@ -165,6 +165,9 @@ pub fn init_index_store(project_root: &Path) {
 
 fn prepare_project_index_store(project_root: &Path) {
     prepare_project_with(&INDEX_STORE_PATHS, project_root, discover_index_store);
+    if let Some(store_path) = index_store_path(project_root) {
+        indexstore::warmup_indexstore(&store_path);
+    }
 }
 
 fn prepare_project_with<F>(

grapha-swift/src/treesitter/extract.rs

@@ -1,3 +1,4 @@
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::path::Path;
 
@@ -9,24 +10,28 @@ use grapha_core::{ExtractionResult, LanguageExtractor};
 use super::common::*;
 use super::swiftui::extract_swiftui_declaration_structure;
 
+thread_local! {
+    static SWIFT_PARSER: RefCell<Parser> = RefCell::new({
+        let mut p = Parser::new();
+        p.set_language(&tree_sitter_swift::LANGUAGE.into()).expect("failed to load Swift grammar");
+        p
+    });
+}
+
 /// Parse Swift source once. Reuse the tree across enrichment passes.
 pub fn parse_swift(source: &[u8]) -> anyhow::Result<Tree> {
-    let mut parser = Parser::new();
-    parser.set_language(&tree_sitter_swift::LANGUAGE.into())?;
-    parser
-        .parse(source, None)
-        .ok_or_else(|| anyhow::anyhow!("tree-sitter failed to parse Swift source"))
+    SWIFT_PARSER.with_borrow_mut(|parser| {
+        parser
+            .parse(source, None)
+            .ok_or_else(|| anyhow::anyhow!("tree-sitter failed to parse Swift source"))
+    })
 }
 
 pub struct SwiftExtractor;
 
 impl LanguageExtractor for SwiftExtractor {
     fn extract(&self, source: &[u8], file_path: &Path) -> anyhow::Result<ExtractionResult> {
-        let mut parser = Parser::new();
-        parser.set_language(&tree_sitter_swift::LANGUAGE.into())?;
-        let tree = parser
-            .parse(source, None)
-            .ok_or_else(|| anyhow::anyhow!("tree-sitter failed to parse Swift source"))?;
+        let tree = parse_swift(source)?;
 
         let mut result = ExtractionResult::new();
         let file_str = file_path.to_string_lossy().to_string();

grapha-swift/swift-bridge/Sources/GraphaSwiftBridge/Bridge.swift

@@ -72,6 +72,14 @@ public func indexstoreExtract(
     return UnsafeRawPointer(ptr)
 }
 
+@c(grapha_indexstore_warmup)
+public func indexstoreWarmup(_ handle: UnsafeMutableRawPointer?) {
+    guard let handle else { return }
+    let key = Int(bitPattern: handle)
+    let reader = _readers.withLock { $0[key] }
+    reader?.warmup()
+}
+
 // MARK: - SwiftSyntax
 
 @c(grapha_swiftsyntax_extract)

grapha-swift/swift-bridge/Sources/GraphaSwiftBridge/IndexStoreReader.swift

@@ -171,22 +171,12 @@ final class IndexStoreReader: @unchecked Sendable {
 
     // MARK: - Public
 
+    func warmup() {
+        ensureFileIndex()
+    }
+
     func extractFile(_ filePath: String) -> (UnsafeMutableRawPointer, UInt32)? {
-        if fileIndex == nil {
-            if let cached = _sharedFileIndexCache.withLock({ $0[storePath] }) {
-                fileIndex = cached
-            } else {
-                _buildIndexLock.withLock { _ in
-                    if let cached = _sharedFileIndexCache.withLock({ $0[storePath] }) {
-                        fileIndex = cached
-                    } else {
-                        fileIndex = buildFileIndex()
-                        let idx = fileIndex!
-                        _sharedFileIndexCache.withLock { $0[storePath] = idx }
-                    }
-                }
-            }
-        }
+        ensureFileIndex()
 
         let resolved = resolvePath(filePath)
         // Fast last-path-component extraction without Foundation URL
@@ -218,6 +208,24 @@ final class IndexStoreReader: @unchecked Sendable {
 
     // MARK: - File Index (built once)
 
+    private func ensureFileIndex() {
+        if fileIndex == nil {
+            if let cached = _sharedFileIndexCache.withLock({ $0[storePath] }) {
+                fileIndex = cached
+            } else {
+                _buildIndexLock.withLock { _ in
+                    if let cached = _sharedFileIndexCache.withLock({ $0[storePath] }) {
+                        fileIndex = cached
+                    } else {
+                        fileIndex = buildFileIndex()
+                        let idx = fileIndex!
+                        _sharedFileIndexCache.withLock { $0[storePath] = idx }
+                    }
+                }
+            }
+        }
+    }
+
     private func buildFileIndex() -> [String: UnitInfo] {
         _cbStore = store
         _cbFileIndex = [:]

grapha/src/app/pipeline.rs

@@ -149,12 +149,18 @@ pub(crate) fn run_pipeline(
     let t_read_ns = AtomicU64::new(0);
     let t_extract_ns = AtomicU64::new(0);
     let t_snippet_ns = AtomicU64::new(0);
+    let t_file_context_ns = AtomicU64::new(0);
+    let t_total_per_file_ns = AtomicU64::new(0);
+    let t_max_single_file_ns = AtomicU64::new(0);
     let extraction_cache_entries = Mutex::new(std::collections::HashMap::new());
 
     let results: Vec<_> = all_files
         .par_iter()
         .filter_map(|file| {
+            let t_file_start = Instant::now();
+            let t_fc = Instant::now();
             let file_context = grapha_core::file_context(&project_context, &module_map, file);
+            t_file_context_ns.fetch_add(t_fc.elapsed().as_nanos() as u64, Ordering::Relaxed);
             let cache_key = extraction_cache_key(&file_context.relative_path);
             if let Some(existing_cache) = existing_extraction_cache
                 && let Some(entry) = existing_cache.get(&cache_key)
@@ -169,6 +175,9 @@ pub(crate) fn run_pipeline(
                     .lock()
                     .expect("extraction cache mutex poisoned")
                     .insert(cache_key, entry.clone());
+                let file_ns = t_file_start.elapsed().as_nanos() as u64;
+                t_total_per_file_ns.fetch_add(file_ns, Ordering::Relaxed);
+                t_max_single_file_ns.fetch_max(file_ns, Ordering::Relaxed);
                 return Some(entry.result.clone());
             }
 
@@ -180,6 +189,9 @@ pub(crate) fn run_pipeline(
                     if let Some(ref pb) = pb {
                         pb.inc(1);
                     }
+                    let file_ns = t_file_start.elapsed().as_nanos() as u64;
+                    t_total_per_file_ns.fetch_add(file_ns, Ordering::Relaxed);
+                    t_max_single_file_ns.fetch_max(file_ns, Ordering::Relaxed);
                     return None;
                 }
             };
@@ -225,6 +237,9 @@ pub(crate) fn run_pipeline(
                             .expect("extraction cache mutex poisoned")
                             .insert(key, entry);
                     }
+                    let file_ns = t_file_start.elapsed().as_nanos() as u64;
+                    t_total_per_file_ns.fetch_add(file_ns, Ordering::Relaxed);
+                    t_max_single_file_ns.fetch_max(file_ns, Ordering::Relaxed);
                     Some(result)
                 }
                 Err(e) => {
@@ -234,6 +249,9 @@ pub(crate) fn run_pipeline(
                             eprintln!("  \x1b[33m!\x1b[0m skipping {}: {e}", file.display())
                         });
                     }
+                    let file_ns = t_file_start.elapsed().as_nanos() as u64;
+                    t_total_per_file_ns.fetch_add(file_ns, Ordering::Relaxed);
+                    t_max_single_file_ns.fetch_max(file_ns, Ordering::Relaxed);
                     None
                 }
             }
@@ -279,10 +297,14 @@ pub(crate) fn run_pipeline(
         let ts_fb_ms = grapha_swift::TIMING_TS_FALLBACK_NS
             .load(std::sync::atomic::Ordering::Relaxed) as f64
             / 1_000_000.0;
+        let fc_ms = t_file_context_ns.load(Ordering::Relaxed) as f64 / 1_000_000.0;
+        let total_per_file_ms = t_total_per_file_ns.load(Ordering::Relaxed) as f64 / 1_000_000.0;
+        let max_single_file_ms = t_max_single_file_ns.load(Ordering::Relaxed) as f64 / 1_000_000.0;
         eprintln!(
-            "    thread-summed: read {:.0}ms, extract {:.0}ms, snippet {:.0}ms",
-            read_ms, extract_ms, snippet_ms
+            "    thread-summed: read {:.0}ms, extract {:.0}ms, snippet {:.0}ms, file_context {:.0}ms, total_per_file {:.0}ms",
+            read_ms, extract_ms, snippet_ms, fc_ms, total_per_file_ms
         );
+        eprintln!("    max_single_file: {:.0}ms", max_single_file_ms);
         eprintln!(
             "    swift: indexstore {:.0}ms, ts-parse {:.0}ms, doc {:.0}ms, swiftui {:.0}ms, l10n {:.0}ms, asset {:.0}ms, swiftsyntax {:.0}ms, ts-fallback {:.0}ms",
             is_ms, ts_parse_ms, doc_ms, swiftui_ms, l10n_ms, asset_ms, ss_ms, ts_fb_ms

grapha/src/extract/rust.rs

@@ -1,3 +1,4 @@
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::path::Path;
 
@@ -9,15 +10,23 @@ use grapha_core::graph::{
 
 use super::{ExtractionResult, LanguageExtractor};
 
+thread_local! {
+    static RUST_PARSER: RefCell<Parser> = RefCell::new({
+        let mut p = Parser::new();
+        p.set_language(&tree_sitter_rust::LANGUAGE.into()).expect("failed to load Rust grammar");
+        p
+    });
+}
+
 pub struct RustExtractor;
 
 impl LanguageExtractor for RustExtractor {
     fn extract(&self, source: &[u8], file_path: &Path) -> anyhow::Result<ExtractionResult> {
-        let mut parser = Parser::new();
-        parser.set_language(&tree_sitter_rust::LANGUAGE.into())?;
-        let tree = parser
-            .parse(source, None)
-            .ok_or_else(|| anyhow::anyhow!("tree-sitter failed to parse source"))?;
+        let tree = RUST_PARSER.with_borrow_mut(|parser| {
+            parser
+                .parse(source, None)
+                .ok_or_else(|| anyhow::anyhow!("tree-sitter failed to parse source"))
+        })?;
 
         let mut result = ExtractionResult::new();
         let file_str = file_path.to_string_lossy().to_string();