From 512f20c39cfe97ef06284bb741843e1e825e9368 Mon Sep 17 00:00:00 2001
From: Felipe Machado <462154+felipou@users.noreply.github.com>
Date: Wed, 15 May 2024 00:15:34 -0300
Subject: [PATCH] Add unit-test-kind option to allow choosing bin or lib

When running with the unit-test option, there was no way to choose
between a bin or lib target, so the result was always an error if
there were both types for the same target name. To allow running
unit tests in that case, this new option unit-test-kind may be used
to specify the type of target that we want.
---
 src/bin/cargo-flamegraph.rs | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/bin/cargo-flamegraph.rs b/src/bin/cargo-flamegraph.rs
index f1acdbe..7fd3e5e 100644
--- a/src/bin/cargo-flamegraph.rs
+++ b/src/bin/cargo-flamegraph.rs
@@ -6,6 +6,13 @@ use clap::{Args, Parser};
 
 use flamegraph::Workload;
 
+#[derive(Debug, Clone, Copy, clap::ValueEnum)]
+#[clap(rename_all = "snake_case")]
+enum UnitTestTargetKind {
+    Bin,
+    Lib,
+}
+
 #[derive(Args, Debug)]
 struct Opt {
     /// Build with the dev profile
@@ -42,6 +49,11 @@ struct Opt {
     #[clap(long, group = "exec-args")]
     unit_test: Option<Option<String>>,
 
+    /// Kind of target (lib or bin) when running with <unit-test> which is may be
+    /// required when we have two targets with the same name.
+    #[clap(long)]
+    unit_test_kind: Option<UnitTestTargetKind>,
+
     /// Crate target to unit benchmark, <bench> may be omitted if crate only has one target
     /// (currently profiles the test harness and all tests in the binary; test selection
     /// can be passed as trailing arguments after `--` as separator)
@@ -436,8 +448,14 @@ fn main() -> anyhow::Result<()> {
         opt.package = Some(target.package);
         target.kind
     } else if let Some(unit_test) = opt.unit_test {
+        let kinds = match opt.unit_test_kind {
+            Some(UnitTestTargetKind::Bin) => &["bin"][..], // get slice to help type inference
+            Some(UnitTestTargetKind::Lib) => &["lib"],
+            None => &["bin", "lib"],
+        };
+
         let target = find_unique_target(
-            &["bin", "lib"],
+            kinds,
             opt.package.as_deref(),
             opt.manifest_path.as_deref(),
             unit_test.as_deref(),