Merge pull request #19685 from fbernier/pgo-install

Add PGO support to install

6 files changed, 175 insertions, 132 deletions

diff --git a/src/tools/rust-analyzer/xtask/src/dist.rs b/src/tools/rust-analyzer/xtask/src/dist.rs
index b3d6f06b073..dbfecdbe112 100644
--- a/src/tools/rust-analyzer/xtask/src/dist.rs
+++ b/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),
diff --git a/src/tools/rust-analyzer/xtask/src/flags.rs b/src/tools/rust-analyzer/xtask/src/flags.rs
index 700806d178c..2fd471b35c7 100644
--- a/src/tools/rust-analyzer/xtask/src/flags.rs
+++ b/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 {}
@@ -110,17 +132,15 @@ pub enum XtaskCmd {
 }
 
 #[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)]
@@ -144,25 +164,6 @@ pub struct RustcPush {
 }
 
 #[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,
     pub jemalloc: 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 {
diff --git a/src/tools/rust-analyzer/xtask/src/install.rs b/src/tools/rust-analyzer/xtask/src/install.rs
index 4e2093f0691..f0cc445dfa2 100644
--- a/src/tools/rust-analyzer/xtask/src/install.rs
+++ b/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(())
 }
diff --git a/src/tools/rust-analyzer/xtask/src/main.rs b/src/tools/rust-analyzer/xtask/src/main.rs
index 52ea896c734..aaa8d0e1d4d 100644
--- a/src/tools/rust-analyzer/xtask/src/main.rs
+++ b/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;
diff --git a/src/tools/rust-analyzer/xtask/src/pgo.rs b/src/tools/rust-analyzer/xtask/src/pgo.rs
new file mode 100644
index 00000000000..7f7b3311d96
--- /dev/null
+++ b/src/tools/rust-analyzer/xtask/src/pgo.rs
@@ -0,0 +1,105 @@
+//! PGO (Profile-Guided Optimization) utilities.
+
+use anyhow::Context;
+use std::env::consts::EXE_EXTENSION;
+use std::ffi::OsStr;
+use std::path::{Path, PathBuf};
+use xshell::{Cmd, Shell, cmd};
+
+use crate::flags::PgoTrainingCrate;
+
+/// Decorates `ra_build_cmd` to add PGO instrumentation, and then runs the PGO instrumented
+/// Rust Analyzer on itself to gather a PGO profile.
+pub(crate) 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)
+}
+
+/// Helper function to create a build command for rust-analyzer
+pub(crate) 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"
+    )
+}
+
+pub(crate) fn apply_pgo_to_cmd<'a>(cmd: Cmd<'a>, profile_path: &Path) -> Cmd<'a> {
+    cmd.env("RUSTFLAGS", format!("-Cprofile-use={}", profile_path.to_str().unwrap()))
+}
diff --git a/src/tools/rust-analyzer/xtask/src/util.rs b/src/tools/rust-analyzer/xtask/src/util.rs
index 39f52938c8c..a740ad6afdc 100644
--- a/src/tools/rust-analyzer/xtask/src/util.rs
+++ b/src/tools/rust-analyzer/xtask/src/util.rs
@@ -1,5 +1,7 @@
 use std::path::{Path, PathBuf};
 
+use xshell::{Shell, cmd};
+
 pub(crate) fn list_rust_files(dir: &Path) -> Vec<PathBuf> {
     let mut res = list_files(dir);
     res.retain(|it| {
@@ -29,3 +31,13 @@ pub(crate) fn list_files(dir: &Path) -> Vec<PathBuf> {
     }
     res
 }
+
+pub(crate) fn detect_target(sh: &Shell) -> String {
+    match std::env::var("RA_TARGET") {
+        Ok(target) => target,
+        _ => match cmd!(sh, "rustc --print=host-tuple").read() {
+            Ok(target) => target,
+            Err(e) => panic!("Failed to detect target: {}\nPlease set RA_TARGET explicitly", e),
+        },
+    }
+}