fix: replace Self with mock in function signatures

2 files changed, 261 insertions, 3 deletions

diff --git a/macros/Cargo.toml b/macros/Cargo.toml
index 818cf8c..6fa4a5f 100644
--- a/macros/Cargo.toml
+++ b/macros/Cargo.toml
@@ -8,7 +8,7 @@ proc-macro = true
 
 [dependencies]
 quote = "1.0.26"
-syn = { version = "1.0.109", features = ["full", "printing"] }
+syn = { version = "1.0.109", features = ["full", "printing", "extra-traits"] }
 proc-macro-error = "1.0.4"
 proc-macro2 = "1.0.52"
 
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index f11c064..7fc062a 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -3,11 +3,25 @@
 use proc_macro::TokenStream;
 use proc_macro_error::{proc_macro_error, ResultExt};
 use quote::{format_ident, quote};
-use syn::{parse, ImplItem};
+use syn::{
+    parse,
+    FnArg,
+    GenericArgument,
+    ImplItem,
+    Path,
+    PathArguments,
+    PathSegment,
+    ReturnType,
+    Type,
+    TypeBareFn,
+    TypeParamBound,
+};
 
 use crate::expectation::Expectation;
 use crate::mock::Mock;
 use crate::mock_input::MockInput;
+use crate::syn_ext::{PathExt, PathSegmentExt, WithLeadingColons};
+use crate::util::create_path;
 
 mod expectation;
 mod mock;
@@ -61,7 +75,46 @@ pub fn mock(input_stream: TokenStream) -> TokenStream
         .items
         .into_iter()
         .filter_map(|item| match item {
-            ImplItem::Method(item_method) => Some(item_method),
+            ImplItem::Method(mut item_method) => {
+                item_method.sig.inputs = item_method
+                    .sig
+                    .inputs
+                    .into_iter()
+                    .map(|fn_arg| match fn_arg {
+                        FnArg::Typed(mut typed_arg) => {
+                            typed_arg.ty = Box::new(replace_path_in_type(
+                                *typed_arg.ty,
+                                &create_path!(Self),
+                                &Path::new(
+                                    WithLeadingColons::No,
+                                    [PathSegment::new(mock_ident.clone(), None)],
+                                ),
+                            ));
+
+                            FnArg::Typed(typed_arg)
+                        }
+
+                        FnArg::Receiver(receiver) => FnArg::Receiver(receiver),
+                    })
+                    .collect();
+
+                item_method.sig.output = match item_method.sig.output {
+                    ReturnType::Type(r_arrow, return_type) => ReturnType::Type(
+                        r_arrow,
+                        Box::new(replace_path_in_type(
+                            *return_type,
+                            &create_path!(Self),
+                            &Path::new(
+                                WithLeadingColons::No,
+                                [PathSegment::new(mock_ident.clone(), None)],
+                            ),
+                        )),
+                    ),
+                    ReturnType::Default => ReturnType::Default,
+                };
+
+                Some(item_method)
+            }
             _ => None,
         })
         .collect::<Vec<_>>();
@@ -94,3 +147,208 @@ pub fn mock(input_stream: TokenStream) -> TokenStream
     }
     .into()
 }
+
+fn replace_path_in_type(ty: Type, target_path: &Path, replacement_path: &Path) -> Type
+{
+    match ty {
+        Type::Ptr(mut type_ptr) => {
+            type_ptr.elem = Box::new(replace_path_in_type(
+                *type_ptr.elem,
+                target_path,
+                replacement_path,
+            ));
+
+            Type::Ptr(type_ptr)
+        }
+        Type::Path(mut type_path) => {
+            if &type_path.path == target_path {
+                type_path.path = replacement_path.clone();
+            } else {
+                type_path.path =
+                    replace_path_args(type_path.path, target_path, replacement_path);
+            }
+
+            Type::Path(type_path)
+        }
+        Type::Array(mut type_array) => {
+            type_array.elem = Box::new(replace_path_in_type(
+                *type_array.elem,
+                target_path,
+                replacement_path,
+            ));
+
+            Type::Array(type_array)
+        }
+        Type::Group(mut type_group) => {
+            type_group.elem = Box::new(replace_path_in_type(
+                *type_group.elem,
+                target_path,
+                replacement_path,
+            ));
+
+            Type::Group(type_group)
+        }
+        Type::BareFn(type_bare_fn) => Type::BareFn(replace_type_bare_fn_type_paths(
+            type_bare_fn,
+            target_path,
+            replacement_path,
+        )),
+        Type::Paren(mut type_paren) => {
+            type_paren.elem = Box::new(replace_path_in_type(
+                *type_paren.elem,
+                target_path,
+                replacement_path,
+            ));
+
+            Type::Paren(type_paren)
+        }
+        Type::Slice(mut type_slice) => {
+            type_slice.elem = Box::new(replace_path_in_type(
+                *type_slice.elem,
+                target_path,
+                replacement_path,
+            ));
+
+            Type::Slice(type_slice)
+        }
+        Type::Tuple(mut type_tuple) => {
+            type_tuple.elems = type_tuple
+                .elems
+                .into_iter()
+                .map(|elem_type| {
+                    replace_path_in_type(elem_type, target_path, replacement_path)
+                })
+                .collect();
+
+            Type::Tuple(type_tuple)
+        }
+        Type::Reference(mut type_reference) => {
+            type_reference.elem = Box::new(replace_path_in_type(
+                *type_reference.elem,
+                target_path,
+                replacement_path,
+            ));
+
+            Type::Reference(type_reference)
+        }
+        Type::TraitObject(mut type_trait_object) => {
+            type_trait_object.bounds = type_trait_object
+                .bounds
+                .into_iter()
+                .map(|bound| match bound {
+                    TypeParamBound::Trait(mut trait_bound) => {
+                        trait_bound.path = replace_path_args(
+                            trait_bound.path,
+                            target_path,
+                            replacement_path,
+                        );
+
+                        TypeParamBound::Trait(trait_bound)
+                    }
+                    TypeParamBound::Lifetime(lifetime) => {
+                        TypeParamBound::Lifetime(lifetime)
+                    }
+                })
+                .collect();
+
+            Type::TraitObject(type_trait_object)
+        }
+        other_type => other_type,
+    }
+}
+
+fn replace_path_args(mut path: Path, target_path: &Path, replacement_path: &Path)
+    -> Path
+{
+    path.segments = path
+        .segments
+        .into_iter()
+        .map(|mut segment| {
+            segment.arguments = match segment.arguments {
+                PathArguments::AngleBracketed(mut generic_args) => {
+                    generic_args.args = generic_args
+                        .args
+                        .into_iter()
+                        .map(|generic_arg| match generic_arg {
+                            GenericArgument::Type(ty) => GenericArgument::Type(
+                                replace_path_in_type(ty, target_path, replacement_path),
+                            ),
+                            GenericArgument::Binding(mut binding) => {
+                                binding.ty = replace_path_in_type(
+                                    binding.ty,
+                                    target_path,
+                                    replacement_path,
+                                );
+
+                                GenericArgument::Binding(binding)
+                            }
+                            generic_arg => generic_arg,
+                        })
+                        .collect();
+
+                    PathArguments::AngleBracketed(generic_args)
+                }
+                PathArguments::Parenthesized(mut generic_args) => {
+                    generic_args.inputs = generic_args
+                        .inputs
+                        .into_iter()
+                        .map(|input_ty| {
+                            replace_path_in_type(input_ty, target_path, replacement_path)
+                        })
+                        .collect();
+
+                    generic_args.output = match generic_args.output {
+                        ReturnType::Type(r_arrow, return_type) => ReturnType::Type(
+                            r_arrow,
+                            Box::new(replace_path_in_type(
+                                *return_type,
+                                target_path,
+                                replacement_path,
+                            )),
+                        ),
+                        ReturnType::Default => ReturnType::Default,
+                    };
+
+                    PathArguments::Parenthesized(generic_args)
+                }
+                PathArguments::None => PathArguments::None,
+            };
+
+            segment
+        })
+        .collect();
+
+    path
+}
+
+fn replace_type_bare_fn_type_paths(
+    mut type_bare_fn: TypeBareFn,
+    target_path: &Path,
+    replacement_path: &Path,
+) -> TypeBareFn
+{
+    type_bare_fn.inputs = type_bare_fn
+        .inputs
+        .into_iter()
+        .map(|mut bare_fn_arg| {
+            bare_fn_arg.ty =
+                replace_path_in_type(bare_fn_arg.ty, target_path, replacement_path);
+
+            bare_fn_arg
+        })
+        .collect();
+
+    type_bare_fn.output = match type_bare_fn.output {
+        ReturnType::Type(r_arrow, return_type) => ReturnType::Type(
+            r_arrow,
+            Box::new(replace_path_in_type(
+                *return_type,
+                target_path,
+                replacement_path,
+            )),
+        ),
+        ReturnType::Default => ReturnType::Default,
+    };
+
+    type_bare_fn
+}