From c9ce9aa474d2e2a65777215a8cbf1aeb135965c2 Mon Sep 17 00:00:00 2001
From: Weihang Lo <me@weihanglo.tw>
Date: Tue, 26 Nov 2024 09:42:31 -0500
Subject: [PATCH] test(pgo): ensure PGO works

This is a regression test to prevent issues like #7416.

The test only run on Linux,
as other platforms have different requirements for PGO,
or emit different PGO function missing warnings.
---
 tests/testsuite/main.rs |   1 +
 tests/testsuite/pgo.rs  | 110 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 tests/testsuite/pgo.rs

diff --git a/tests/testsuite/main.rs b/tests/testsuite/main.rs
index 4b52381e3e45..72146c4ab791 100644
--- a/tests/testsuite/main.rs
+++ b/tests/testsuite/main.rs
@@ -136,6 +136,7 @@ mod package_features;
 mod patch;
 mod path;
 mod paths;
+mod pgo;
 mod pkgid;
 mod precise_pre_release;
 mod proc_macro;
diff --git a/tests/testsuite/pgo.rs b/tests/testsuite/pgo.rs
new file mode 100644
index 000000000000..f0441b5889aa
--- /dev/null
+++ b/tests/testsuite/pgo.rs
@@ -0,0 +1,110 @@
+//! Test if PGO works.
+
+use std::path::PathBuf;
+use std::process::Command;
+
+use cargo_test_support::prelude::*;
+use cargo_test_support::project;
+use cargo_test_support::str;
+
+fn llvm_profdata() -> Option<PathBuf> {
+    let output = Command::new("rustc")
+        .arg("--print=target-libdir")
+        .output()
+        .expect("rustc to run");
+    assert!(output.status.success());
+    let mut libdir = PathBuf::from(String::from_utf8(output.stdout).unwrap());
+    assert!(libdir.pop());
+    let mut bin = libdir.join("bin").join("llvm-profdata");
+    bin.exists().then(|| bin.clone()).or_else(|| {
+        bin.set_extension("exe");
+        bin.exists().then_some(bin)
+    })
+}
+
+#[cargo_test]
+fn pgo_works() {
+    if cfg!(not(target_os = "linux")) {
+        // macOS may emit different LLVM PGO warnings.
+        // Windows LLVM has different requirements.
+        return;
+    }
+
+    let Some(llvm_profdata) = llvm_profdata() else {
+        return;
+    };
+
+    let p = project()
+        .file(
+            "Cargo.toml",
+            r#"
+            [package]
+            name = "foo"
+            edition = "2021"
+            "#,
+        )
+        .file(
+            "src/main.rs",
+            r#"
+                fn fibonacci(n: u64) -> u64 {
+                    match n {
+                        0 => 0,
+                        1 => 1,
+                        _ => fibonacci(n - 1) + fibonacci(n - 2),
+                    }
+                }
+
+                fn main() {
+                    for i in [15, 20, 25] {
+                        let _ = fibonacci(i);
+                    }
+                }
+            "#,
+        )
+        .build();
+
+    let target_dir = p.build_dir();
+    let release_bin = target_dir.join("release").join("foo");
+    let pgo_data_dir = target_dir.join("pgo-data");
+    let profdata_path = target_dir.join("merged.profdata");
+
+    // Build the instrumented binary
+    p.cargo("build --release")
+        .env(
+            "RUSTFLAGS",
+            format!("-Cprofile-generate={}", pgo_data_dir.display()),
+        )
+        .run();
+    // Run the instrumented binary
+    cargo_test_support::execs()
+        .with_process_builder(cargo_test_support::process(release_bin))
+        .run();
+
+    cargo_test_support::process(llvm_profdata)
+        .arg("merge")
+        .arg("-o")
+        .arg(&profdata_path)
+        .arg(pgo_data_dir)
+        .status()
+        .unwrap();
+
+    // Use merged profdata during optimization.
+    //
+    // -Cllvm-args=-pgo-warn-missing-function is essential.
+    // If there are LLVM warnings, there might be something wrong.
+    p.cargo("build --release -v")
+        .env(
+            "RUSTFLAGS",
+            format!(
+                "-Cprofile-use={} -Cllvm-args=-pgo-warn-missing-function",
+                profdata_path.display()
+            ),
+        )
+        .with_stderr_data(str![[r#"
+[COMPILING] foo v0.0.0 ([ROOT]/foo)
+[RUNNING] `rustc [..]-Cprofile-use=[ROOT]/foo/target/merged.profdata -Cllvm-args=-pgo-warn-missing-function`
+[FINISHED] `release` profile [optimized] target(s) in [ELAPSED]s
+
+"#]])
+        .run();
+}