fix: replace Self in generics

1 files changed, 106 insertions, 10 deletions

diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index bcb4449..13037e4 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -10,6 +10,8 @@ use syn::{
     Block,
     FnArg,
     GenericArgument,
+    GenericParam,
+    Generics,
     ImplItem,
     ImplItemMethod,
     ItemTrait,
@@ -22,6 +24,7 @@ use syn::{
     TypeBareFn,
     TypeParamBound,
     Visibility,
+    WherePredicate,
 };
 
 use crate::expectation::Expectation;
@@ -174,6 +177,13 @@ fn get_type_replaced_impl_item_methods(
     mock_ident: &Ident,
 ) -> Vec<ImplItemMethod>
 {
+    let target_path = create_path!(Self);
+
+    let replacement_path = Path::new(
+        WithLeadingColons::No,
+        [PathSegment::new(mock_ident.clone(), None)],
+    );
+
     impl_items
         .into_iter()
         .filter_map(|item| match item {
@@ -186,11 +196,8 @@ fn get_type_replaced_impl_item_methods(
                         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)],
-                                ),
+                                &target_path,
+                                &replacement_path,
                             ));
 
                             FnArg::Typed(typed_arg)
@@ -205,16 +212,19 @@ fn get_type_replaced_impl_item_methods(
                         r_arrow,
                         Box::new(replace_path_in_type(
                             *return_type,
-                            &create_path!(Self),
-                            &Path::new(
-                                WithLeadingColons::No,
-                                [PathSegment::new(mock_ident.clone(), None)],
-                            ),
+                            &target_path,
+                            &replacement_path,
                         )),
                     ),
                     ReturnType::Default => ReturnType::Default,
                 };
 
+                item_method.sig.generics = replace_path_in_generics(
+                    item_method.sig.generics,
+                    &target_path,
+                    &replacement_path,
+                );
+
                 Some(item_method)
             }
             _ => None,
@@ -222,6 +232,71 @@ fn get_type_replaced_impl_item_methods(
         .collect()
 }
 
+fn replace_path_in_generics(
+    mut generics: Generics,
+    target_path: &Path,
+    replacement_path: &Path,
+) -> Generics
+{
+    generics.params = generics
+        .params
+        .into_iter()
+        .map(|generic_param| match generic_param {
+            GenericParam::Type(mut type_param) => {
+                type_param.bounds = type_param
+                    .bounds
+                    .into_iter()
+                    .map(|bound| {
+                        replace_type_param_bound_paths(
+                            bound,
+                            target_path,
+                            replacement_path,
+                        )
+                    })
+                    .collect();
+
+                GenericParam::Type(type_param)
+            }
+            generic_param => generic_param,
+        })
+        .collect();
+
+    generics.where_clause = generics.where_clause.map(|mut where_clause| {
+        where_clause.predicates = where_clause
+            .predicates
+            .into_iter()
+            .map(|predicate| match predicate {
+                WherePredicate::Type(mut predicate_type) => {
+                    predicate_type.bounded_ty = replace_path_in_type(
+                        predicate_type.bounded_ty,
+                        target_path,
+                        replacement_path,
+                    );
+
+                    predicate_type.bounds = predicate_type
+                        .bounds
+                        .into_iter()
+                        .map(|bound| {
+                            replace_type_param_bound_paths(
+                                bound,
+                                target_path,
+                                replacement_path,
+                            )
+                        })
+                        .collect();
+
+                    WherePredicate::Type(predicate_type)
+                }
+                predicate => predicate,
+            })
+            .collect();
+
+        where_clause
+    });
+
+    generics
+}
+
 fn replace_path_in_type(ty: Type, target_path: &Path, replacement_path: &Path) -> Type
 {
     match ty {
@@ -426,3 +501,24 @@ fn replace_type_bare_fn_type_paths(
 
     type_bare_fn
 }
+
+fn replace_type_param_bound_paths(
+    type_param_bound: TypeParamBound,
+    target_path: &Path,
+    replacement_path: &Path,
+) -> TypeParamBound
+{
+    match type_param_bound {
+        TypeParamBound::Trait(mut trait_bound) => {
+            if &trait_bound.path == target_path {
+                trait_bound.path = replacement_path.clone();
+            } else {
+                trait_bound.path =
+                    replace_path_args(trait_bound.path, target_path, replacement_path);
+            }
+
+            TypeParamBound::Trait(trait_bound)
+        }
+        TypeParamBound::Lifetime(lifetime) => TypeParamBound::Lifetime(lifetime),
+    }
+}