perf: parallel cold checks that actually scale + flat LSP latency

baml_language/crates/baml_cli/Cargo.toml

@@ -50,6 +50,7 @@ indicatif = { workspace = true }
 libsui = { workspace = true }
 mimalloc = { workspace = true }
 object = { workspace = true }
+rayon = { workspace = true }
 regex = { workspace = true }
 reqwest = { workspace = true, features = [ "blocking", "json", "rustls" ] }
 serde = { workspace = true }

baml_language/crates/baml_cli/src/project_load.rs

@@ -280,8 +280,12 @@ pub(crate) fn resolve_project_sources(from: Option<&Path>) -> Result<ResolvedPro
     };
 
     let walk_root = project_source_root(&canonical);
+    // Read sources across worker threads; `collect` on an indexed parallel
+    // iterator preserves discovery order, so `FileId` assignment (and with it
+    // diagnostic ordering) is identical to a serial read.
+    use rayon::prelude::*;
     let files = discover_baml_files(&walk_root)
-        .into_iter()
+        .into_par_iter()
         .map(|path| {
             let content = std::fs::read_to_string(&path)
                 .with_context(|| format!("Failed to read {}", path.display()))?;
@@ -303,10 +307,17 @@ pub(crate) fn build_db_from_sources(
 ) -> ProjectDatabase {
     let mut db = ProjectDatabase::new();
     db.set_project_root(&resolved.root);
-    for (path, content) in &resolved.files {
+    for (path, _) in &resolved.files {
         on_file(path);
-        db.add_or_update_file(path, content);
     }
+    // Bulk registration: one project-file-list write instead of one per file
+    // (the per-file path is O(files²) Vec copies + one salsa revision each).
+    db.add_or_update_files(
+        resolved
+            .files
+            .iter()
+            .map(|(path, content)| (path.as_path(), content.as_str())),
+    );
     db
 }
 

baml_language/crates/baml_project/src/check.rs

@@ -79,59 +79,90 @@ pub fn collect_compiler2_diagnostics(db: &ProjectDatabase) -> Vec<Diagnostic> {
 
 /// Run `check_file` for every file across worker threads.
 ///
-/// Every query `check_file` reaches is read-only, and `ProjectDatabase`'s
-/// `Clone` produces a shared-storage salsa handle (the rust-analyzer
-/// concurrency model): workers share one memo table, so a scope inferred by
-/// one thread is a cache hit for every other. Output order does not matter —
-/// the caller sorts diagnostics by (file, span, message) — so workers just
-/// pull the next file off a shared atomic counter (files vary a lot in size;
-/// work-stealing beats fixed chunks).
-///
-/// The first file is checked serially before fanning out: it warms the
-/// package-level queries (package items / resolution context / alias maps)
-/// that every file reads, so cold workers don't all block on the same shared
-/// memo slots.
+/// Output order does not matter — the caller sorts diagnostics by
+/// (file, span, message).
 fn collect_file_diagnostics_parallel(
     db: &ProjectDatabase,
     source_files: &[SourceFile],
     diagnostics: &mut Vec<Diagnostic>,
 ) {
-    // `ProjectDatabase` is `Send` but deliberately not `Sync` (each salsa
-    // handle carries thread-confined query-stack state), so tasks cannot share
-    // `&db` — every task MOVES its own cloned shared-storage handle instead
-    // (an Arc bump; all clones share one memo table). Small chunks keep
-    // work-stealing effective on rayon's global pool — files vary a lot in
-    // check cost — while amortizing the per-task clone.
+    for file_diagnostics in check_files_parallel(db, source_files) {
+        diagnostics.extend(file_diagnostics);
+    }
+}
+
+/// Prime every compiler2 file's PPIR semantic index across worker threads.
+///
+/// Whole-package aggregate queries (`package_items` / `namespace_items`)
+/// fold over **every** file's semantic index, and every file's check demands
+/// them. On a cold database, the first worker to claim such an aggregate
+/// memo computes parse + lowering + indexing for the entire project inline
+/// on its own thread while every other worker parks on that memo's sync
+/// slot — the whole compile degenerates to ~1 effective core regardless of
+/// thread count. Priming the per-file indexes first keeps all workers busy
+/// on file-local work; the aggregate fold itself is then cheap for whichever
+/// worker claims it.
+fn prime_file_indexes_parallel(db: &ProjectDatabase) {
+    const CHUNK: usize = 4;
+    let all_files = baml_compiler2_hir::compiler2_all_files(db);
+    let chunks: Vec<&[SourceFile]> = all_files.chunks(CHUNK).collect();
+    let handles: Vec<ProjectDatabase> = chunks.iter().map(|_| db.clone()).collect();
+    rayon::scope(move |s| {
+        for (chunk, db) in chunks.into_iter().zip(handles) {
+            s.spawn(move |_| {
+                for file in chunk {
+                    let _ = baml_compiler2_ppir::file_semantic_index(&db, *file);
+                }
+            });
+        }
+    });
+}
+
+/// Check `files` across worker threads, returning each file's diagnostics in
+/// input order (so callers that group per file — the LSP's diagnostics
+/// candidate — can zip results back to their inputs).
+///
+/// Every query `check_file` reaches is read-only, and `ProjectDatabase`'s
+/// `Clone` produces a shared-storage salsa handle (the rust-analyzer
+/// concurrency model): workers share one memo table, so a scope inferred by
+/// one thread is a cache hit for every other. `ProjectDatabase` is `Send`
+/// but deliberately not `Sync` (each salsa handle carries thread-confined
+/// query-stack state), so tasks cannot share `&db` — every task MOVES its
+/// own cloned handle instead (an Arc bump; all clones share one memo table).
+/// Small chunks keep work-stealing effective on rayon's global pool — files
+/// vary a lot in check cost — while amortizing the per-task clone.
+pub fn check_files_parallel(db: &ProjectDatabase, files: &[SourceFile]) -> Vec<Vec<Diagnostic>> {
     const CHUNK: usize = 4;
 
-    let Some((first, rest)) = source_files.split_first() else {
-        return;
-    };
-    diagnostics.extend(lsp2_check_file(db, *first));
+    prime_file_indexes_parallel(db);
 
-    let chunks: Vec<&[SourceFile]> = rest.chunks(CHUNK).collect();
+    let chunks: Vec<(usize, &[SourceFile])> = files.chunks(CHUNK).enumerate().collect();
     // Handles are cloned OUTSIDE the rayon scope: a `!Sync` database can't be
     // borrowed by the (Send) scope closure, so each chunk's handle is created
     // up front and MOVED into its task.
     let handles: Vec<ProjectDatabase> = chunks.iter().map(|_| db.clone()).collect();
-    let (tx, rx) = std::sync::mpsc::channel::<Vec<Diagnostic>>();
+    let (tx, rx) = std::sync::mpsc::channel::<(usize, Vec<Vec<Diagnostic>>)>();
     rayon::scope(move |s| {
-        for (chunk, db) in chunks.into_iter().zip(handles) {
+        for ((chunk_index, chunk), db) in chunks.into_iter().zip(handles) {
             let tx = tx.clone();
             s.spawn(move |_| {
-                let mut out = Vec::new();
-                for file in chunk {
-                    out.extend(lsp2_check_file(&db, *file));
-                }
+                let out: Vec<Vec<Diagnostic>> = chunk
+                    .iter()
+                    .map(|file| lsp2_check_file(&db, *file))
+                    .collect();
                 // Receiver outlives the scope; a send only fails if it
                 // dropped early, which would mean a panic elsewhere.
-                let _ = tx.send(out);
+                let _ = tx.send((chunk_index, out));
             });
         }
     });
-    for out in rx {
-        diagnostics.extend(out);
+    let mut slots: Vec<Option<Vec<Diagnostic>>> = (0..files.len()).map(|_| None).collect();
+    for (chunk_index, out) in rx {
+        for (offset, file_diagnostics) in out.into_iter().enumerate() {
+            slots[chunk_index * CHUNK + offset] = Some(file_diagnostics);
+        }
     }
+    slots.into_iter().map(Option::unwrap_or_default).collect()
 }
 
 /// The per-checked-file split alongside the merged, honest-ordered set produced

baml_language/crates/baml_project/src/db.rs

@@ -474,6 +474,44 @@ impl ProjectDatabase {
         }
     }
 
+    /// Bulk [`Self::add_or_update_file`]: identical per-file semantics
+    /// (canonicalization, tombstone revival, map registration), but the
+    /// project file list is written once at the end instead of once per new
+    /// file. The per-file path clones and re-sets the whole `files` Vec and
+    /// bumps the salsa revision each time — O(files²) copies plus one
+    /// revision per file during initial project load.
+    pub fn add_or_update_files<'a, I>(&mut self, files: I)
+    where
+        I: IntoIterator<Item = (&'a std::path::Path, &'a str)>,
+    {
+        let mut new_files = Vec::new();
+        for (path, content) in files {
+            let canonical_path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
+
+            if let Some(&existing_file) = self.file_map.get(&canonical_path) {
+                existing_file.set_text(self).to(content.to_string());
+                continue;
+            }
+            let file = if let Some(file) = self.removed_file_tombstones.remove(&canonical_path) {
+                file.set_text(self).to(content.to_string());
+                file
+            } else {
+                self.add_file_internal(&canonical_path, content)
+            };
+            let file_id = file.file_id(self);
+            Arc::make_mut(&mut self.file_map).insert(canonical_path.clone(), file);
+            Arc::make_mut(&mut self.file_id_to_path).insert(file_id, canonical_path);
+            new_files.push(file);
+        }
+        if !new_files.is_empty()
+            && let Some(project) = self.project
+        {
+            let mut project_files: Vec<SourceFile> = project.files(self).clone();
+            project_files.extend(new_files);
+            project.set_files(self).to(project_files);
+        }
+    }
+
     /// Remove a file from the database.
     ///
     /// Note: Salsa doesn't support true removal. The input is emptied (so its

baml_language/crates/baml_project/src/lib.rs

@@ -32,7 +32,7 @@ pub mod symbols;
 pub mod testing;
 
 pub use check::{
-    CheckResult, NarrowedDiagnostics, collect_compiler2_diagnostics,
+    CheckResult, NarrowedDiagnostics, check_files_parallel, collect_compiler2_diagnostics,
     collect_compiler2_diagnostics_narrowed, collect_diagnostics, collect_package_level_diagnostics,
 };
 pub use client_codegen::build_symbol_pool;

baml_language/crates/bex_project/src/project.rs

@@ -224,15 +224,19 @@ pub(crate) fn collect_diagnostic_candidate(guard: &SourceGuard<'_>) -> Diagnosti
     let mut documents = Vec::new();
     let mut has_errors = false;
 
-    for file in &source_files {
+    // Check across worker threads (input order preserved, so per-file zip is
+    // sound). This runs under the held source guard, so no mutation can race
+    // the cloned worker handles; the guard is the same exclusion the serial
+    // loop relied on.
+    let per_file = baml_project::check_files_parallel(db, &source_files);
+    for (file, diagnostics) in source_files.iter().zip(per_file) {
         let file_id = file.file_id(db);
         let Some(path) = db.file_id_to_path(file_id).cloned() else {
             continue;
         };
         let text = file.text(db).clone();
         file_texts.insert(file_id, (path.clone(), text));
 
-        let diagnostics = baml_lsp2_actions::check_file(db, *file);
         has_errors |= diagnostics
             .iter()
             .any(|d| d.severity == baml_compiler_diagnostics::Severity::Error);