refactor: rewrite the mock macro as a procedural macro

1 files changed, 313 insertions, 0 deletions

diff --git a/macros/src/mock.rs b/macros/src/mock.rs
new file mode 100644
index 0000000..88d3434
--- /dev/null
+++ b/macros/src/mock.rs
@@ -0,0 +1,313 @@
+use proc_macro2::{Ident, Span, TokenStream};
+use proc_macro_error::{abort_call_site, ResultExt};
+use quote::{format_ident, quote, ToTokens};
+use syn::{
+    parse2,
+    AngleBracketedGenericArguments,
+    FnArg,
+    GenericArgument,
+    GenericParam,
+    ImplItemMethod,
+    Lifetime,
+    Pat,
+    Path,
+    PathSegment,
+    Receiver,
+    ReturnType,
+    Signature,
+    Token,
+    TraitItemMethod,
+    Type,
+    TypePath,
+    TypeReference,
+    Visibility,
+};
+
+use crate::expectation::create_expectation_ident;
+use crate::syn_ext::{
+    AngleBracketedGenericArgumentsExt,
+    GenericArgumentExt,
+    IsMut,
+    PathExt,
+    PathSegmentExt,
+    ReturnTypeExt,
+    SignatureExt,
+    TypePathExt,
+    TypeReferenceExt,
+    VisibilityExt,
+    WithColons,
+    WithLeadingColons,
+};
+use crate::util::create_unit_type_tuple;
+
+pub struct Mock
+{
+    ident: Ident,
+    mocked_trait: TypePath,
+    expectations_fields: Vec<ExpectationsField>,
+    item_methods: Vec<TraitItemMethod>,
+}
+
+impl Mock
+{
+    pub fn new(
+        ident: Ident,
+        mocked_trait: TypePath,
+        item_methods: &[TraitItemMethod],
+    ) -> Self
+    {
+        let expectations_fields = item_methods
+            .iter()
+            .map(|method_item| {
+                let generic_args = method_item
+                    .sig
+                    .generics
+                    .params
+                    .iter()
+                    .filter_map(|generic_param| match generic_param {
+                        GenericParam::Type(_) => Some(GenericArgument::Type(
+                            Type::Tuple(create_unit_type_tuple()),
+                        )),
+                        GenericParam::Lifetime(_) => {
+                            Some(GenericArgument::Lifetime(Lifetime {
+                                apostrophe: Span::call_site(),
+                                ident: format_ident!("static"),
+                            }))
+                        }
+                        GenericParam::Const(_) => None,
+                    })
+                    .collect::<Vec<_>>();
+
+                ExpectationsField {
+                    field_ident: format_ident!("{}_expectations", method_item.sig.ident),
+                    expectation_ident: create_expectation_ident(
+                        &ident,
+                        &method_item.sig.ident,
+                    ),
+                    generic_args,
+                }
+            })
+            .collect::<Vec<_>>();
+
+        Self {
+            ident,
+            mocked_trait,
+            expectations_fields,
+            item_methods: item_methods.to_vec(),
+        }
+    }
+}
+
+impl ToTokens for Mock
+{
+    fn to_tokens(&self, tokens: &mut TokenStream)
+    {
+        let Self {
+            ident,
+            mocked_trait,
+            expectations_fields,
+            item_methods,
+        } = self;
+
+        let expectations_field_idents = expectations_fields
+            .iter()
+            .map(|expectations_field| expectations_field.field_ident.clone());
+
+        let mock_functions = item_methods
+            .iter()
+            .map(|item_method| create_mock_function(item_method.clone()));
+
+        let expect_functions = item_methods
+            .iter()
+            .map(|item_method| create_expect_function(&self.ident, &item_method.clone()))
+            .collect::<Vec<_>>();
+
+        quote! {
+            pub struct #ident
+            {
+                #(#expectations_fields),*
+            }
+
+            impl #ident
+            {
+                pub fn new() -> Self
+                {
+                    Self {
+                        #(
+                            #expectations_field_idents: ::std::collections::HashMap::new()
+                        ),*
+                    }
+                }
+
+                #(#expect_functions)*
+            }
+
+            impl #mocked_trait for #ident {
+                #(
+                    #mock_functions
+                )*
+            }
+        }
+        .to_tokens(tokens);
+    }
+}
+
+struct ExpectationsField
+{
+    field_ident: Ident,
+    expectation_ident: Ident,
+    generic_args: Vec<GenericArgument>,
+}
+
+impl ToTokens for ExpectationsField
+{
+    fn to_tokens(&self, tokens: &mut TokenStream)
+    {
+        let Self {
+            field_ident,
+            expectation_ident,
+            generic_args,
+        } = self;
+
+        quote! {
+            #field_ident: ::std::collections::HashMap<
+                Vec<::ridicule::__private::type_id::TypeID>,
+                #expectation_ident<#(#generic_args),*>
+            >
+        }
+        .to_tokens(tokens);
+    }
+}
+
+fn create_mock_function(item_method: TraitItemMethod) -> ImplItemMethod
+{
+    let func_ident = &item_method.sig.ident;
+
+    let type_param_idents = item_method
+        .sig
+        .generics
+        .type_params()
+        .map(|type_param| type_param.ident.clone())
+        .collect::<Vec<_>>();
+
+    let args = item_method
+        .sig
+        .inputs
+        .iter()
+        .map(|fn_arg| match fn_arg {
+            FnArg::Receiver(_) => format_ident!("self"),
+            FnArg::Typed(pat_type) => {
+                let Pat::Ident(pat_ident) = pat_type.pat.as_ref() else {
+                    abort_call_site!("Unsupport argument pattern");
+                };
+
+                pat_ident.ident.clone()
+            }
+        })
+        .collect::<Vec<_>>();
+
+    let expectations_field = format_ident!("{func_ident}_expectations");
+
+    let ids = quote! {
+        let ids = vec![
+            #(::ridicule::__private::type_id::TypeID::of::<#type_param_idents>()),*
+        ];
+    };
+
+    ImplItemMethod {
+        attrs: item_method.attrs,
+        vis: Visibility::Inherited,
+        defaultness: None,
+        sig: item_method.sig.clone(),
+        block: parse2(quote! {
+            {
+                #ids
+
+                let expectation = self
+                    .#expectations_field
+                    .get(&ids)
+                    .expect(concat!(
+                        "No expectation found for function ",
+                        stringify!(#func_ident)
+                    ))
+                    .with_generic_params::<#(#type_param_idents),*>();
+
+                let Some(returning) = &expectation.returning else {
+                    panic!(concat!(
+                        "Expectation for function",
+                        stringify!(#func_ident),
+                        " is missing a function to call")
+                    );
+                };
+
+                returning(#(#args),*)
+            }
+        })
+        .unwrap_or_abort(),
+    }
+}
+
+fn create_expect_function(mock: &Ident, item_method: &TraitItemMethod) -> ImplItemMethod
+{
+    let signature = Signature::new(
+        format_ident!("expect_{}", item_method.sig.ident),
+        item_method.sig.generics.clone(),
+        [FnArg::Receiver(Receiver {
+            attrs: vec![],
+            reference: Some((<Token![&]>::default(), None)),
+            mutability: Some(<Token![mut]>::default()),
+            self_token: <Token![self]>::default(),
+        })],
+        ReturnType::new(Type::Reference(TypeReference::new(
+            None,
+            IsMut::Yes,
+            Type::Path(TypePath::new(Path::new(
+                WithLeadingColons::No,
+                [PathSegment::new(
+                    create_expectation_ident(mock, &item_method.sig.ident),
+                    Some(AngleBracketedGenericArguments::new(
+                        WithColons::No,
+                        item_method.sig.generics.params.iter().map(|generic_param| {
+                            GenericArgument::from_generic_param(generic_param.clone())
+                        }),
+                    )),
+                )],
+            ))),
+        ))),
+    );
+
+    let type_param_idents = item_method
+        .sig
+        .generics
+        .type_params()
+        .map(|type_param| type_param.ident.clone())
+        .collect::<Vec<_>>();
+
+    let expectation = create_expectation_ident(mock, &item_method.sig.ident);
+
+    let expectations_field = format_ident!("{}_expectations", item_method.sig.ident);
+
+    ImplItemMethod {
+        attrs: item_method.attrs.clone(),
+        vis: Visibility::new_pub_crate(),
+        defaultness: None,
+        sig: signature,
+        block: parse2(quote! {{
+            let ids = vec![
+                #(::ridicule::__private::type_id::TypeID::of::<#type_param_idents>()),*
+            ];
+
+            let expectation =
+                #expectation::<#(#type_param_idents),*>::new()
+                    .strip_generic_params();
+
+            self.#expectations_field.insert(ids.clone(), expectation);
+
+            self.#expectations_field
+                .get_mut(&ids)
+                .unwrap()
+                .with_generic_params_mut()
+        }})
+        .unwrap_or_abort(),
+    }
+}