Add `xdist dist --pgo` command to build Rust Analyzer with PGO optimizations

2 files changed, 98 insertions, 6 deletions

diff --git a/src/tools/rust-analyzer/xtask/src/dist.rs b/src/tools/rust-analyzer/xtask/src/dist.rs
index 35ae7965f33..a07244252bf 100644
--- a/src/tools/rust-analyzer/xtask/src/dist.rs
+++ b/src/tools/rust-analyzer/xtask/src/dist.rs
@@ -1,13 +1,15 @@
+use anyhow::Context;
+use flate2::{Compression, write::GzEncoder};
+use std::env::consts::EXE_EXTENSION;
+use std::ffi::OsStr;
 use std::{
     env,
     fs::File,
     io::{self, BufWriter},
     path::{Path, PathBuf},
 };
-
-use flate2::{Compression, write::GzEncoder};
 use time::OffsetDateTime;
-use xshell::{Shell, cmd};
+use xshell::{Cmd, Shell, cmd};
 use zip::{DateTime, ZipWriter, write::SimpleFileOptions};
 
 use crate::{
@@ -38,11 +40,18 @@ impl flags::Dist {
                 // A hack to make VS Code prefer nightly over stable.
                 format!("{VERSION_NIGHTLY}.{patch_version}")
             };
-            dist_server(sh, &format!("{version}-standalone"), &target, allocator, self.zig)?;
+            dist_server(
+                sh,
+                &format!("{version}-standalone"),
+                &target,
+                allocator,
+                self.zig,
+                self.pgo,
+            )?;
             let release_tag = if stable { date_iso(sh)? } else { "nightly".to_owned() };
             dist_client(sh, &version, &release_tag, &target)?;
         } else {
-            dist_server(sh, "0.0.0-standalone", &target, allocator, self.zig)?;
+            dist_server(sh, "0.0.0-standalone", &target, allocator, self.zig, self.pgo)?;
         }
         Ok(())
     }
@@ -84,6 +93,7 @@ fn dist_server(
     target: &Target,
     allocator: Malloc,
     zig: bool,
+    pgo: bool,
 ) -> anyhow::Result<()> {
     let _e = sh.push_env("CFG_RELEASE", release);
     let _e = sh.push_env("CARGO_PROFILE_RELEASE_LTO", "thin");
@@ -100,7 +110,22 @@ fn dist_server(
     };
     let features = allocator.to_features();
     let command = if linux_target && zig { "zigbuild" } else { "build" };
-    cmd!(sh, "cargo {command} --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --target {target_name} {features...} --release").run()?;
+
+    let pgo_profile = if pgo {
+        Some(gather_pgo_profile(
+            sh,
+            build_command(sh, command, &target_name, features),
+            &target_name,
+        )?)
+    } else {
+        None
+    };
+
+    let mut cmd = build_command(sh, command, &target_name, features);
+    if let Some(profile) = pgo_profile {
+        cmd = cmd.env("RUSTFLAGS", format!("-Cprofile-use={}", profile.to_str().unwrap()));
+    }
+    cmd.run().context("cannot build Rust Analyzer")?;
 
     let dst = Path::new("dist").join(&target.artifact_name);
     if target_name.contains("-windows-") {
@@ -112,6 +137,70 @@ fn dist_server(
     Ok(())
 }
 
+fn build_command<'a>(
+    sh: &'a Shell,
+    command: &str,
+    target_name: &str,
+    features: &[&str],
+) -> Cmd<'a> {
+    cmd!(
+        sh,
+        "cargo {command} --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --target {target_name} {features...} --release"
+    )
+}
+
+/// Decorates `ra_build_cmd` to add PGO instrumentation, and then runs the PGO instrumented
+/// Rust Analyzer on itself to gather a PGO profile.
+fn gather_pgo_profile<'a>(
+    sh: &'a Shell,
+    ra_build_cmd: Cmd<'a>,
+    target: &str,
+) -> anyhow::Result<PathBuf> {
+    let pgo_dir = std::path::absolute("ra-pgo-profiles")?;
+    // Clear out any stale profiles
+    if pgo_dir.is_dir() {
+        std::fs::remove_dir_all(&pgo_dir)?;
+    }
+    std::fs::create_dir_all(&pgo_dir)?;
+
+    // Figure out a path to `llvm-profdata`
+    let target_libdir = cmd!(sh, "rustc --print=target-libdir")
+        .read()
+        .context("cannot resolve target-libdir from rustc")?;
+    let target_bindir = PathBuf::from(target_libdir).parent().unwrap().join("bin");
+    let llvm_profdata = target_bindir.join(format!("llvm-profdata{}", EXE_EXTENSION));
+
+    // Build RA with PGO instrumentation
+    let cmd_gather =
+        ra_build_cmd.env("RUSTFLAGS", format!("-Cprofile-generate={}", pgo_dir.to_str().unwrap()));
+    cmd_gather.run().context("cannot build rust-analyzer with PGO instrumentation")?;
+
+    // Run RA on itself to gather profiles
+    let train_crate = ".";
+    cmd!(
+        sh,
+        "target/{target}/release/rust-analyzer analysis-stats {train_crate} --run-all-ide-things"
+    )
+    .run()
+    .context("cannot generate PGO profiles")?;
+
+    // Merge profiles into a single file
+    let merged_profile = pgo_dir.join("merged.profdata");
+    let profile_files = std::fs::read_dir(pgo_dir)?.filter_map(|entry| {
+        let entry = entry.ok()?;
+        if entry.path().extension() == Some(OsStr::new("profraw")) {
+            Some(entry.path().to_str().unwrap().to_owned())
+        } else {
+            None
+        }
+    });
+    cmd!(sh, "{llvm_profdata} merge {profile_files...} -o {merged_profile}").run().context(
+        "cannot merge PGO profiles. Do you have the rustup `llvm-tools` component installed?",
+    )?;
+
+    Ok(merged_profile)
+}
+
 fn gzip(src_path: &Path, dest_path: &Path) -> anyhow::Result<()> {
     let mut encoder = GzEncoder::new(File::create(dest_path)?, Compression::best());
     let mut input = io::BufReader::new(File::open(src_path)?);
diff --git a/src/tools/rust-analyzer/xtask/src/flags.rs b/src/tools/rust-analyzer/xtask/src/flags.rs
index d03e2f8437e..81b6a1b7d2a 100644
--- a/src/tools/rust-analyzer/xtask/src/flags.rs
+++ b/src/tools/rust-analyzer/xtask/src/flags.rs
@@ -59,6 +59,8 @@ xflags::xflags! {
             optional --client-patch-version version: String
             /// Use cargo-zigbuild
             optional --zig
+            /// Apply PGO optimizations
+            optional --pgo
         }
         /// Read a changelog AsciiDoc file and update the GitHub Releases entry in Markdown.
         cmd publish-release-notes {
@@ -147,6 +149,7 @@ pub struct Dist {
     pub jemalloc: bool,
     pub client_patch_version: Option<String>,
     pub zig: bool,
+    pub pgo: bool,
 }
 
 #[derive(Debug)]