Add PGO support to install

Enable PGO optimizations when installing rust-analyzer with the
--pgo flag. This mirrors functionality already available in dist
command, allowing developers to create optimized local builds.

Example:
cargo xtask install --server --pgo clap-rs/clap

PGO code has been extracted to a dedicated module for reuse.

Author: fbernier

src/tools/rust-analyzer/xtask/src/dist.rs

@@ -1,9 +1,6 @@
 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},
@@ -12,11 +9,11 @@ use time::OffsetDateTime;
 use xshell::{Cmd, Shell, cmd};
 use zip::{DateTime, ZipWriter, write::SimpleFileOptions};
 
-use crate::flags::PgoTrainingCrate;
 use crate::{
     date_iso,
-    flags::{self, Malloc},
+    flags::{self, Malloc, PgoTrainingCrate},
     project_root,
+    util::detect_target,
 };
 
 const VERSION_STABLE: &str = "0.3";
@@ -28,7 +25,7 @@ impl flags::Dist {
         let stable = sh.var("GITHUB_REF").unwrap_or_default().as_str() == "refs/heads/release";
 
         let project_root = project_root();
-        let target = Target::get(&project_root);
+        let target = Target::get(&project_root, sh);
         let allocator = self.allocator();
         let dist = project_root.join("dist");
         sh.remove_path(&dist)?;
@@ -113,9 +110,9 @@ fn dist_server(
     let command = if linux_target && zig { "zigbuild" } else { "build" };
 
     let pgo_profile = if let Some(train_crate) = pgo {
-        Some(gather_pgo_profile(
+        Some(crate::pgo::gather_pgo_profile(
             sh,
-            build_command(sh, command, &target_name, features),
+            crate::pgo::build_command(sh, command, &target_name, features),
             &target_name,
             train_crate,
         )?)
@@ -151,85 +148,6 @@ fn build_command<'a>(
     )
 }
 
-/// 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,
-    train_crate: PgoTrainingCrate,
-) -> anyhow::Result<PathBuf> {
-    let pgo_dir = std::path::absolute("rust-analyzer-pgo")?;
-    // 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("llvm-profdata").with_extension(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")?;
-
-    let (train_path, label) = match &train_crate {
-        PgoTrainingCrate::RustAnalyzer => (PathBuf::from("."), "itself"),
-        PgoTrainingCrate::GitHub(repo) => {
-            (download_crate_for_training(sh, &pgo_dir, repo)?, repo.as_str())
-        }
-    };
-
-    // Run RA either on itself or on a downloaded crate
-    eprintln!("Training RA on {label}...");
-    cmd!(
-        sh,
-        "target/{target}/release/rust-analyzer analysis-stats -q --run-all-ide-things {train_path}"
-    )
-    .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)
-}
-
-/// Downloads a crate from GitHub, stores it into `pgo_dir` and returns a path to it.
-fn download_crate_for_training(sh: &Shell, pgo_dir: &Path, repo: &str) -> anyhow::Result<PathBuf> {
-    let mut it = repo.splitn(2, '@');
-    let repo = it.next().unwrap();
-    let revision = it.next();
-
-    // FIXME: switch to `--revision` here around 2035 or so
-    let revision =
-        if let Some(revision) = revision { &["--branch", revision] as &[&str] } else { &[] };
-
-    let normalized_path = repo.replace("/", "-");
-    let target_path = pgo_dir.join(normalized_path);
-    cmd!(sh, "git clone --depth 1 https://github.com/{repo} {revision...} {target_path}")
-        .run()
-        .with_context(|| "cannot download PGO training crate from {repo}")?;
-
-    Ok(target_path)
-}
-
 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)?);
@@ -283,21 +201,8 @@ struct Target {
 }
 
 impl Target {
-    fn get(project_root: &Path) -> Self {
-        let name = match env::var("RA_TARGET") {
-            Ok(target) => target,
-            _ => {
-                if cfg!(target_os = "linux") {
-                    "x86_64-unknown-linux-gnu".to_owned()
-                } else if cfg!(target_os = "windows") {
-                    "x86_64-pc-windows-msvc".to_owned()
-                } else if cfg!(target_os = "macos") {
-                    "x86_64-apple-darwin".to_owned()
-                } else {
-                    panic!("Unsupported OS, maybe try setting RA_TARGET")
-                }
-            }
-        };
+    fn get(project_root: &Path, sh: &Shell) -> Self {
+        let name = detect_target(sh);
         let (name, libc_suffix) = match name.split_once('.') {
             Some((l, r)) => (l.to_owned(), Some(r.to_owned())),
             None => (name, None),

src/tools/rust-analyzer/xtask/src/flags.rs

@@ -4,6 +4,25 @@ use std::{fmt, str::FromStr};
 
 use crate::install::{ClientOpt, ProcMacroServerOpt, ServerOpt};
 
+#[derive(Debug, Clone)]
+pub enum PgoTrainingCrate {
+    // Use RA's own sources for PGO training
+    RustAnalyzer,
+    // Download a Rust crate from `https://github.com/{0}` and use it for PGO training.
+    GitHub(String),
+}
+
+impl FromStr for PgoTrainingCrate {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "rust-analyzer" => Ok(Self::RustAnalyzer),
+            url => Ok(Self::GitHub(url.to_owned())),
+        }
+    }
+}
+
 xflags::xflags! {
     src "./src/flags.rs"
 
@@ -29,6 +48,9 @@ xflags::xflags! {
 
             /// build in release with debug info set to 2.
             optional --dev-rel
+
+            /// Apply PGO optimizations
+            optional --pgo pgo: PgoTrainingCrate
         }
 
         cmd fuzz-tests {}
@@ -109,18 +131,16 @@ pub enum XtaskCmd {
     Tidy(Tidy),
 }
 
-#[derive(Debug)]
-pub struct Tidy {}
-
 #[derive(Debug)]
 pub struct Install {
     pub client: bool,
     pub code_bin: Option<String>,
     pub server: bool,
-    pub proc_macro_server: bool,
     pub mimalloc: bool,
     pub jemalloc: bool,
+    pub proc_macro_server: bool,
     pub dev_rel: bool,
+    pub pgo: Option<PgoTrainingCrate>,
 }
 
 #[derive(Debug)]
@@ -143,25 +163,6 @@ pub struct RustcPush {
     pub branch: Option<String>,
 }
 
-#[derive(Debug)]
-pub enum PgoTrainingCrate {
-    // Use RA's own sources for PGO training
-    RustAnalyzer,
-    // Download a Rust crate from `https://github.com/{0}` and use it for PGO training.
-    GitHub(String),
-}
-
-impl FromStr for PgoTrainingCrate {
-    type Err = String;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "rust-analyzer" => Ok(Self::RustAnalyzer),
-            url => Ok(Self::GitHub(url.to_owned())),
-        }
-    }
-}
-
 #[derive(Debug)]
 pub struct Dist {
     pub mimalloc: bool,
@@ -195,6 +196,9 @@ pub struct Codegen {
     pub check: bool,
 }
 
+#[derive(Debug)]
+pub struct Tidy;
+
 impl Xtask {
     #[allow(dead_code)]
     pub fn from_env_or_exit() -> Self {
@@ -324,7 +328,7 @@ impl Install {
         } else {
             Malloc::System
         };
-        Some(ServerOpt { malloc, dev_rel: self.dev_rel })
+        Some(ServerOpt { malloc, dev_rel: self.dev_rel, pgo: self.pgo.clone() })
     }
     pub(crate) fn proc_macro_server(&self) -> Option<ProcMacroServerOpt> {
         if !self.proc_macro_server {

src/tools/rust-analyzer/xtask/src/install.rs

@@ -5,7 +5,10 @@ use std::{env, path::PathBuf, str};
 use anyhow::{Context, bail, format_err};
 use xshell::{Shell, cmd};
 
-use crate::flags::{self, Malloc};
+use crate::{
+    flags::{self, Malloc, PgoTrainingCrate},
+    util::detect_target,
+};
 
 impl flags::Install {
     pub(crate) fn run(self, sh: &Shell) -> anyhow::Result<()> {
@@ -35,6 +38,7 @@ const VS_CODES: &[&str] = &["code", "code-exploration", "code-insiders", "codium
 pub(crate) struct ServerOpt {
     pub(crate) malloc: Malloc,
     pub(crate) dev_rel: bool,
+    pub(crate) pgo: Option<PgoTrainingCrate>,
 }
 
 pub(crate) struct ProcMacroServerOpt {
@@ -135,21 +139,33 @@ fn install_server(sh: &Shell, opts: ServerOpt) -> anyhow::Result<()> {
     let features = opts.malloc.to_features();
     let profile = if opts.dev_rel { "dev-rel" } else { "release" };
 
-    let cmd = cmd!(
+    let mut install_cmd = cmd!(
         sh,
         "cargo install --path crates/rust-analyzer --profile={profile} --locked --force --features force-always-assert {features...}"
     );
-    cmd.run()?;
+
+    if let Some(train_crate) = opts.pgo {
+        let build_cmd = cmd!(
+            sh,
+            "cargo build --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --profile={profile} --locked --features force-always-assert {features...}"
+        );
+
+        let target = detect_target(sh);
+        let profile = crate::pgo::gather_pgo_profile(sh, build_cmd, &target, train_crate)?;
+        install_cmd = crate::pgo::apply_pgo_to_cmd(install_cmd, &profile);
+    }
+
+    install_cmd.run()?;
     Ok(())
 }
 
 fn install_proc_macro_server(sh: &Shell, opts: ProcMacroServerOpt) -> anyhow::Result<()> {
     let profile = if opts.dev_rel { "dev-rel" } else { "release" };
 
-    let cmd = cmd!(
+    cmd!(
         sh,
         "cargo +nightly install --path crates/proc-macro-srv-cli --profile={profile} --locked --force --features sysroot-abi"
-    );
-    cmd.run()?;
+    ).run()?;
+
     Ok(())
 }

src/tools/rust-analyzer/xtask/src/main.rs

@@ -22,6 +22,7 @@ mod codegen;
 mod dist;
 mod install;
 mod metrics;
+mod pgo;
 mod publish;
 mod release;
 mod tidy;