Merge #5142

5142: analysis-stats: allow parallel type inference r=matklad a=jonas-schievink

This is mostly just for testing/fun, but it looks like type inference can be sped up massively with little to no effort (since it runs after the serial phases are already done).

Without `--parallel`:

```
Item Collection: 16.43597698s, 683mb allocated 720mb resident
Inference: 25.429774879s, 1720mb allocated 1781mb resident
Total: 41.865866352s, 1720mb allocated 1781mb resident
```

With `--parallel`:

```
Item Collection: 16.380369815s, 683mb allocated 735mb resident
Parallel Inference: 7.449166445s, 1721mb allocated 1812mb resident
Inference: 143.437157ms, 1721mb allocated 1812mb resident
Total: 23.973303611s, 1721mb allocated 1812mb resident
```

Co-authored-by: Jonas Schievink <jonas.schievink@ferrous-systems.com>

Cargo.lock

@@ -1461,6 +1461,7 @@ dependencies = [
  "ra_toolchain",
  "ra_tt",
  "rand",
+ "rayon",
  "rustc-hash",
  "serde",
  "serde_json",

crates/rust-analyzer/Cargo.toml

@@ -28,6 +28,7 @@ rustc-hash = "1.1.0"
 serde = { version = "1.0.106", features = ["derive"] }
 serde_json = "1.0.48"
 threadpool = "1.7.1"
+rayon = "1.3.1"
 
 stdx = { path = "../stdx" }
 

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

@@ -25,6 +25,7 @@ pub(crate) enum Command {
     },
     Stats {
         randomize: bool,
+        parallel: bool,
         memory_usage: bool,
         only: Option<String>,
         with_deps: bool,
@@ -157,10 +158,14 @@ USAGE:
     rust-analyzer analysis-stats [FLAGS] [OPTIONS] [PATH]
 
 FLAGS:
+    -o, --only              Only analyze items matching this path
     -h, --help              Prints help information
-        --memory-usage
+        --memory-usage      Collect memory usage statistics (requires `--feature jemalloc`)
+        --randomize         Randomize order in which crates, modules, and items are processed
+        --parallel          Run type inference in parallel
         --load-output-dirs  Load OUT_DIR values by running `cargo check` before analysis
         --with-proc-macro   Use ra-proc-macro-srv for proc-macro expanding
+        --with-deps         Also analyze all dependencies
     -v, --verbose
     -q, --quiet
 
@@ -174,6 +179,7 @@ ARGS:
                 }
 
                 let randomize = matches.contains("--randomize");
+                let parallel = matches.contains("--parallel");
                 let memory_usage = matches.contains("--memory-usage");
                 let only: Option<String> = matches.opt_value_from_str(["-o", "--only"])?;
                 let with_deps: bool = matches.contains("--with-deps");
@@ -189,6 +195,7 @@ ARGS:
 
                 Command::Stats {
                     randomize,
+                    parallel,
                     memory_usage,
                     only,
                     with_deps,

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

@@ -32,6 +32,7 @@ fn main() -> Result<()> {
         args::Command::Highlight { rainbow } => cli::highlight(rainbow)?,
         args::Command::Stats {
             randomize,
+            parallel,
             memory_usage,
             only,
             with_deps,
@@ -45,6 +46,7 @@ fn main() -> Result<()> {
             only.as_ref().map(String::as_ref),
             with_deps,
             randomize,
+            parallel,
             load_output_dirs,
             with_proc_macro,
         )?,

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

@@ -5,6 +5,7 @@ use std::{path::Path, time::Instant};
 
 use itertools::Itertools;
 use rand::{seq::SliceRandom, thread_rng};
+use rayon::prelude::*;
 use rustc_hash::FxHashSet;
 
 use hir::{
@@ -13,12 +14,23 @@ use hir::{
 };
 use hir_def::FunctionId;
 use hir_ty::{Ty, TypeWalk};
-use ra_db::SourceDatabaseExt;
+use ra_db::{
+    salsa::{self, ParallelDatabase},
+    SourceDatabaseExt,
+};
 use ra_syntax::AstNode;
 use stdx::format_to;
 
 use crate::cli::{load_cargo::load_cargo, progress_report::ProgressReport, Result, Verbosity};
 
+/// Need to wrap Snapshot to provide `Clone` impl for `map_with`
+struct Snap<DB>(DB);
+impl<DB: ParallelDatabase> Clone for Snap<salsa::Snapshot<DB>> {
+    fn clone(&self) -> Snap<salsa::Snapshot<DB>> {
+        Snap(self.0.snapshot())
+    }
+}
+
 pub fn analysis_stats(
     verbosity: Verbosity,
     memory_usage: bool,
@@ -26,6 +38,7 @@ pub fn analysis_stats(
     only: Option<&str>,
     with_deps: bool,
     randomize: bool,
+    parallel: bool,
     load_output_dirs: bool,
     with_proc_macro: bool,
 ) -> Result<()> {
@@ -91,12 +104,26 @@ pub fn analysis_stats(
         funcs.shuffle(&mut thread_rng());
     }
 
-    let inference_time = Instant::now();
     let mut bar = match verbosity {
         Verbosity::Quiet | Verbosity::Spammy => ProgressReport::hidden(),
         _ => ProgressReport::new(funcs.len() as u64),
     };
 
+    if parallel {
+        let inference_time = Instant::now();
+        let snap = Snap(db.snapshot());
+        funcs
+            .par_iter()
+            .map_with(snap, |snap, &f| {
+                let f_id = FunctionId::from(f);
+                snap.0.body(f_id.into());
+                snap.0.infer(f_id.into());
+            })
+            .count();
+        println!("Parallel Inference: {:?}, {}", inference_time.elapsed(), ra_prof::memory_usage());
+    }
+
+    let inference_time = Instant::now();
     bar.tick();
     let mut num_exprs = 0;
     let mut num_exprs_unknown = 0;