Extend capabilities of TypeFoldable_Generic

Author: Jamesbarford

compiler/rustc_type_ir_macros/src/lib.rs

@@ -17,11 +17,24 @@ decl_derive!(
     [GenericTypeVisitable] => customizable_type_visitable_derive
 );
 
-struct LiftedTy {
+struct TransformedTy {
     ty: syn::Type,
     generic_parameter_bounds: Vec<syn::Ident>,
 }
 
+enum TypeParameterPath {
+    Interner,
+    GenericParameter(syn::Ident),
+}
+
+enum TypeParameterTransform {
+    Continue,
+    Stop,
+}
+
+type TypeParameterVisitor =
+    fn(TypeParameterPath, &mut syn::TypePath, &mut Vec<syn::Ident>) -> TypeParameterTransform;
+
 fn has_ignore_attr(attrs: &[Attribute], name: &'static str, meta: &'static str) -> bool {
     let mut ignored = false;
     attrs.iter().for_each(|attr| {
@@ -91,6 +104,9 @@ fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_macro2::Toke
 
     s.add_where_predicate(parse_quote! { I: Interner });
     s.add_bounds(synstructure::AddBounds::Fields);
+    let generic_parameters =
+        s.ast().generics.type_params().map(|ty| ty.ident.clone()).collect::<Vec<_>>();
+    let mut generic_parameter_bounds = vec![];
     s.bind_with(|_| synstructure::BindStyle::Move);
     let body_try_fold = s.each_variant(|vi| {
         let bindings = vi.bindings();
@@ -101,6 +117,12 @@ fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_macro2::Toke
             if has_ignore_attr(&bind.ast().attrs, "type_foldable", "identity") {
                 bind.to_token_stream()
             } else {
+                for param in
+                    type_foldable_generic_parameters(bind.ast().ty.clone(), &generic_parameters)
+                {
+                    push_unique(&mut generic_parameter_bounds, param);
+                }
+
                 quote! {
                     ::rustc_type_ir::TypeFoldable::try_fold_with(#bind, __folder)?
                 }
@@ -129,6 +151,9 @@ fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_macro2::Toke
     // to generate code for them.
     s.filter(|bi| !has_ignore_attr(&bi.ast().attrs, "type_foldable", "identity"));
     s.add_bounds(synstructure::AddBounds::Fields);
+    for param in generic_parameter_bounds {
+        s.add_where_predicate(parse_quote! { #param: ::rustc_type_ir::TypeFoldable<I> });
+    }
     s.bound_impl(
         quote!(::rustc_type_ir::TypeFoldable<I>),
         quote! {
@@ -149,6 +174,19 @@ fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_macro2::Toke
     )
 }
 
+fn type_foldable_generic_parameters(
+    ty: syn::Type,
+    generic_parameters: &[syn::Ident],
+) -> Vec<syn::Ident> {
+    transform_type_parameters(ty, generic_parameters, |path, _, generic_parameter_bounds| {
+        if let TypeParameterPath::GenericParameter(param) = path {
+            push_unique(generic_parameter_bounds, param);
+        }
+        TypeParameterTransform::Continue
+    })
+    .generic_parameter_bounds
+}
+
 /// `Lift_Generic` is specialised for structs/enums parameterised by an interner
 /// `I: Interner`. It derives `Lift<J>` by rewriting interner associated types
 /// from `I::Assoc` to `J::Assoc`. The required associated type lift bounds are
@@ -251,40 +289,75 @@ fn is_type_phantom(ty: &syn::Type) -> bool {
     get_first_path_segment(ty).is_some_and(|segment| segment.ident == "PhantomData")
 }
 
-fn lift(mut ty: syn::Type, generic_parameters: &[syn::Ident]) -> LiftedTy {
-    struct ItoJ<'a> {
+fn lift(ty: syn::Type, generic_parameters: &[syn::Ident]) -> TransformedTy {
+    transform_type_parameters(ty, generic_parameters, |path, ty, generic_parameter_bounds| {
+        match path {
+            TypeParameterPath::Interner => {
+                *ty.path.segments.first_mut().unwrap() = parse_quote! { J };
+                TypeParameterTransform::Continue
+            }
+            TypeParameterPath::GenericParameter(param) => {
+                push_unique(generic_parameter_bounds, param.clone());
+                *ty = parse_quote! { <#param as ::rustc_type_ir::lift::Lift<J>>::Lifted };
+                TypeParameterTransform::Stop
+            }
+        }
+    })
+}
+
+fn transform_type_parameters(
+    mut ty: syn::Type,
+    generic_parameters: &[syn::Ident],
+    visit: TypeParameterVisitor,
+) -> TransformedTy {
+    struct TypeParameterTransformer<'a> {
         generic_parameters: &'a [syn::Ident],
         generic_parameter_bounds: Vec<syn::Ident>,
+        visit: TypeParameterVisitor,
     }
 
-    impl VisitMut for ItoJ<'_> {
+    impl VisitMut for TypeParameterTransformer<'_> {
         fn visit_type_path_mut(&mut self, i: &mut syn::TypePath) {
-            if i.qself.is_none() {
+            let path = if i.qself.is_none() {
                 let segments_len = i.path.segments.len();
-                if let Some(first) = i.path.segments.first_mut() {
-                    // Turn paths from `I` into `J`
+                i.path.segments.first().and_then(|first| {
                     if first.ident == "I" {
-                        *first = parse_quote! { J };
+                        Some(TypeParameterPath::Interner)
                     } else if segments_len == 1
                         && matches!(first.arguments, syn::PathArguments::None)
                         && self.generic_parameters.iter().any(|param| first.ident == *param)
                     {
-                        let ident = first.ident.clone();
-                        if !self.generic_parameter_bounds.iter().any(|param| *param == ident) {
-                            self.generic_parameter_bounds.push(ident.clone());
-                        }
-
-                        *i = parse_quote! { <#ident as ::rustc_type_ir::lift::Lift<J>>::Lifted };
-                        return;
+                        Some(TypeParameterPath::GenericParameter(first.ident.clone()))
+                    } else {
+                        None
                     }
+                })
+            } else {
+                None
+            };
+
+            if let Some(path) = path {
+                if let TypeParameterTransform::Stop =
+                    (self.visit)(path, i, &mut self.generic_parameter_bounds)
+                {
+                    return;
                 }
             }
+
             syn::visit_mut::visit_type_path_mut(self, i);
         }
     }
-    let mut visitor = ItoJ { generic_parameters, generic_parameter_bounds: Vec::new() };
+
+    let mut visitor =
+        TypeParameterTransformer { generic_parameters, generic_parameter_bounds: vec![], visit };
     visitor.visit_type_mut(&mut ty);
-    LiftedTy { ty, generic_parameter_bounds: visitor.generic_parameter_bounds }
+    TransformedTy { ty, generic_parameter_bounds: visitor.generic_parameter_bounds }
+}
+
+fn push_unique(params: &mut Vec<syn::Ident>, param: syn::Ident) {
+    if !params.iter().any(|prev| *prev == param) {
+        params.push(param);
+    }
 }
 
 #[cfg(not(feature = "nightly"))]