Improve prime_caches and display its progress

5 files changed, 119 insertions, 13 deletions

diff --git a/crates/base_db/src/input.rs b/crates/base_db/src/input.rs
index c330314d4..215ac4b41 100644
--- a/crates/base_db/src/input.rs
+++ b/crates/base_db/src/input.rs
@@ -221,6 +221,34 @@ impl CrateGraph {
         deps.into_iter()
     }
 
+    /// Returns all crates in the graph, sorted in topological order (ie. dependencies of a crate
+    /// come before the crate itself).
+    pub fn crates_in_topological_order(&self) -> Vec<CrateId> {
+        let mut res = Vec::new();
+        let mut visited = FxHashSet::default();
+
+        for krate in self.arena.keys().copied() {
+            go(self, &mut visited, &mut res, krate);
+        }
+
+        return res;
+
+        fn go(
+            graph: &CrateGraph,
+            visited: &mut FxHashSet<CrateId>,
+            res: &mut Vec<CrateId>,
+            source: CrateId,
+        ) {
+            if !visited.insert(source) {
+                return;
+            }
+            for dep in graph[source].dependencies.iter() {
+                go(graph, visited, res, dep.crate_id)
+            }
+            res.push(source)
+        }
+    }
+
     // FIXME: this only finds one crate with the given root; we could have multiple
     pub fn crate_id_for_crate_root(&self, file_id: FileId) -> Option<CrateId> {
         let (&crate_id, _) =
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 686cee3a1..aaf9b3b4b 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -77,6 +77,7 @@ pub use crate::{
     hover::{HoverAction, HoverConfig, HoverGotoTypeData, HoverResult},
     inlay_hints::{InlayHint, InlayHintsConfig, InlayKind},
     markup::Markup,
+    prime_caches::PrimeCachesProgress,
     references::{
         Declaration, Reference, ReferenceAccess, ReferenceKind, ReferenceSearchResult, RenameError,
     },
@@ -223,8 +224,11 @@ impl Analysis {
         self.with_db(|db| status::status(&*db, file_id))
     }
 
-    pub fn prime_caches(&self, files: Vec<FileId>) -> Cancelable<()> {
-        self.with_db(|db| prime_caches::prime_caches(db, files))
+    pub fn prime_caches<F>(&self, cb: F) -> Cancelable<()>
+    where
+        F: Fn(PrimeCachesProgress) + Sync + std::panic::UnwindSafe,
+    {
+        self.with_db(move |db| prime_caches::prime_caches(db, &cb))
     }
 
     /// Gets the text of the source file.
diff --git a/crates/ide/src/prime_caches.rs b/crates/ide/src/prime_caches.rs
index c5ab5a1d8..9687c2734 100644
--- a/crates/ide/src/prime_caches.rs
+++ b/crates/ide/src/prime_caches.rs
@@ -3,10 +3,45 @@
 //! request takes longer to compute. This modules implemented prepopulating of
 //! various caches, it's not really advanced at the moment.
 
-use crate::{FileId, RootDatabase};
+use base_db::SourceDatabase;
+use hir::db::DefDatabase;
 
-pub(crate) fn prime_caches(db: &RootDatabase, files: Vec<FileId>) {
-    for file in files {
-        let _ = crate::syntax_highlighting::highlight(db, file, None, false);
+use crate::RootDatabase;
+
+#[derive(Debug)]
+pub enum PrimeCachesProgress {
+    Started,
+    /// We started indexing a crate.
+    StartedOnCrate {
+        on_crate: String,
+        n_done: usize,
+        n_total: usize,
+    },
+    /// We finished indexing all crates.
+    Finished,
+}
+
+pub(crate) fn prime_caches(db: &RootDatabase, cb: &(dyn Fn(PrimeCachesProgress) + Sync)) {
+    let _p = profile::span("prime_caches");
+    let graph = db.crate_graph();
+    let topo = &graph.crates_in_topological_order();
+
+    cb(PrimeCachesProgress::Started);
+
+    // FIXME: This would be easy to parallelize, since it's in the ideal ordering for that.
+    // Unfortunately rayon prevents panics from propagation out of a `scope`, which breaks
+    // cancellation, so we cannot use rayon.
+    for (i, krate) in topo.iter().enumerate() {
+        let crate_name =
+            graph[*krate].declaration_name.as_ref().map(ToString::to_string).unwrap_or_default();
+
+        cb(PrimeCachesProgress::StartedOnCrate {
+            on_crate: crate_name,
+            n_done: i,
+            n_total: topo.len(),
+        });
+        db.crate_def_map(*krate);
     }
+
+    cb(PrimeCachesProgress::Finished);
 }
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index 06b38d99c..fb18f9014 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -7,6 +7,7 @@ use std::{
 
 use base_db::VfsPath;
 use crossbeam_channel::{select, Receiver};
+use ide::PrimeCachesProgress;
 use ide::{Canceled, FileId};
 use lsp_server::{Connection, Notification, Request, Response};
 use lsp_types::notification::Notification as _;
@@ -61,7 +62,7 @@ pub(crate) enum Task {
     Response(Response),
     Diagnostics(Vec<(FileId, Vec<lsp_types::Diagnostic>)>),
     Workspaces(Vec<anyhow::Result<ProjectWorkspace>>),
-    Unit,
+    PrimeCaches(PrimeCachesProgress),
 }
 
 impl fmt::Debug for Event {
@@ -197,7 +198,28 @@ impl GlobalState {
                     }
                 }
                 Task::Workspaces(workspaces) => self.switch_workspaces(workspaces),
-                Task::Unit => (),
+                Task::PrimeCaches(progress) => {
+                    let (state, message, fraction);
+                    match progress {
+                        PrimeCachesProgress::Started => {
+                            state = Progress::Begin;
+                            message = None;
+                            fraction = 0.0;
+                        }
+                        PrimeCachesProgress::StartedOnCrate { on_crate, n_done, n_total } => {
+                            state = Progress::Report;
+                            message = Some(format!("{}/{} ({})", n_done, n_total, on_crate));
+                            fraction = Progress::fraction(n_done, n_total);
+                        }
+                        PrimeCachesProgress::Finished => {
+                            state = Progress::End;
+                            message = None;
+                            fraction = 1.0;
+                        }
+                    };
+
+                    self.report_progress("indexing", state, message, Some(fraction));
+                }
             },
             Event::Vfs(mut task) => {
                 let _p = profile::span("GlobalState::handle_event/vfs");
@@ -573,12 +595,18 @@ impl GlobalState {
                 Task::Diagnostics(diagnostics)
             })
         }
-        self.task_pool.handle.spawn({
-            let subs = subscriptions;
+        self.task_pool.handle.spawn_with_sender({
             let snap = self.snapshot();
-            move || {
-                snap.analysis.prime_caches(subs).unwrap_or_else(|_: Canceled| ());
-                Task::Unit
+            move |sender| {
+                snap.analysis
+                    .prime_caches(|progress| {
+                        sender.send(Task::PrimeCaches(progress)).unwrap();
+                    })
+                    .unwrap_or_else(|_: Canceled| {
+                        // Pretend that we're done, so that the progress bar is removed. Otherwise
+                        // the editor may complain about it already existing.
+                        sender.send(Task::PrimeCaches(PrimeCachesProgress::Finished)).unwrap()
+                    });
             }
         });
     }
diff --git a/crates/rust-analyzer/src/thread_pool.rs b/crates/rust-analyzer/src/thread_pool.rs
index 4fa502925..833893739 100644
--- a/crates/rust-analyzer/src/thread_pool.rs
+++ b/crates/rust-analyzer/src/thread_pool.rs
@@ -23,6 +23,17 @@ impl<T> TaskPool<T> {
         })
     }
 
+    pub(crate) fn spawn_with_sender<F>(&mut self, task: F)
+    where
+        F: FnOnce(Sender<T>) + Send + 'static,
+        T: Send + 'static,
+    {
+        self.inner.execute({
+            let sender = self.sender.clone();
+            move || task(sender)
+        })
+    }
+
     pub(crate) fn len(&self) -> usize {
         self.inner.queued_count()
     }