Respect --target-dir

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)).

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()?;

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()?;

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 {
@@ -21,6 +23,25 @@ pub struct RunningCargo {
     message_iter: MessageIter<BufReader<ChildStdout>>,
 }
 
+#[derive(Debug, Copy, Clone, clap::ValueEnum)]
+pub enum CargoCommand {
+    Build,
+    Test,
+    Run,
+    Bench,
+}
+
+impl CargoCommand {
+    pub fn to_str(&self) -> &str {
+        match self {
+            CargoCommand::Build => "build",
+            CargoCommand::Test => "test",
+            CargoCommand::Run => "run",
+            CargoCommand::Bench => "bench",
+        }
+    }
+}
+
 impl RunningCargo {
     pub fn messages(&mut self) -> &mut MessageIter<BufReader<ChildStdout>> {
         &mut self.message_iter
@@ -111,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();
@@ -121,21 +142,62 @@ fn parse_cargo_args(cargo_args: Vec<String>) -> CargoArgs {
             "--release" => {
                 log::warn!("Do not pass `--release` manually, it will be added automatically by `cargo-pgo`");
             }
-            // Skip `--message-format`, we need it to be JSON.
-            "--message-format" => {
-                log::warn!("Do not pass `--message-format` manually, it will be added automatically by `cargo-pgo`");
-                iterator.next(); // skip flag value
+            _ => {
+                if get_key_value("--message-format", arg.as_str(), &mut iterator).is_some() {
+                    // Skip `--message-format`, we need it to be JSON.
+                    log::warn!("Do not pass `--message-format` manually, it will be added automatically by `cargo-pgo`");
+                } else if let Some(value) = get_key_value("--target", arg.as_str(), &mut iterator) {
+                    // Check if `--target` was passed
+                    args.contains_target = true;
+                    args.filtered.push("--target".to_string());
+                    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);
+                }
             }
-            "--target" => {
-                args.contains_target = true;
-                args.filtered.push(arg);
-            }
-            _ => args.filtered.push(arg),
         }
     }
     args
 }
 
+/// Parses a `--key=<value>` or `--key <value>` key/value CLI argument pair.
+fn get_key_value<Iter: Iterator<Item = String>>(
+    key: &str,
+    arg: &str,
+    iter: &mut Iter,
+) -> Option<Option<String>> {
+    // A different argument was passed, nothing to be seen here
+    if !arg.starts_with(key) {
+        return None;
+    }
+    // --key was passed exactly, we should extract the value from the following argument
+    if arg == key {
+        let value = iter.next();
+        return Some(value);
+    }
+
+    // --key<suffix> was passed, let's try to split it into --key=value
+    if let Some((parsed_key, value)) = arg.split_once('=') {
+        // if --keyfoo=value was passed, ignore it
+        if parsed_key == key {
+            return Some(Some(value.to_string()));
+        }
+    }
+
+    None
+}
+
 pub fn handle_metadata_message(message: Message) {
     let stdout = std::io::stdout();
     let mut stdout = stdout.lock();
@@ -185,10 +247,11 @@ pub fn get_artifact_kind(artifact: &Artifact) -> &str {
 
 #[cfg(test)]
 mod tests {
-    use crate::build::parse_cargo_args;
+    use crate::build::{get_key_value, parse_cargo_args};
+    use std::path::PathBuf;
 
     #[test]
-    fn test_parse_cargo_args_filter_release() {
+    fn parse_cargo_args_filter_release() {
         let args = parse_cargo_args(vec![
             "foo".to_string(),
             "--release".to_string(),
@@ -198,7 +261,7 @@ mod tests {
     }
 
     #[test]
-    fn test_parse_cargo_args_filter_message_format() {
+    fn parse_cargo_args_filter_message_format() {
         let args = parse_cargo_args(vec![
             "foo".to_string(),
             "--message-format".to_string(),
@@ -209,7 +272,17 @@ mod tests {
     }
 
     #[test]
-    fn test_parse_cargo_args_find_target() {
+    fn parse_cargo_args_filter_message_format_equals() {
+        let args = parse_cargo_args(vec![
+            "foo".to_string(),
+            "--message-format=json".to_string(),
+            "bar".to_string(),
+        ]);
+        assert_eq!(args.filtered, vec!["foo".to_string(), "bar".to_string()]);
+    }
+
+    #[test]
+    fn parse_cargo_args_find_target() {
         let args = parse_cargo_args(vec![
             "--target".to_string(),
             "x64".to_string(),
@@ -221,23 +294,86 @@ mod tests {
         );
         assert!(args.contains_target);
     }
-}
 
-#[derive(Debug, Copy, Clone, clap::ValueEnum)]
-pub enum CargoCommand {
-    Build,
-    Test,
-    Run,
-    Bench,
-}
+    #[test]
+    fn parse_cargo_args_find_target_equals() {
+        let args = parse_cargo_args(vec!["--target=x64".to_string(), "bar".to_string()]);
+        assert_eq!(
+            args.filtered,
+            vec!["--target".to_string(), "x64".to_string(), "bar".to_string()]
+        );
+        assert!(args.contains_target);
+    }
 
-impl CargoCommand {
-    pub fn to_str(&self) -> &str {
-        match self {
-            CargoCommand::Build => "build",
-            CargoCommand::Test => "test",
-            CargoCommand::Run => "run",
-            CargoCommand::Bench => "bench",
-        }
+    #[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!(
+            get_key_value("--foo", "--bar", &mut std::iter::empty()),
+            None
+        );
+    }
+
+    #[test]
+    fn get_key_value_exact_key_missing_value() {
+        assert_eq!(
+            get_key_value("--foo", "--foo", &mut std::iter::empty()),
+            Some(None)
+        );
+    }
+
+    #[test]
+    fn get_key_value_exact_key_value() {
+        assert_eq!(
+            get_key_value("--foo", "--foo", &mut vec!["bar".to_string()].into_iter()),
+            Some(Some("bar".to_string()))
+        );
+    }
+
+    #[test]
+    fn get_key_value_equals_wrong_prefix() {
+        assert_eq!(
+            get_key_value("--foo", "--foox=bar", &mut std::iter::empty()),
+            None
+        );
+    }
+
+    #[test]
+    fn get_key_value_equals() {
+        assert_eq!(
+            get_key_value("--foo", "--foo=bar", &mut std::iter::empty()),
+            Some(Some("bar".to_string()))
+        );
     }
 }

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 {

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 {

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)?;