Split TypeFolder and FallibleTypeFolder

Author: compiler-errors

compiler/rustc_macros/src/type_foldable.rs

@@ -14,7 +14,7 @@ pub(super) fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_m
 
     s.add_bounds(synstructure::AddBounds::Generics);
     s.bind_with(|_| synstructure::BindStyle::Move);
-    let body_fold = s.each_variant(|vi| {
+    let try_body_fold = s.each_variant(|vi| {
         let bindings = vi.bindings();
         vi.construct(|_, index| {
             let bind = &bindings[index];
@@ -44,14 +44,51 @@ pub(super) fn type_foldable_derive(mut s: synstructure::Structure<'_>) -> proc_m
         })
     });
 
+    let body_fold = s.each_variant(|vi| {
+        let bindings = vi.bindings();
+        vi.construct(|_, index| {
+            let bind = &bindings[index];
+
+            let mut fixed = false;
+
+            // retain value of fields with #[type_foldable(identity)]
+            bind.ast().attrs.iter().for_each(|x| {
+                if !x.path().is_ident("type_foldable") {
+                    return;
+                }
+                let _ = x.parse_nested_meta(|nested| {
+                    if nested.path.is_ident("identity") {
+                        fixed = true;
+                    }
+                    Ok(())
+                });
+            });
+
+            if fixed {
+                bind.to_token_stream()
+            } else {
+                quote! {
+                    ::rustc_middle::ty::TypeFoldable::fold_with(#bind, __folder)
+                }
+            }
+        })
+    });
+
     s.bound_impl(
         quote!(::rustc_middle::ty::TypeFoldable<::rustc_middle::ty::TyCtxt<'tcx>>),
         quote! {
             fn try_fold_with<__F: ::rustc_middle::ty::FallibleTypeFolder<::rustc_middle::ty::TyCtxt<'tcx>>>(
                 self,
                 __folder: &mut __F
             ) -> Result<Self, __F::Error> {
-                Ok(match self { #body_fold })
+                Ok(match self { #try_body_fold })
+            }
+
+            fn fold_with<__F: ::rustc_middle::ty::TypeFolder<::rustc_middle::ty::TyCtxt<'tcx>>>(
+                self,
+                __folder: &mut __F
+            ) -> Self {
+                match self { #body_fold }
             }
         },
     )

compiler/rustc_middle/src/infer/canonical.rs

@@ -44,6 +44,10 @@ impl<'tcx> ty::TypeFoldable<TyCtxt<'tcx>> for CanonicalVarInfos<'tcx> {
         self,
         folder: &mut F,
     ) -> Result<Self, F::Error> {
+        ty::util::try_fold_list(self, folder, |tcx, v| tcx.mk_canonical_var_infos(v))
+    }
+
+    fn fold_with<F: ty::TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
         ty::util::fold_list(self, folder, |tcx, v| tcx.mk_canonical_var_infos(v))
     }
 }

compiler/rustc_middle/src/ty/generic_args.rs

@@ -16,7 +16,7 @@ use smallvec::SmallVec;
 use crate::ty::codec::{TyDecoder, TyEncoder};
 use crate::ty::{
     self, ClosureArgs, CoroutineArgs, CoroutineClosureArgs, FallibleTypeFolder, InlineConstArgs,
-    Lift, List, Ty, TyCtxt, TypeFoldable, TypeVisitable, TypeVisitor, VisitorResult,
+    Lift, List, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeVisitable, TypeVisitor, VisitorResult,
     walk_visitable_list,
 };
 
@@ -322,6 +322,14 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for GenericArg<'tcx> {
             GenericArgKind::Const(ct) => ct.try_fold_with(folder).map(Into::into),
         }
     }
+
+    fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        match self.unpack() {
+            GenericArgKind::Lifetime(lt) => lt.fold_with(folder).into(),
+            GenericArgKind::Type(ty) => ty.fold_with(folder).into(),
+            GenericArgKind::Const(ct) => ct.fold_with(folder).into(),
+        }
+    }
 }
 
 impl<'tcx> TypeVisitable<TyCtxt<'tcx>> for GenericArg<'tcx> {
@@ -591,6 +599,32 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for GenericArgsRef<'tcx> {
                 }
             }
             0 => Ok(self),
+            _ => ty::util::try_fold_list(self, folder, |tcx, v| tcx.mk_args(v)),
+        }
+    }
+
+    fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        // This code is hot enough that it's worth specializing for the most
+        // common length lists, to avoid the overhead of `SmallVec` creation.
+        // The match arms are in order of frequency. The 1, 2, and 0 cases are
+        // typically hit in 90--99.99% of cases. When folding doesn't change
+        // the args, it's faster to reuse the existing args rather than
+        // calling `mk_args`.
+        match self.len() {
+            1 => {
+                let param0 = self[0].fold_with(folder);
+                if param0 == self[0] { self } else { folder.cx().mk_args(&[param0]) }
+            }
+            2 => {
+                let param0 = self[0].fold_with(folder);
+                let param1 = self[1].fold_with(folder);
+                if param0 == self[0] && param1 == self[1] {
+                    self
+                } else {
+                    folder.cx().mk_args(&[param0, param1])
+                }
+            }
+            0 => self,
             _ => ty::util::fold_list(self, folder, |tcx, v| tcx.mk_args(v)),
         }
     }
@@ -626,6 +660,36 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for &'tcx ty::List<Ty<'tcx>> {
                     Ok(folder.cx().mk_type_list(&[param0, param1]))
                 }
             }
+            _ => ty::util::try_fold_list(self, folder, |tcx, v| tcx.mk_type_list(v)),
+        }
+    }
+
+    fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        // This code is fairly hot, though not as hot as `GenericArgsRef`.
+        //
+        // When compiling stage 2, I get the following results:
+        //
+        // len |   total   |   %
+        // --- | --------- | -----
+        //  2  |  15083590 |  48.1
+        //  3  |   7540067 |  24.0
+        //  1  |   5300377 |  16.9
+        //  4  |   1351897 |   4.3
+        //  0  |   1256849 |   4.0
+        //
+        // I've tried it with some private repositories and got
+        // close to the same result, with 4 and 0 swapping places
+        // sometimes.
+        match self.len() {
+            2 => {
+                let param0 = self[0].fold_with(folder);
+                let param1 = self[1].fold_with(folder);
+                if param0 == self[0] && param1 == self[1] {
+                    self
+                } else {
+                    folder.cx().mk_type_list(&[param0, param1])
+                }
+            }
             _ => ty::util::fold_list(self, folder, |tcx, v| tcx.mk_type_list(v)),
         }
     }

compiler/rustc_middle/src/ty/mod.rs

@@ -539,6 +539,13 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for Term<'tcx> {
             ty::TermKind::Const(ct) => ct.try_fold_with(folder).map(Into::into),
         }
     }
+
+    fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        match self.unpack() {
+            ty::TermKind::Ty(ty) => ty.fold_with(folder).into(),
+            ty::TermKind::Const(ct) => ct.fold_with(folder).into(),
+        }
+    }
 }
 
 impl<'tcx> TypeVisitable<TyCtxt<'tcx>> for Term<'tcx> {