Respect --target-dir

Kobzol · Dec 31, 2023 · 819ef38 · 819ef38
1 parent aadc2ec
commit 819ef38
Show file tree

Hide file tree

Showing 9 changed files with 151 additions and 13 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+# Dev
+## Fixes
+- Respect the `--target-dir` cargo argument and improve parsing of cargo arguments in general.
+
 # 0.2.4 (16. 1. 2023)
 ## Fixes
 - Do not close `stdin` when executing `cargo pgo run` ([#28](https://github.com/Kobzol/cargo-pgo/issues/28)).

diff --git a/src/bolt/instrument.rs b/src/bolt/instrument.rs
@@ -34,6 +34,12 @@ pub struct BoltInstrumentArgs {
     cargo_args: Vec<String>,
 }
 
+impl BoltInstrumentArgs {
+    pub fn cargo_args(&self) -> &[String] {
+        &self.cargo_args
+    }
+}
+
 pub fn bolt_instrument(ctx: CargoContext, args: BoltInstrumentArgs) -> anyhow::Result<()> {
     let bolt_dir = ctx.get_bolt_directory()?;
     let bolt_env = find_bolt_env()?;

diff --git a/src/bolt/optimize.rs b/src/bolt/optimize.rs
@@ -32,6 +32,12 @@ pub struct BoltOptimizeArgs {
     cargo_args: Vec<String>,
 }
 
+impl BoltOptimizeArgs {
+    pub fn cargo_args(&self) -> &[String] {
+        &self.cargo_args
+    }
+}
+
 pub fn bolt_optimize(ctx: CargoContext, args: BoltOptimizeArgs) -> anyhow::Result<()> {
     let bolt_dir = ctx.get_bolt_directory()?;
     let bolt_env = find_bolt_env()?;

diff --git a/src/build.rs b/src/build.rs
@@ -3,12 +3,14 @@ use cargo_metadata::{Artifact, Message, MessageIter};
 use std::collections::HashMap;
 use std::fmt::Write as WriteFmt;
 use std::io::{BufReader, Write};
+use std::path::PathBuf;
 use std::process::{Child, ChildStdout, Command, Stdio};
 
 #[derive(Debug, Default)]
-struct CargoArgs {
-    filtered: Vec<String>,
-    contains_target: bool,
+pub struct CargoArgs {
+    pub filtered: Vec<String>,
+    pub contains_target: bool,
+    pub target_dir: Option<PathBuf>,
 }
 
 enum ReleaseMode {
@@ -130,7 +132,7 @@ fn cargo_command(
     Ok(command.spawn()?)
 }
 
-fn parse_cargo_args(cargo_args: Vec<String>) -> CargoArgs {
+pub fn parse_cargo_args(cargo_args: Vec<String>) -> CargoArgs {
     let mut args = CargoArgs::default();
 
     let mut iterator = cargo_args.into_iter();
@@ -151,6 +153,15 @@ fn parse_cargo_args(cargo_args: Vec<String>) -> CargoArgs {
                     if let Some(value) = value {
                         args.filtered.push(value);
                     }
+                } else if let Some(value) =
+                    get_key_value("--target-dir", arg.as_str(), &mut iterator)
+                {
+                    // Extract `--target-dir`
+                    args.target_dir = value.clone().map(PathBuf::from);
+                    args.filtered.push("--target-dir".to_string());
+                    if let Some(value) = value {
+                        args.filtered.push(value);
+                    }
                 } else {
                     args.filtered.push(arg);
                 }
@@ -237,6 +248,7 @@ pub fn get_artifact_kind(artifact: &Artifact) -> &str {
 #[cfg(test)]
 mod tests {
     use crate::build::{get_key_value, parse_cargo_args};
+    use std::path::PathBuf;
 
     #[test]
     fn parse_cargo_args_filter_release() {
@@ -293,6 +305,38 @@ mod tests {
         assert!(args.contains_target);
     }
 
+    #[test]
+    fn parse_cargo_args_target_dir() {
+        let args = parse_cargo_args(vec![
+            "--target-dir".to_string(),
+            "/tmp/foo".to_string(),
+            "bar".to_string(),
+        ]);
+        assert_eq!(
+            args.filtered,
+            vec![
+                "--target-dir".to_string(),
+                "/tmp/foo".to_string(),
+                "bar".to_string()
+            ]
+        );
+        assert_eq!(args.target_dir, Some(PathBuf::from("/tmp/foo")));
+    }
+
+    #[test]
+    fn parse_cargo_args_target_dir_equals() {
+        let args = parse_cargo_args(vec!["--target-dir=/tmp/foo".to_string(), "bar".to_string()]);
+        assert_eq!(
+            args.filtered,
+            vec![
+                "--target-dir".to_string(),
+                "/tmp/foo".to_string(),
+                "bar".to_string()
+            ]
+        );
+        assert_eq!(args.target_dir, Some(PathBuf::from("/tmp/foo")));
+    }
+
     #[test]
     fn get_key_value_wrong_key() {
         assert_eq!(

diff --git a/src/main.rs b/src/main.rs
@@ -58,10 +58,38 @@ enum BoltArgs {
     Optimize(BoltOptimizeArgs),
 }
 
+impl BoltArgs {
+    fn cargo_args(&self) -> &[String] {
+        match self {
+            BoltArgs::Build(args) => args.cargo_args(),
+            BoltArgs::Optimize(args) => args.cargo_args(),
+        }
+    }
+}
+
+impl Args {
+    fn cargo_args(&self) -> &[String] {
+        match self {
+            Args::Pgo(args) => match args {
+                Subcommand::Info => &[],
+                Subcommand::Instrument(args) => args.cargo_args(),
+                Subcommand::Build(args)
+                | Subcommand::Run(args)
+                | Subcommand::Test(args)
+                | Subcommand::Bench(args) => args.cargo_args(),
+                Subcommand::Optimize(args) => args.cargo_args(),
+                Subcommand::Bolt(args) => args.cargo_args(),
+                Subcommand::Clean => &[],
+            },
+        }
+    }
+}
+
 fn run() -> anyhow::Result<()> {
     let args = Args::parse();
 
-    let ctx = get_cargo_ctx()?;
+    let cargo_args = args.cargo_args();
+    let ctx = get_cargo_ctx(cargo_args)?;
 
     let Args::Pgo(args) = args;
     match args {

diff --git a/src/pgo/instrument.rs b/src/pgo/instrument.rs
@@ -22,6 +22,12 @@ pub struct PgoInstrumentArgs {
     cargo_args: Vec<String>,
 }
 
+impl PgoInstrumentArgs {
+    pub fn cargo_args(&self) -> &[String] {
+        &self.cargo_args
+    }
+}
+
 #[derive(clap::Parser, Debug)]
 #[clap(trailing_var_arg(true))]
 pub struct PgoInstrumentShortcutArgs {
@@ -33,6 +39,12 @@ pub struct PgoInstrumentShortcutArgs {
     cargo_args: Vec<String>,
 }
 
+impl PgoInstrumentShortcutArgs {
+    pub fn cargo_args(&self) -> &[String] {
+        &self.cargo_args
+    }
+}
+
 impl PgoInstrumentShortcutArgs {
     pub fn into_full_args(self, command: CargoCommand) -> PgoInstrumentArgs {
         let PgoInstrumentShortcutArgs {

diff --git a/src/pgo/optimize.rs b/src/pgo/optimize.rs
@@ -30,6 +30,12 @@ pub struct PgoOptimizeArgs {
     cargo_args: Vec<String>,
 }
 
+impl PgoOptimizeArgs {
+    pub fn cargo_args(&self) -> &[String] {
+        &self.cargo_args
+    }
+}
+
 /// Merges PGO profiles and creates RUSTFLAGS that use them.
 pub fn prepare_pgo_optimization_flags(pgo_env: &PgoEnv, pgo_dir: &Path) -> anyhow::Result<String> {
     let stats = gather_pgo_profile_stats(pgo_dir)?;

diff --git a/src/workspace.rs b/src/workspace.rs
@@ -1,3 +1,4 @@
+use crate::build::parse_cargo_args;
 use crate::ensure_directory;
 use std::path::{Path, PathBuf};
 
@@ -22,12 +23,18 @@ impl CargoContext {
 }
 
 /// Finds Cargo metadata from the current directory.
-pub fn get_cargo_ctx() -> anyhow::Result<CargoContext> {
-    let cmd = cargo_metadata::MetadataCommand::new();
-    let metadata = cmd
-        .exec()
-        .map_err(|error| anyhow::anyhow!("Cannot get cargo metadata: {:?}", error))?;
-    Ok(CargoContext {
-        target_directory: metadata.target_directory.into_std_path_buf(),
-    })
+pub fn get_cargo_ctx(cargo_args: &[String]) -> anyhow::Result<CargoContext> {
+    let cargo_args = parse_cargo_args(cargo_args.to_vec());
+    let target_directory = match cargo_args.target_dir {
+        Some(dir) => dir,
+        None => {
+            let cmd = cargo_metadata::MetadataCommand::new();
+            let metadata = cmd
+                .exec()
+                .map_err(|error| anyhow::anyhow!("Cannot get cargo metadata: {:?}", error))?;
+            metadata.target_directory.into_std_path_buf()
+        }
+    };
+
+    Ok(CargoContext { target_directory })
 }
diff --git a/tests/integration/pgo.rs b/tests/integration/pgo.rs
@@ -1,4 +1,5 @@
 use crate::utils::{get_dir_files, init_cargo_project, run_command};
+use cargo_pgo::get_default_target;
 
 use crate::utils::OutputExt;
 
@@ -204,3 +205,27 @@ fn main() {
 
     Ok(())
 }
+
+#[test]
+fn test_respect_target_dir() -> anyhow::Result<()> {
+    let project = init_cargo_project()?;
+    let target_dir = project.path("custom-target-dir");
+    let profile_dir = target_dir.join("pgo-profiles");
+
+    project
+        .run(&["build", "--", "--target-dir", target_dir.to_str().unwrap()])?
+        .assert_ok();
+    assert!(!project.default_pgo_profile_dir().is_dir());
+    assert!(profile_dir.is_dir());
+
+    run_command(
+        target_dir
+            .join(get_default_target()?)
+            .join("release")
+            .join("foo"),
+    )?;
+
+    assert!(!get_dir_files(&profile_dir)?.is_empty());
+
+    Ok(())
+}