feat: add argument matching

2 files changed, 182 insertions, 0 deletions

diff --git a/macros/src/expectation.rs b/macros/src/expectation.rs
index af604ae..8120cfd 100644
--- a/macros/src/expectation.rs
+++ b/macros/src/expectation.rs
@@ -16,13 +16,20 @@ use syn::{
     ImplItemMethod,
     ItemStruct,
     Lifetime,
+    Pat,
+    PatIdent,
+    PatType,
     Path,
     PathSegment,
     Receiver,
     ReturnType,
     Token,
+    TraitBound,
+    TraitBoundModifier,
     Type,
     TypeBareFn,
+    TypeImplTrait,
+    TypeParamBound,
     TypePath,
     TypeReference,
     Visibility,
@@ -174,6 +181,7 @@ impl Expectation
         generics: Generics,
         phantom_fields: &[PhantomField],
         returning_fn: &Type,
+        boxed_predicate_types: &[Type],
     ) -> ItemStruct
     {
         ItemStruct {
@@ -225,6 +233,24 @@ impl Expectation
                     },
                 ]
                 .into_iter()
+                .chain(boxed_predicate_types.iter().enumerate().map(
+                    |(index, boxed_predicate_type)| Field {
+                        attrs: vec![],
+                        vis: Visibility::Inherited,
+                        ident: Some(format_ident!("predicate_{index}")),
+                        colon_token: Some(<Token![:]>::default()),
+                        ty: Type::Path(TypePath::new(Path::new(
+                            WithLeadingColons::No,
+                            [PathSegment::new(
+                                format_ident!("Option"),
+                                Some(AngleBracketedGenericArguments::new(
+                                    WithColons::No,
+                                    [GenericArgument::Type(boxed_predicate_type.clone())],
+                                )),
+                            )],
+                        ))),
+                    },
+                ))
                 .chain(phantom_fields.iter().cloned().map(Field::from))
                 .collect(),
             }),
@@ -267,11 +293,60 @@ impl ToTokens for Expectation
 
         let method_ident = &self.method_ident;
 
+        let arg_types_no_refs = self
+            .arg_types
+            .iter()
+            .map(|arg_type| match arg_type {
+                Type::Reference(type_ref) => &*type_ref.elem,
+                ty => ty,
+            })
+            .collect::<Vec<_>>();
+
+        let predicate_paths = arg_types_no_refs
+            .iter()
+            .map(|arg_type| {
+                Path::new(
+                    WithLeadingColons::Yes,
+                    [
+                        PathSegment::new(format_ident!("ridicule"), None),
+                        PathSegment::new(
+                            format_ident!("Predicate"),
+                            Some(AngleBracketedGenericArguments::new(
+                                WithColons::No,
+                                [GenericArgument::Type((*arg_type).clone())],
+                            )),
+                        ),
+                    ],
+                )
+            })
+            .collect::<Vec<_>>();
+
+        let boxed_predicate_types = arg_types_no_refs
+            .iter()
+            .map(|arg_type| {
+                Type::Path(TypePath::new(Path::new(
+                    WithLeadingColons::Yes,
+                    [
+                        PathSegment::new(format_ident!("ridicule"), None),
+                        PathSegment::new(format_ident!("__private"), None),
+                        PathSegment::new(
+                            format_ident!("BoxPredicate"),
+                            Some(AngleBracketedGenericArguments::new(
+                                WithColons::No,
+                                [GenericArgument::Type((*arg_type).clone())],
+                            )),
+                        ),
+                    ],
+                )))
+            })
+            .collect::<Vec<_>>();
+
         let expectation_struct = Self::create_struct(
             self.ident.clone(),
             generics.clone(),
             phantom_fields,
             &returning_fn,
+            &boxed_predicate_types,
         );
 
         let boundless_generics = generics.clone().strip_where_clause_and_bounds();
@@ -284,6 +359,70 @@ impl ToTokens for Expectation
             quote! { unsafe { std::mem::transmute(self) } }
         };
 
+        let with_arg_names = (0..self.arg_types.len())
+            .map(|index| format_ident!("predicate_{index}"))
+            .collect::<Vec<_>>();
+
+        let with_args =
+            predicate_paths
+                .iter()
+                .enumerate()
+                .map(|(index, predicate_path)| {
+                    FnArg::Typed(PatType {
+                        attrs: vec![],
+                        pat: Box::new(Pat::Ident(PatIdent {
+                            attrs: vec![],
+                            by_ref: None,
+                            mutability: None,
+                            ident: format_ident!("predicate_{index}"),
+                            subpat: None,
+                        })),
+                        colon_token: <Token![:]>::default(),
+                        ty: Box::new(Type::ImplTrait(TypeImplTrait {
+                            impl_token: <Token![impl]>::default(),
+                            bounds: [
+                                TypeParamBound::Trait(TraitBound {
+                                    paren_token: None,
+                                    modifier: TraitBoundModifier::None,
+                                    lifetimes: None,
+                                    path: predicate_path.clone(),
+                                }),
+                                TypeParamBound::Trait(TraitBound {
+                                    paren_token: None,
+                                    modifier: TraitBoundModifier::None,
+                                    lifetimes: None,
+                                    path: create_path!(Send),
+                                }),
+                                TypeParamBound::Trait(TraitBound {
+                                    paren_token: None,
+                                    modifier: TraitBoundModifier::None,
+                                    lifetimes: None,
+                                    path: create_path!(Sync),
+                                }),
+                                TypeParamBound::Lifetime(Lifetime::create(
+                                    format_ident!("static"),
+                                )),
+                            ]
+                            .into_iter()
+                            .collect(),
+                        })),
+                    })
+                });
+
+        let check_predicates_arg_names = (0..self.arg_types.len())
+            .map(|index| format_ident!("arg_{index}"))
+            .collect::<Vec<_>>();
+
+        let arg_types = &self.arg_types;
+
+        let predicate_field_inits = (0..boxed_predicate_types.len())
+            .map(|index| {
+                let ident = format_ident!("predicate_{index}");
+
+                quote! { #ident: None }
+            })
+            .collect::<Vec<_>>();
+
         quote! {
             #expectation_struct
 
@@ -295,6 +434,7 @@ impl ToTokens for Expectation
                         call_cnt: ::std::sync::atomic::AtomicU32::new(0),
                         call_cnt_expectation:
                             ::ridicule::__private::CallCountExpectation::Unlimited,
+                        #(#predicate_field_inits,)*
                         #(#phantom_fields),*
                     }
                 }
@@ -324,6 +464,30 @@ impl ToTokens for Expectation
                     self
                 }
 
+                pub fn with(&mut self, #(#with_args),*) -> &mut Self
+                {
+                    #(
+                        self.#with_arg_names = Some(
+                            ::ridicule::__private::BoxPredicate::new(#with_arg_names)
+                        );
+                    )*
+
+                    self
+                }
+
+                fn check_predicates(&self, #(#check_predicates_arg_names: &#arg_types),*)
+                {
+                    use ::ridicule::Predicate;
+
+                    #(
+                        if let Some(predicate) = &self.#with_arg_names {
+                            if !predicate.eval(&#check_predicates_arg_names) {
+                                panic!("Predicate '{}' evaluated to false", predicate);
+                            }
+                        }
+                    )*
+                }
+
                 #[allow(unused)]
                 fn strip_generic_params(
                     self,
diff --git a/macros/src/mock.rs b/macros/src/mock.rs
index 9aa03f4..4d040a9 100644
--- a/macros/src/mock.rs
+++ b/macros/src/mock.rs
@@ -226,6 +226,22 @@ fn create_mock_function(
         })
         .collect::<Vec<_>>();
 
+    let typed_args = item_method
+        .sig
+        .inputs
+        .iter()
+        .filter_map(|fn_arg| match fn_arg {
+            FnArg::Typed(arg) => {
+                let Pat::Ident(pat_ident) = arg.pat.as_ref() else {
+                    return None;
+                };
+
+                Some(pat_ident.ident.clone())
+            }
+            FnArg::Receiver(_) => None,
+        })
+        .collect::<Vec<_>>();
+
     let expectations_field = format_ident!("{func_ident}_expectations");
 
     ImplItemMethod {
@@ -257,6 +273,8 @@ fn create_mock_function(
                     ))
                     .with_generic_params::<#(#type_param_idents,)*>();
 
+                expectation.check_predicates(#(&#typed_args),*);
+
                 (expectation.get_returning())(#(#args),*)
             }
         })