Auto merge of #107627 - nnethercote:optimize-fold_ty, r=compiler-errors

Optimize `fold_ty`

Micro-optimizing the heck out of the important `fold_ty` methods.

r? `@oli-obk`

3 files changed, 104 insertions, 99 deletions

diff --git a/compiler/rustc_infer/src/infer/freshen.rs b/compiler/rustc_infer/src/infer/freshen.rs
index 83d71edc2ab..2355234637c 100644
--- a/compiler/rustc_infer/src/infer/freshen.rs
+++ b/compiler/rustc_infer/src/infer/freshen.rs
@@ -140,79 +140,21 @@ impl<'a, 'tcx> TypeFolder<'tcx> for TypeFreshener<'a, 'tcx> {
         }
     }
 
+    #[inline]
     fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
         if !t.needs_infer() && !t.has_erasable_regions() {
-            return t;
-        }
-
-        let tcx = self.infcx.tcx;
-
-        match *t.kind() {
-            ty::Infer(ty::TyVar(v)) => {
-                let opt_ty = self.infcx.inner.borrow_mut().type_variables().probe(v).known();
-                self.freshen_ty(opt_ty, ty::TyVar(v), ty::FreshTy)
-            }
+            t
+        } else {
+            match *t.kind() {
+                ty::Infer(v) => self.fold_infer_ty(v).unwrap_or(t),
 
-            ty::Infer(ty::IntVar(v)) => self.freshen_ty(
-                self.infcx
-                    .inner
-                    .borrow_mut()
-                    .int_unification_table()
-                    .probe_value(v)
-                    .map(|v| v.to_type(tcx)),
-                ty::IntVar(v),
-                ty::FreshIntTy,
-            ),
+                // This code is hot enough that a non-debug assertion here makes a noticeable
+                // difference on benchmarks like `wg-grammar`.
+                #[cfg(debug_assertions)]
+                ty::Placeholder(..) | ty::Bound(..) => bug!("unexpected type {:?}", t),
 
-            ty::Infer(ty::FloatVar(v)) => self.freshen_ty(
-                self.infcx
-                    .inner
-                    .borrow_mut()
-                    .float_unification_table()
-                    .probe_value(v)
-                    .map(|v| v.to_type(tcx)),
-                ty::FloatVar(v),
-                ty::FreshFloatTy,
-            ),
-
-            ty::Infer(ty::FreshTy(ct) | ty::FreshIntTy(ct) | ty::FreshFloatTy(ct)) => {
-                if ct >= self.ty_freshen_count {
-                    bug!(
-                        "Encountered a freshend type with id {} \
-                          but our counter is only at {}",
-                        ct,
-                        self.ty_freshen_count
-                    );
-                }
-                t
+                _ => t.super_fold_with(self),
             }
-
-            ty::Generator(..)
-            | ty::Bool
-            | ty::Char
-            | ty::Int(..)
-            | ty::Uint(..)
-            | ty::Float(..)
-            | ty::Adt(..)
-            | ty::Str
-            | ty::Error(_)
-            | ty::Array(..)
-            | ty::Slice(..)
-            | ty::RawPtr(..)
-            | ty::Ref(..)
-            | ty::FnDef(..)
-            | ty::FnPtr(_)
-            | ty::Dynamic(..)
-            | ty::Never
-            | ty::Tuple(..)
-            | ty::Alias(..)
-            | ty::Foreign(..)
-            | ty::Param(..)
-            | ty::Closure(..)
-            | ty::GeneratorWitnessMIR(..)
-            | ty::GeneratorWitness(..) => t.super_fold_with(self),
-
-            ty::Placeholder(..) | ty::Bound(..) => bug!("unexpected type {:?}", t),
         }
     }
 
@@ -253,3 +195,54 @@ impl<'a, 'tcx> TypeFolder<'tcx> for TypeFreshener<'a, 'tcx> {
         }
     }
 }
+
+impl<'a, 'tcx> TypeFreshener<'a, 'tcx> {
+    // This is separate from `fold_ty` to keep that method small and inlinable.
+    #[inline(never)]
+    fn fold_infer_ty(&mut self, v: ty::InferTy) -> Option<Ty<'tcx>> {
+        match v {
+            ty::TyVar(v) => {
+                let opt_ty = self.infcx.inner.borrow_mut().type_variables().probe(v).known();
+                Some(self.freshen_ty(opt_ty, ty::TyVar(v), ty::FreshTy))
+            }
+
+            ty::IntVar(v) => Some(
+                self.freshen_ty(
+                    self.infcx
+                        .inner
+                        .borrow_mut()
+                        .int_unification_table()
+                        .probe_value(v)
+                        .map(|v| v.to_type(self.infcx.tcx)),
+                    ty::IntVar(v),
+                    ty::FreshIntTy,
+                ),
+            ),
+
+            ty::FloatVar(v) => Some(
+                self.freshen_ty(
+                    self.infcx
+                        .inner
+                        .borrow_mut()
+                        .float_unification_table()
+                        .probe_value(v)
+                        .map(|v| v.to_type(self.infcx.tcx)),
+                    ty::FloatVar(v),
+                    ty::FreshFloatTy,
+                ),
+            ),
+
+            ty::FreshTy(ct) | ty::FreshIntTy(ct) | ty::FreshFloatTy(ct) => {
+                if ct >= self.ty_freshen_count {
+                    bug!(
+                        "Encountered a freshend type with id {} \
+                          but our counter is only at {}",
+                        ct,
+                        self.ty_freshen_count
+                    );
+                }
+                None
+            }
+        }
+    }
+}
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index f39170bb291..8e0bcff8d0a 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -30,7 +30,7 @@ use rustc_middle::ty::relate::RelateResult;
 use rustc_middle::ty::subst::{GenericArg, GenericArgKind, InternalSubsts, SubstsRef};
 use rustc_middle::ty::visit::TypeVisitable;
 pub use rustc_middle::ty::IntVarValue;
-use rustc_middle::ty::{self, GenericParamDefKind, InferConst, Ty, TyCtxt};
+use rustc_middle::ty::{self, GenericParamDefKind, InferConst, InferTy, Ty, TyCtxt};
 use rustc_middle::ty::{ConstVid, FloatVid, IntVid, TyVid};
 use rustc_span::symbol::Symbol;
 use rustc_span::Span;
@@ -1389,8 +1389,8 @@ impl<'tcx> InferCtxt<'tcx> {
     where
         T: TypeFoldable<'tcx>,
     {
-        if !value.needs_infer() {
-            return value; // Avoid duplicated subst-folding.
+        if !value.has_non_region_infer() {
+            return value;
         }
         let mut r = resolve::OpportunisticVarResolver::new(self);
         value.fold_with(&mut r)
@@ -1870,9 +1870,33 @@ impl<'a, 'tcx> TypeFolder<'tcx> for ShallowResolver<'a, 'tcx> {
     /// If `ty` is a type variable of some kind, resolve it one level
     /// (but do not resolve types found in the result). If `typ` is
     /// not a type variable, just return it unmodified.
+    #[inline]
     fn fold_ty(&mut self, ty: Ty<'tcx>) -> Ty<'tcx> {
-        match *ty.kind() {
-            ty::Infer(ty::TyVar(v)) => {
+        if let ty::Infer(v) = ty.kind() { self.fold_infer_ty(*v).unwrap_or(ty) } else { ty }
+    }
+
+    fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
+        if let ty::ConstKind::Infer(InferConst::Var(vid)) = ct.kind() {
+            self.infcx
+                .inner
+                .borrow_mut()
+                .const_unification_table()
+                .probe_value(vid)
+                .val
+                .known()
+                .unwrap_or(ct)
+        } else {
+            ct
+        }
+    }
+}
+
+impl<'a, 'tcx> ShallowResolver<'a, 'tcx> {
+    // This is separate from `fold_ty` to keep that method small and inlinable.
+    #[inline(never)]
+    fn fold_infer_ty(&mut self, v: InferTy) -> Option<Ty<'tcx>> {
+        match v {
+            ty::TyVar(v) => {
                 // Not entirely obvious: if `typ` is a type variable,
                 // it can be resolved to an int/float variable, which
                 // can then be recursively resolved, hence the
@@ -1886,41 +1910,26 @@ impl<'a, 'tcx> TypeFolder<'tcx> for ShallowResolver<'a, 'tcx> {
                 // Note: if these two lines are combined into one we get
                 // dynamic borrow errors on `self.inner`.
                 let known = self.infcx.inner.borrow_mut().type_variables().probe(v).known();
-                known.map_or(ty, |t| self.fold_ty(t))
+                known.map(|t| self.fold_ty(t))
             }
 
-            ty::Infer(ty::IntVar(v)) => self
+            ty::IntVar(v) => self
                 .infcx
                 .inner
                 .borrow_mut()
                 .int_unification_table()
                 .probe_value(v)
-                .map_or(ty, |v| v.to_type(self.infcx.tcx)),
+                .map(|v| v.to_type(self.infcx.tcx)),
 
-            ty::Infer(ty::FloatVar(v)) => self
+            ty::FloatVar(v) => self
                 .infcx
                 .inner
                 .borrow_mut()
                 .float_unification_table()
                 .probe_value(v)
-                .map_or(ty, |v| v.to_type(self.infcx.tcx)),
+                .map(|v| v.to_type(self.infcx.tcx)),
 
-            _ => ty,
-        }
-    }
-
-    fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
-        if let ty::ConstKind::Infer(InferConst::Var(vid)) = ct.kind() {
-            self.infcx
-                .inner
-                .borrow_mut()
-                .const_unification_table()
-                .probe_value(vid)
-                .val
-                .known()
-                .unwrap_or(ct)
-        } else {
-            ct
+            ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_) => None,
         }
     }
 }
diff --git a/compiler/rustc_infer/src/infer/resolve.rs b/compiler/rustc_infer/src/infer/resolve.rs
index 65b90aa3d79..a39a40cf9ab 100644
--- a/compiler/rustc_infer/src/infer/resolve.rs
+++ b/compiler/rustc_infer/src/infer/resolve.rs
@@ -16,26 +16,29 @@ use std::ops::ControlFlow;
 /// useful for printing messages etc but also required at various
 /// points for correctness.
 pub struct OpportunisticVarResolver<'a, 'tcx> {
-    infcx: &'a InferCtxt<'tcx>,
+    // The shallow resolver is used to resolve inference variables at every
+    // level of the type.
+    shallow_resolver: crate::infer::ShallowResolver<'a, 'tcx>,
 }
 
 impl<'a, 'tcx> OpportunisticVarResolver<'a, 'tcx> {
     #[inline]
     pub fn new(infcx: &'a InferCtxt<'tcx>) -> Self {
-        OpportunisticVarResolver { infcx }
+        OpportunisticVarResolver { shallow_resolver: crate::infer::ShallowResolver { infcx } }
     }
 }
 
 impl<'a, 'tcx> TypeFolder<'tcx> for OpportunisticVarResolver<'a, 'tcx> {
     fn tcx<'b>(&'b self) -> TyCtxt<'tcx> {
-        self.infcx.tcx
+        TypeFolder::tcx(&self.shallow_resolver)
     }
 
+    #[inline]
     fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
         if !t.has_non_region_infer() {
             t // micro-optimize -- if there is nothing in this type that this fold affects...
         } else {
-            let t = self.infcx.shallow_resolve(t);
+            let t = self.shallow_resolver.fold_ty(t);
             t.super_fold_with(self)
         }
     }
@@ -44,7 +47,7 @@ impl<'a, 'tcx> TypeFolder<'tcx> for OpportunisticVarResolver<'a, 'tcx> {
         if !ct.has_non_region_infer() {
             ct // micro-optimize -- if there is nothing in this const that this fold affects...
         } else {
-            let ct = self.infcx.shallow_resolve(ct);
+            let ct = self.shallow_resolver.fold_const(ct);
             ct.super_fold_with(self)
         }
     }