Split cache priming into distinct phases

2 files changed, 95 insertions, 34 deletions

diff --git a/src/tools/rust-analyzer/crates/ide-db/src/prime_caches.rs b/src/tools/rust-analyzer/crates/ide-db/src/prime_caches.rs
index 0002fda0ba7..238b2b3c7d5 100644
--- a/src/tools/rust-analyzer/crates/ide-db/src/prime_caches.rs
+++ b/src/tools/rust-analyzer/crates/ide-db/src/prime_caches.rs
@@ -7,11 +7,12 @@ mod topologic_sort;
 use std::time::Duration;
 
 use hir::db::DefDatabase;
+use itertools::Itertools;
 
 use crate::{
     base_db::{
         ra_salsa::{Database, ParallelDatabase, Snapshot},
-        Cancelled, CrateId, SourceDatabase, SourceRootDatabase,
+        Cancelled, CrateId, SourceDatabase,
     },
     symbol_index::SymbolsDatabase,
     FxIndexMap, RootDatabase,
@@ -26,6 +27,7 @@ pub struct ParallelPrimeCachesProgress {
     pub crates_total: usize,
     /// the total number of crates that have finished priming
     pub crates_done: usize,
+    pub work_type: &'static str,
 }
 
 pub fn parallel_prime_caches(
@@ -51,37 +53,28 @@ pub fn parallel_prime_caches(
         EndCrate { crate_id: CrateId },
     }
 
+    // We split off def map computation from other work,
+    // as the def map is the relevant one. Once the defmaps are computed
+    // the project is ready to go, the other indices are just nice to have for some IDE features.
+    #[derive(PartialOrd, Ord, PartialEq, Eq, Copy, Clone)]
+    enum PrimingPhase {
+        DefMap,
+        ImportMap,
+        CrateSymbols,
+    }
+
     let (work_sender, progress_receiver) = {
         let (progress_sender, progress_receiver) = crossbeam_channel::unbounded();
         let (work_sender, work_receiver) = crossbeam_channel::unbounded();
-        let graph = graph.clone();
-        let local_roots = db.local_roots();
         let prime_caches_worker = move |db: Snapshot<RootDatabase>| {
-            while let Ok((crate_id, crate_name)) = work_receiver.recv() {
+            while let Ok((crate_id, crate_name, kind)) = work_receiver.recv() {
                 progress_sender
                     .send(ParallelPrimeCacheWorkerProgress::BeginCrate { crate_id, crate_name })?;
 
-                // Compute the DefMap and possibly ImportMap
-                let file_id = graph[crate_id].root_file_id;
-                let root_id = db.file_source_root(file_id);
-                if db.source_root(root_id).is_library {
-                    db.crate_def_map(crate_id);
-                } else {
-                    // This also computes the DefMap
-                    db.import_map(crate_id);
-                }
-
-                // Compute the symbol search index.
-                // This primes the cache for `ide_db::symbol_index::world_symbols()`.
-                //
-                // We do this for workspace crates only (members of local_roots), because doing it
-                // for all dependencies could be *very* unnecessarily slow in a large project.
-                //
-                // FIXME: We should do it unconditionally if the configuration is set to default to
-                // searching dependencies (rust-analyzer.workspace.symbol.search.scope), but we
-                // would need to pipe that configuration information down here.
-                if local_roots.contains(&root_id) {
-                    db.crate_symbols(crate_id.into());
+                match kind {
+                    PrimingPhase::DefMap => _ = db.crate_def_map(crate_id),
+                    PrimingPhase::ImportMap => _ = db.import_map(crate_id),
+                    PrimingPhase::CrateSymbols => _ = db.crate_symbols(crate_id.into()),
                 }
 
                 progress_sender.send(ParallelPrimeCacheWorkerProgress::EndCrate { crate_id })?;
@@ -112,16 +105,30 @@ pub fn parallel_prime_caches(
     let mut crates_currently_indexing =
         FxIndexMap::with_capacity_and_hasher(num_worker_threads, Default::default());
 
+    let mut additional_phases = vec![];
+
     while crates_done < crates_total {
         db.unwind_if_cancelled();
 
         for crate_id in &mut crates_to_prime {
-            work_sender
-                .send((
-                    crate_id,
-                    graph[crate_id].display_name.as_deref().unwrap_or_default().to_owned(),
-                ))
-                .ok();
+            let krate = &graph[crate_id];
+            let name = krate.display_name.as_deref().unwrap_or_default().to_owned();
+            if krate.origin.is_lang() {
+                additional_phases.push((crate_id, name.clone(), PrimingPhase::ImportMap));
+            } else if krate.origin.is_local() {
+                // Compute the symbol search index.
+                // This primes the cache for `ide_db::symbol_index::world_symbols()`.
+                //
+                // We do this for workspace crates only (members of local_roots), because doing it
+                // for all dependencies could be *very* unnecessarily slow in a large project.
+                //
+                // FIXME: We should do it unconditionally if the configuration is set to default to
+                // searching dependencies (rust-analyzer.workspace.symbol.search.scope), but we
+                // would need to pipe that configuration information down here.
+                additional_phases.push((crate_id, name.clone(), PrimingPhase::CrateSymbols));
+            }
+
+            work_sender.send((crate_id, name, PrimingPhase::DefMap)).ok();
         }
 
         // recv_timeout is somewhat a hack, we need a way to from this thread check to see if the current salsa revision
@@ -153,6 +160,51 @@ pub fn parallel_prime_caches(
             crates_currently_indexing: crates_currently_indexing.values().cloned().collect(),
             crates_done,
             crates_total,
+            work_type: "Indexing",
+        };
+
+        cb(progress);
+    }
+
+    let mut crates_done = 0;
+    let crates_total = additional_phases.len();
+    for w in additional_phases.into_iter().sorted_by_key(|&(_, _, phase)| phase) {
+        work_sender.send(w).ok();
+    }
+
+    while crates_done < crates_total {
+        db.unwind_if_cancelled();
+
+        // recv_timeout is somewhat a hack, we need a way to from this thread check to see if the current salsa revision
+        // is cancelled on a regular basis. workers will only exit if they are processing a task that is cancelled, or
+        // if this thread exits, and closes the work channel.
+        let worker_progress = match progress_receiver.recv_timeout(Duration::from_millis(10)) {
+            Ok(p) => p,
+            Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
+                continue;
+            }
+            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
+                // our workers may have died from a cancelled task, so we'll check and re-raise here.
+                db.unwind_if_cancelled();
+                break;
+            }
+        };
+        match worker_progress {
+            ParallelPrimeCacheWorkerProgress::BeginCrate { crate_id, crate_name } => {
+                crates_currently_indexing.insert(crate_id, crate_name);
+            }
+            ParallelPrimeCacheWorkerProgress::EndCrate { crate_id } => {
+                crates_currently_indexing.swap_remove(&crate_id);
+                crates_to_prime.mark_done(crate_id);
+                crates_done += 1;
+            }
+        };
+
+        let progress = ParallelPrimeCachesProgress {
+            crates_currently_indexing: crates_currently_indexing.values().cloned().collect(),
+            crates_done,
+            crates_total,
+            work_type: "Populating symbols",
         };
 
         cb(progress);
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
index de0ee639ca5..d4a90181efe 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
@@ -325,17 +325,19 @@ impl GlobalState {
                 }
 
                 for progress in prime_caches_progress {
-                    let (state, message, fraction);
+                    let (state, message, fraction, title);
                     match progress {
                         PrimeCachesProgress::Begin => {
                             state = Progress::Begin;
                             message = None;
                             fraction = 0.0;
+                            title = "Indexing";
                         }
                         PrimeCachesProgress::Report(report) => {
                             state = Progress::Report;
+                            title = report.work_type;
 
-                            message = match &report.crates_currently_indexing[..] {
+                            message = match &*report.crates_currently_indexing {
                                 [crate_name] => Some(format!(
                                     "{}/{} ({crate_name})",
                                     report.crates_done, report.crates_total
@@ -356,6 +358,7 @@ impl GlobalState {
                             state = Progress::End;
                             message = None;
                             fraction = 1.0;
+                            title = "Indexing";
 
                             self.prime_caches_queue.op_completed(());
                             if cancelled {
@@ -365,7 +368,13 @@ impl GlobalState {
                         }
                     };
 
-                    self.report_progress("Indexing", state, message, Some(fraction), None);
+                    self.report_progress(
+                        title,
+                        state,
+                        message,
+                        Some(fraction),
+                        Some("rustAnalyzer/cachePriming".to_owned()),
+                    );
                 }
             }
             Event::Vfs(message) => {