From 820afa2171ed3f6c5d9832c2b06386cfe010e181 Mon Sep 17 00:00:00 2001
From: Gerald Pinder <gmpinder@gmail.com>
Date: Sat, 11 Jan 2025 15:52:13 -0500
Subject: [PATCH] feat: Add more support for patterns and conditional match

---
 src/cmd.rs   | 45 +++++++++++++++++++++++++++++++++++++--------
 src/lib.rs   | 23 +++++++++++++++++++++++
 tests/cmd.rs | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 111 insertions(+), 8 deletions(-)

diff --git a/src/cmd.rs b/src/cmd.rs
index 9626ad2..9b6fed4 100644
--- a/src/cmd.rs
+++ b/src/cmd.rs
@@ -64,7 +64,10 @@ enum LogicArg {
 impl Parse for LogicArg {
     fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
         if input.peek(Token![for]) {
-            if input.peek3(Token![in]) {
+            let pat_fork = input.fork();
+            _ = pat_fork.parse::<Token![for]>()?;
+
+            if pat_fork.call(syn::Pat::parse_single).is_ok() && pat_fork.peek(Token![in]) {
                 input.parse().map(Self::ForIn)
             } else {
                 input.parse().map(Self::ForIter)
@@ -263,28 +266,54 @@ impl ToTokens for Match {
 
 struct MatchArm {
     pattern: syn::Pat,
+    if_expr: Option<Value>,
     args: Arguments,
 }
 
 impl Parse for MatchArm {
     fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
         let pattern = input.call(syn::Pat::parse_multi)?;
+        let if_expr = if input.peek(Token![if]) {
+            _ = input.parse::<Token![if]>()?;
+            Some(input.parse()?)
+        } else {
+            None
+        };
         _ = input.parse::<Token![=>]>()?;
         let args = input.parse()?;
 
-        Ok(Self { pattern, args })
+        Ok(Self {
+            pattern,
+            if_expr,
+            args,
+        })
     }
 }
 
 impl ToTokens for MatchArm {
     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
-        let Self { pattern, args } = self;
+        let Self {
+            pattern,
+            if_expr,
+            args,
+        } = self;
 
-        tokens.extend(quote! {
-            #pattern => {
-                #args
-            }
-        });
+        tokens.extend(if_expr.as_ref().map_or_else(
+            || {
+                quote! {
+                    #pattern => {
+                        #args
+                    }
+                }
+            },
+            |if_expr| {
+                quote! {
+                    #pattern if #if_expr => {
+                        #args
+                    }
+                }
+            },
+        ));
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index ebc0299..5464b4a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -86,6 +86,29 @@ mod cmd;
 /// assert_eq!(format!("{command:?}"), r#""echo" "test" "arg1" "arg2" "multi" "multi1" "multi" "multi2""#.to_string());
 /// ```
 ///
+/// ## Match Statements
+/// ```
+/// use comlexr::cmd;
+///
+/// enum TestArgs {
+///     Arg1,
+///     Arg2,
+///     Arg3,
+/// }
+///
+/// let match_arg = TestArgs::Arg2;
+/// let command = cmd!(
+///     "echo",
+///     "test",
+///     match match_arg {
+///         TestArgs::Arg1 => "arg1",
+///         TestArgs::Arg2 => ["arg1", "arg2"],
+///         TestArgs::Arg3 => ["arg1", "arg2", "arg3"],
+///     }
+/// );
+/// assert_eq!(format!("{command:?}"), r#""echo" "test" "arg1" "arg2""#.to_string());
+/// ```
+///
 /// ## Dynamic Closures
 /// ```
 /// use comlexr::cmd;
diff --git a/tests/cmd.rs b/tests/cmd.rs
index ae61c99..5551769 100644
--- a/tests/cmd.rs
+++ b/tests/cmd.rs
@@ -62,6 +62,31 @@ fn for_in(#[case] single_iter: &[&str], #[case] multi_iter: &[&str], #[case] exp
     assert_eq!(format!("{command:?}"), expected.to_string());
 }
 
+struct TestStruct(&'static str);
+
+#[rstest]
+#[case(&[], &[], r#""echo" "test""#)]
+#[case(&[TestStruct("1"), TestStruct("2")], &[], r#""echo" "test" "1" "2""#)]
+#[case(&[], &[TestStruct("3"), TestStruct("4")], r#""echo" "test" "multi" "3" "multi" "4""#)]
+#[case(&[TestStruct("1"), TestStruct("2")], &[TestStruct("3"), TestStruct("4")], r#""echo" "test" "1" "2" "multi" "3" "multi" "4""#)]
+fn for_in_pat(
+    #[case] single_iter: &[TestStruct],
+    #[case] multi_iter: &[TestStruct],
+    #[case] expected: &str,
+) {
+    let command = cmd!(
+        "echo",
+        "test",
+        for TestStruct(arg) in single_iter => arg,
+        for TestStruct(arg) in multi_iter => [
+            "multi",
+            arg,
+        ],
+    );
+
+    assert_eq!(format!("{command:?}"), expected.to_string());
+}
+
 #[rstest]
 #[case(None, None, r#""echo" "test""#)]
 #[case(Some("arg"), None, r#""echo" "test" "arg""#)]
@@ -105,6 +130,32 @@ fn match_statement(#[case] match_arg: TestArgs, #[case] expected: &str) {
     assert_eq!(format!("{command:?}"), expected.to_string());
 }
 
+#[rstest]
+#[case(TestArgs::Arg1, true, r#""echo" "test" "arg1""#)]
+#[case(TestArgs::Arg1, false, r#""echo" "test""#)]
+#[case(TestArgs::Arg2, true, r#""echo" "test" "arg1" "arg2""#)]
+#[case(TestArgs::Arg2, false, r#""echo" "test""#)]
+#[case(TestArgs::Arg3, true, r#""echo" "test" "arg1" "arg2" "arg3""#)]
+#[case(TestArgs::Arg3, false, r#""echo" "test""#)]
+fn match_statement_conditional(
+    #[case] match_arg: TestArgs,
+    #[case] flag: bool,
+    #[case] expected: &str,
+) {
+    let command = cmd!(
+        "echo",
+        "test",
+        match match_arg {
+            TestArgs::Arg1 if flag => "arg1",
+            TestArgs::Arg2 if flag => ["arg1", "arg2"],
+            TestArgs::Arg3 if flag => ["arg1", "arg2", "arg3"],
+            _ => [],
+        }
+    );
+
+    assert_eq!(format!("{command:?}"), expected.to_string());
+}
+
 #[rstest]
 #[case(None, None, r#""echo" "test""#)]
 #[case(Some("arg"), None, r#""echo" "test" "arg""#)]