From 2e010cd20a3b1cf2586923ee2863dd1be5f5497d Mon Sep 17 00:00:00 2001
From: Benjy Weinberger <benjyw@gmail.com>
Date: Tue, 2 Apr 2024 22:42:23 -0400
Subject: [PATCH] Support passthough args in the rust args parser (#20738)

---
 src/rust/engine/options/src/args.rs       | 37 +++++++++++++++--------
 src/rust/engine/options/src/args_tests.rs | 35 +++++++++++++++++++++
 2 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/src/rust/engine/options/src/args.rs b/src/rust/engine/options/src/args.rs
index faa29f746b8..6457c635efc 100644
--- a/src/rust/engine/options/src/args.rs
+++ b/src/rust/engine/options/src/args.rs
@@ -97,6 +97,7 @@ impl Arg {
 #[derive(Debug)]
 pub struct Args {
     args: Vec<Arg>,
+    passthrough_args: Option<Vec<String>>,
 }
 
 impl Args {
@@ -104,21 +105,22 @@ impl Args {
     // argv[0] process name.
     pub fn new<I: IntoIterator<Item = String>>(arg_strs: I) -> Self {
         let mut args: Vec<Arg> = vec![];
+        let mut passthrough_args: Option<Vec<String>> = None;
         let mut scope = Scope::Global;
-        for arg_str in arg_strs.into_iter() {
-            if arg_str.starts_with("--") {
+        let mut args_iter = arg_strs.into_iter();
+        while let Some(arg_str) = args_iter.next() {
+            if arg_str == "--" {
+                // We've hit the passthrough args delimiter (`--`).
+                passthrough_args = Some(args_iter.collect::<Vec<String>>());
+                break;
+            } else if arg_str.starts_with("--") {
                 let mut components = arg_str.splitn(2, '=');
                 let flag = components.next().unwrap();
-                if flag.is_empty() {
-                    // We've hit the passthrough args delimiter (`--`), so don't look further.
-                    break;
-                } else {
-                    args.push(Arg {
-                        context: scope.clone(),
-                        flag: flag.to_string(),
-                        value: components.next().map(str::to_string),
-                    });
-                }
+                args.push(Arg {
+                    context: scope.clone(),
+                    flag: flag.to_string(),
+                    value: components.next().map(str::to_string),
+                });
             } else if arg_str.starts_with('-') && arg_str.len() >= 2 {
                 let (flag, mut value) = arg_str.split_at(2);
                 // We support -ldebug and -l=debug, so strip that extraneous equals sign.
@@ -139,7 +141,10 @@ impl Args {
             }
         }
 
-        Self { args }
+        Self {
+            args,
+            passthrough_args,
+        }
     }
 
     pub fn argv() -> Self {
@@ -148,6 +153,12 @@ impl Args {
         Self::new(env::args().collect::<Vec<_>>())
     }
 
+    pub fn get_passthrough_args(&self) -> Option<Vec<&str>> {
+        self.passthrough_args
+            .as_ref()
+            .map(|v| Vec::from_iter(v.iter().map(String::as_str)))
+    }
+
     fn get_list<T: Parseable>(&self, id: &OptionId) -> Result<Option<Vec<ListEdit<T>>>, String> {
         let mut edits = vec![];
         for arg in &self.args {
diff --git a/src/rust/engine/options/src/args_tests.rs b/src/rust/engine/options/src/args_tests.rs
index fe81a7c0f53..2dd49237474 100644
--- a/src/rust/engine/options/src/args_tests.rs
+++ b/src/rust/engine/options/src/args_tests.rs
@@ -52,6 +52,7 @@ fn test_string() {
     assert_string("quux", option_id!(["scope"], "quux"));
 
     assert!(args.get_string(&option_id!("dne")).unwrap().is_none());
+    assert!(args.get_passthrough_args().is_none());
 }
 
 #[test]
@@ -84,6 +85,7 @@ fn test_bool() {
     assert_bool(true, option_id!(["scope"], "quuxt"));
 
     assert!(args.get_bool(&option_id!("dne")).unwrap().is_none());
+    assert!(args.get_passthrough_args().is_none());
     assert_eq!(
         "Problem parsing -c bool value:\n1:swallow\n  ^\nExpected 'true' or 'false' at line 1 column 1".to_owned(),
         args.get_bool(&option_id!(-'c', "unladen", "capacity"))
@@ -109,6 +111,7 @@ fn test_float() {
     assert_float(1.137, option_id!("baz", "spam"));
 
     assert!(args.get_float(&option_id!("dne")).unwrap().is_none());
+    assert!(args.get_passthrough_args().is_none());
 
     assert_eq!(
         "Problem parsing --bad float value:\n1:swallow\n  ^\n\
@@ -171,6 +174,7 @@ fn test_string_list() {
     );
 
     assert!(args.get_string_list(&option_id!("dne")).unwrap().is_none());
+    assert!(args.get_passthrough_args().is_none());
 
     let expected_error_msg = "\
 Problem parsing --bad string list value:
@@ -350,3 +354,34 @@ fn test_nonexistent_optional_fromfile() {
     let args = Args::new(vec!["--foo=@?/does/not/exist".to_string()]);
     assert!(args.get_string(&option_id!("foo")).unwrap().is_none());
 }
+
+#[test]
+fn test_passthrough_args() {
+    let args = mk_args([
+        "-ldebug",
+        "--foo=bar",
+        "--",
+        "--passthrough0",
+        "passthrough1",
+        "-p",
+    ]);
+
+    let assert_string = |expected: &str, id: OptionId| {
+        assert_eq!(expected.to_owned(), args.get_string(&id).unwrap().unwrap())
+    };
+
+    assert_string("bar", option_id!("foo"));
+    assert_string("debug", option_id!(-'l', "level"));
+
+    assert_eq!(
+        Some(vec!["--passthrough0", "passthrough1", "-p",]),
+        args.get_passthrough_args()
+    );
+}
+
+#[test]
+fn test_empty_passthrough_args() {
+    let args = mk_args(["-ldebug", "--foo=bar", "--"]);
+
+    assert_eq!(Some(vec![]), args.get_passthrough_args());
+}