Merge pull request #19565 from davidbarsky/davidbarsky/add-prime-caches-subcommand

internal: add `prime-caches` subcommand

crates/rust-analyzer/src/bin/main.rs

@@ -88,6 +88,7 @@ fn actual_main() -> anyhow::Result<ExitCode> {
         flags::RustAnalyzerCmd::Scip(cmd) => cmd.run()?,
         flags::RustAnalyzerCmd::RunTests(cmd) => cmd.run()?,
         flags::RustAnalyzerCmd::RustcTests(cmd) => cmd.run()?,
+        flags::RustAnalyzerCmd::PrimeCaches(cmd) => cmd.run()?,
     }
     Ok(ExitCode::SUCCESS)
 }

crates/rust-analyzer/src/cli.rs

@@ -8,6 +8,7 @@ pub mod flags;
 mod highlight;
 mod lsif;
 mod parse;
+mod prime_caches;
 mod run_tests;
 mod rustc_tests;
 mod scip;

crates/rust-analyzer/src/cli/analysis_stats.rs

@@ -437,7 +437,7 @@ impl flags::AnalysisStats {
         let mut bar = match verbosity {
             Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
             _ if self.parallel || self.output.is_some() => ProgressReport::hidden(),
-            _ => ProgressReport::new(file_ids.len() as u64),
+            _ => ProgressReport::new(file_ids.len()),
         };
 
         file_ids.sort();
@@ -646,7 +646,7 @@ impl flags::AnalysisStats {
         let mut bar = match verbosity {
             Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
             _ if self.parallel || self.output.is_some() => ProgressReport::hidden(),
-            _ => ProgressReport::new(bodies.len() as u64),
+            _ => ProgressReport::new(bodies.len()),
         };
         let mut sw = self.stop_watch();
         let mut all = 0;
@@ -692,7 +692,7 @@ impl flags::AnalysisStats {
         let mut bar = match verbosity {
             Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
             _ if self.parallel || self.output.is_some() => ProgressReport::hidden(),
-            _ => ProgressReport::new(bodies.len() as u64),
+            _ => ProgressReport::new(bodies.len()),
         };
 
         if self.parallel {
@@ -1023,7 +1023,7 @@ impl flags::AnalysisStats {
         let mut bar = match verbosity {
             Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
             _ if self.output.is_some() => ProgressReport::hidden(),
-            _ => ProgressReport::new(bodies.len() as u64),
+            _ => ProgressReport::new(bodies.len()),
         };
 
         let mut sw = self.stop_watch();

crates/rust-analyzer/src/cli/flags.rs

@@ -53,7 +53,7 @@ xflags::xflags! {
 
         /// Batch typecheck project and print summary statistics
         cmd analysis-stats {
-            /// Directory with Cargo.toml.
+            /// Directory with Cargo.toml or rust-project.json.
             required path: PathBuf
 
             optional --output format: OutputFormat
@@ -74,7 +74,7 @@ xflags::xflags! {
 
             /// Don't run build scripts or load `OUT_DIR` values by running `cargo check` before analysis.
             optional --disable-build-scripts
-            /// Don't use expand proc macros.
+            /// Don't expand proc macros.
             optional --disable-proc-macros
             /// Run the proc-macro-srv binary at the specified path.
             optional --proc-macro-srv path: PathBuf
@@ -101,7 +101,7 @@ xflags::xflags! {
 
         /// Run unit tests of the project using mir interpreter
         cmd run-tests {
-            /// Directory with Cargo.toml.
+            /// Directory with Cargo.toml or rust-project.json.
             required path: PathBuf
         }
 
@@ -115,12 +115,12 @@ xflags::xflags! {
         }
 
         cmd diagnostics {
-            /// Directory with Cargo.toml.
+            /// Directory with Cargo.toml or rust-project.json.
             required path: PathBuf
 
             /// Don't run build scripts or load `OUT_DIR` values by running `cargo check` before analysis.
             optional --disable-build-scripts
-            /// Don't use expand proc macros.
+            /// Don't expand proc macros.
             optional --disable-proc-macros
             /// Run the proc-macro-srv binary at the specified path.
             optional --proc-macro-srv path: PathBuf
@@ -128,17 +128,32 @@ xflags::xflags! {
 
         /// Report unresolved references
         cmd unresolved-references {
-            /// Directory with Cargo.toml.
+            /// Directory with Cargo.toml or rust-project.json.
             required path: PathBuf
 
             /// Don't run build scripts or load `OUT_DIR` values by running `cargo check` before analysis.
             optional --disable-build-scripts
-            /// Don't use expand proc macros.
+            /// Don't expand proc macros.
             optional --disable-proc-macros
             /// Run the proc-macro-srv binary at the specified path.
             optional --proc-macro-srv path: PathBuf
         }
 
+        /// Prime caches, as rust-analyzer does typically at startup in interactive sessions.
+        cmd prime-caches {
+            /// Directory with Cargo.toml or rust-project.json.
+            required path: PathBuf
+
+            /// Don't run build scripts or load `OUT_DIR` values by running `cargo check` before analysis.
+            optional --disable-build-scripts
+            /// Don't expand proc macros.
+            optional --disable-proc-macros
+            /// Run the proc-macro-srv binary at the specified path.
+            optional --proc-macro-srv path: PathBuf
+            /// Run cache priming in parallel.
+            optional --parallel
+        }
+
         cmd ssr {
             /// A structured search replace rule (`$a.foo($b) ==>> bar($a, $b)`)
             repeated rule: SsrRule
@@ -197,6 +212,7 @@ pub enum RustAnalyzerCmd {
     RustcTests(RustcTests),
     Diagnostics(Diagnostics),
     UnresolvedReferences(UnresolvedReferences),
+    PrimeCaches(PrimeCaches),
     Ssr(Ssr),
     Search(Search),
     Lsif(Lsif),
@@ -276,6 +292,16 @@ pub struct UnresolvedReferences {
     pub proc_macro_srv: Option<PathBuf>,
 }
 
+#[derive(Debug)]
+pub struct PrimeCaches {
+    pub path: PathBuf,
+
+    pub disable_build_scripts: bool,
+    pub disable_proc_macros: bool,
+    pub proc_macro_srv: Option<PathBuf>,
+    pub parallel: bool,
+}
+
 #[derive(Debug)]
 pub struct Ssr {
     pub rule: Vec<SsrRule>,

crates/rust-analyzer/src/cli/prime_caches.rs

@@ -0,0 +1,67 @@
+//! Load the project and run cache priming.
+//!
+//! Unlike `analysis-stats`, this command is intended to be used for
+//! benchmarking rust-analyzer's default startup configuration. It *does not*
+//! attempt to simulate the full IDE experience through the lifetime of the
+//! an editing session.
+
+use load_cargo::{LoadCargoConfig, ProcMacroServerChoice, load_workspace};
+use profile::StopWatch;
+use project_model::{ProjectManifest, ProjectWorkspace};
+use vfs::AbsPathBuf;
+
+use crate::cli::flags;
+
+impl flags::PrimeCaches {
+    pub fn run(self) -> anyhow::Result<()> {
+        let root =
+            vfs::AbsPathBuf::assert_utf8(std::env::current_dir()?.join(&self.path)).normalize();
+        let config = crate::config::Config::new(
+            root.clone(),
+            lsp_types::ClientCapabilities::default(),
+            vec![],
+            None,
+        );
+        let mut stop_watch = StopWatch::start();
+
+        let cargo_config = config.cargo(None);
+        let with_proc_macro_server = if let Some(p) = &self.proc_macro_srv {
+            let path = vfs::AbsPathBuf::assert_utf8(std::env::current_dir()?.join(p));
+            ProcMacroServerChoice::Explicit(path)
+        } else {
+            ProcMacroServerChoice::Sysroot
+        };
+        let load_cargo_config = LoadCargoConfig {
+            load_out_dirs_from_check: !self.disable_build_scripts,
+            with_proc_macro_server,
+            // while this command is nominally focused on cache priming,
+            // we want to ensure that this command, not `load_workspace_at`,
+            // is responsible for that work.
+            prefill_caches: false,
+        };
+
+        let root = AbsPathBuf::assert_utf8(std::env::current_dir()?.join(root));
+        let root = ProjectManifest::discover_single(&root)?;
+        let workspace = ProjectWorkspace::load(root, &cargo_config, &|_| {})?;
+
+        let (db, _, _) = load_workspace(workspace, &cargo_config.extra_env, &load_cargo_config)?;
+        let elapsed = stop_watch.elapsed();
+        eprintln!(
+            "Load time: {:?}ms, memory allocated: {}MB",
+            elapsed.time.as_millis(),
+            elapsed.memory.allocated.megabytes() as u64
+        );
+
+        let threads = if self.parallel { num_cpus::get() } else { 1 };
+        ide_db::prime_caches::parallel_prime_caches(&db, threads, &|_| ());
+
+        let elapsed = stop_watch.elapsed();
+        eprintln!(
+            "Cache priming time: {:?}ms, total memory allocated: {}MB",
+            elapsed.time.as_millis(),
+            elapsed.memory.allocated.megabytes() as u64
+        );
+
+        Ok(())
+    }
+}

crates/rust-analyzer/src/cli/progress_report.rs

@@ -9,13 +9,13 @@ pub(crate) struct ProgressReport<'a> {
     text: String,
     hidden: bool,
 
-    len: u64,
+    len: usize,
     pos: u64,
     msg: Option<Box<dyn Fn() -> String + 'a>>,
 }
 
 impl<'a> ProgressReport<'a> {
-    pub(crate) fn new(len: u64) -> ProgressReport<'a> {
+    pub(crate) fn new(len: usize) -> ProgressReport<'a> {
         ProgressReport { curr: 0.0, text: String::new(), hidden: false, len, pos: 0, msg: None }
     }