Auto merge of #142322 - compiler-errors:perf-instantiate-uwu2, r=<try>

Perf instantiate uwu2

This is horrible

Author: rust-bors[bot]

compiler/rustc_infer/src/infer/canonical/instantiate.rs

@@ -7,8 +7,11 @@
 //! [c]: https://rust-lang.github.io/chalk/book/canonical_queries/canonicalization.html
 
 use rustc_macros::extension;
-use rustc_middle::bug;
-use rustc_middle::ty::{self, FnMutDelegate, GenericArgKind, TyCtxt, TypeFoldable};
+use rustc_middle::ty::{
+    self, DelayedMap, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeSuperVisitable,
+    TypeVisitableExt, TypeVisitor,
+};
+use rustc_type_ir::TypeVisitable;
 
 use crate::infer::canonical::{Canonical, CanonicalVarValues};
 
@@ -18,11 +21,15 @@ use crate::infer::canonical::{Canonical, CanonicalVarValues};
 impl<'tcx, V> Canonical<'tcx, V> {
     /// Instantiate the wrapped value, replacing each canonical value
     /// with the value given in `var_values`.
-    fn instantiate(&self, tcx: TyCtxt<'tcx>, var_values: &CanonicalVarValues<'tcx>) -> V
+    fn instantiate<const UWU: bool>(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        var_values: &CanonicalVarValues<'tcx>,
+    ) -> V
     where
         V: TypeFoldable<TyCtxt<'tcx>>,
     {
-        self.instantiate_projected(tcx, var_values, |value| value.clone())
+        self.instantiate_projected::<_, UWU>(tcx, var_values, |value| value.clone())
     }
 
     /// Allows one to apply a instantiation to some subset of
@@ -31,7 +38,7 @@ impl<'tcx, V> Canonical<'tcx, V> {
     /// variables bound in `self` (usually this extracts from subset
     /// of `self`). Apply the instantiation `var_values` to this value
     /// V, replacing each of the canonical variables.
-    fn instantiate_projected<T>(
+    fn instantiate_projected<T, const UWU: bool>(
         &self,
         tcx: TyCtxt<'tcx>,
         var_values: &CanonicalVarValues<'tcx>,
@@ -42,14 +49,14 @@ impl<'tcx, V> Canonical<'tcx, V> {
     {
         assert_eq!(self.variables.len(), var_values.len());
         let value = projection_fn(&self.value);
-        instantiate_value(tcx, var_values, value)
+        instantiate_value_0::<_, UWU>(tcx, var_values, value)
     }
 }
 
 /// Instantiate the values from `var_values` into `value`. `var_values`
 /// must be values for the set of canonical variables that appear in
 /// `value`.
-pub(super) fn instantiate_value<'tcx, T>(
+pub(super) fn instantiate_value_0<'tcx, T, const UWU: bool>(
     tcx: TyCtxt<'tcx>,
     var_values: &CanonicalVarValues<'tcx>,
     value: T,
@@ -58,23 +65,160 @@ where
     T: TypeFoldable<TyCtxt<'tcx>>,
 {
     if var_values.var_values.is_empty() {
-        value
-    } else {
-        let delegate = FnMutDelegate {
-            regions: &mut |br: ty::BoundRegion| match var_values[br.var].kind() {
-                GenericArgKind::Lifetime(l) => l,
-                r => bug!("{:?} is a region but value is {:?}", br, r),
-            },
-            types: &mut |bound_ty: ty::BoundTy| match var_values[bound_ty.var].kind() {
-                GenericArgKind::Type(ty) => ty,
-                r => bug!("{:?} is a type but value is {:?}", bound_ty, r),
-            },
-            consts: &mut |bound_ct: ty::BoundVar| match var_values[bound_ct].kind() {
-                GenericArgKind::Const(ct) => ct,
-                c => bug!("{:?} is a const but value is {:?}", bound_ct, c),
-            },
-        };
-
-        tcx.replace_escaping_bound_vars_uncached(value, delegate)
+        return value;
+    }
+
+    value.fold_with(&mut BoundVarReplacer::<UWU> {
+        tcx,
+        current_index: ty::INNERMOST,
+        var_values: var_values.var_values,
+        cache: Default::default(),
+    })
+}
+
+/// Replaces the escaping bound vars (late bound regions or bound types) in a type.
+struct BoundVarReplacer<'tcx, const UWU: bool> {
+    tcx: TyCtxt<'tcx>,
+
+    /// As with `RegionFolder`, represents the index of a binder *just outside*
+    /// the ones we have visited.
+    current_index: ty::DebruijnIndex,
+
+    var_values: ty::GenericArgsRef<'tcx>,
+
+    /// This cache only tracks the `DebruijnIndex` and assumes that it does not matter
+    /// for the delegate how often its methods get used.
+    cache: DelayedMap<(ty::DebruijnIndex, Ty<'tcx>), Ty<'tcx>>,
+}
+
+impl<'tcx, const UWU: bool> TypeFolder<TyCtxt<'tcx>> for BoundVarReplacer<'tcx, UWU> {
+    fn cx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn fold_binder<T: TypeFoldable<TyCtxt<'tcx>>>(
+        &mut self,
+        t: ty::Binder<'tcx, T>,
+    ) -> ty::Binder<'tcx, T> {
+        self.current_index.shift_in(1);
+        let t = t.super_fold_with(self);
+        self.current_index.shift_out(1);
+        t
+    }
+
+    fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
+        match *t.kind() {
+            ty::Bound(debruijn, bound_ty) if debruijn == self.current_index => {
+                self.var_values[bound_ty.var.as_usize()].expect_ty()
+            }
+            _ => {
+                if !t.has_vars_bound_at_or_above(self.current_index) {
+                    t
+                } else if let Some(&t) = self.cache.get(&(self.current_index, t)) {
+                    t
+                } else {
+                    let res = t.super_fold_with(self);
+                    assert!(self.cache.insert((self.current_index, t), res));
+                    res
+                }
+            }
+        }
+    }
+
+    fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
+        match r.kind() {
+            ty::ReBound(debruijn, br) if debruijn == self.current_index => {
+                self.var_values[br.var.as_usize()].expect_region()
+            }
+            _ => r,
+        }
+    }
+
+    fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
+        match ct.kind() {
+            ty::ConstKind::Bound(debruijn, bound_const) if debruijn == self.current_index => {
+                self.var_values[bound_const.as_usize()].expect_const()
+            }
+            _ => ct.super_fold_with(self),
+        }
+    }
+
+    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+        if p.has_vars_bound_at_or_above(self.current_index) { p.super_fold_with(self) } else { p }
+    }
+
+    fn fold_clauses(&mut self, c: ty::Clauses<'tcx>) -> ty::Clauses<'tcx> {
+        if c.has_vars_bound_at_or_above(self.current_index) {
+            if self.current_index == ty::INNERMOST {
+                if UWU {
+                    let index = *self
+                        .tcx
+                        .highest_var_in_clauses_cache
+                        .lock()
+                        .entry(c)
+                        .or_insert_with(|| highest_var_in_clauses(c));
+                    let c_args = &self.var_values[..=index];
+                    if let Some(c2) = self.tcx.clauses_cache.lock().get(&(c, c_args)) {
+                        c2
+                    } else {
+                        let folded = c.super_fold_with(self);
+                        self.tcx.clauses_cache.lock().insert((c, c_args), folded);
+                        folded
+                    }
+                } else {
+                    c.super_fold_with(self)
+                }
+            } else {
+                c.super_fold_with(self)
+            }
+        } else {
+            c
+        }
+    }
+}
+
+fn highest_var_in_clauses<'tcx>(c: ty::Clauses<'tcx>) -> usize {
+    struct HighestVarInClauses {
+        max_var: usize,
+        current_index: ty::DebruijnIndex,
+    }
+    impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for HighestVarInClauses {
+        fn visit_binder<T: TypeFoldable<TyCtxt<'tcx>>>(
+            &mut self,
+            t: &ty::Binder<'tcx, T>,
+        ) -> Self::Result {
+            self.current_index.shift_in(1);
+            let t = t.super_visit_with(self);
+            self.current_index.shift_out(1);
+            t
+        }
+        fn visit_ty(&mut self, t: Ty<'tcx>) {
+            if let ty::Bound(debruijn, bound_ty) = *t.kind()
+                && debruijn == self.current_index
+            {
+                self.max_var = self.max_var.max(bound_ty.var.as_usize());
+            } else if t.has_vars_bound_at_or_above(self.current_index) {
+                t.super_visit_with(self);
+            }
+        }
+        fn visit_region(&mut self, r: ty::Region<'tcx>) {
+            if let ty::ReBound(debruijn, bound_region) = r.kind()
+                && debruijn == self.current_index
+            {
+                self.max_var = self.max_var.max(bound_region.var.as_usize());
+            }
+        }
+        fn visit_const(&mut self, ct: ty::Const<'tcx>) {
+            if let ty::ConstKind::Bound(debruijn, bound_const) = ct.kind()
+                && debruijn == self.current_index
+            {
+                self.max_var = self.max_var.max(bound_const.as_usize());
+            } else if ct.has_vars_bound_at_or_above(self.current_index) {
+                ct.super_visit_with(self);
+            }
+        }
     }
+    let mut visitor = HighestVarInClauses { max_var: 0, current_index: ty::INNERMOST };
+    c.visit_with(&mut visitor);
+    visitor.max_var
 }

compiler/rustc_infer/src/infer/canonical/mod.rs

@@ -69,7 +69,7 @@ impl<'tcx> InferCtxt<'tcx> {
 
         let canonical_inference_vars =
             self.instantiate_canonical_vars(span, canonical.variables, |ui| universes[ui]);
-        let result = canonical.instantiate(self.tcx, &canonical_inference_vars);
+        let result = canonical.instantiate::<true>(self.tcx, &canonical_inference_vars);
         (result, canonical_inference_vars)
     }
 

compiler/rustc_infer/src/infer/canonical/query_response.rs

@@ -17,7 +17,7 @@ use rustc_middle::mir::ConstraintCategory;
 use rustc_middle::ty::{self, BoundVar, GenericArg, GenericArgKind, Ty, TyCtxt, TypeFoldable};
 use tracing::{debug, instrument};
 
-use crate::infer::canonical::instantiate::{CanonicalExt, instantiate_value};
+use crate::infer::canonical::instantiate::{CanonicalExt, instantiate_value_0};
 use crate::infer::canonical::{
     Canonical, CanonicalQueryResponse, CanonicalVarValues, Certainty, OriginalQueryValues,
     QueryRegionConstraints, QueryResponse,
@@ -170,12 +170,13 @@ impl<'tcx> InferCtxt<'tcx> {
             self.query_response_instantiation(cause, param_env, original_values, query_response)?;
 
         for (predicate, _category) in &query_response.value.region_constraints.outlives {
-            let predicate = instantiate_value(self.tcx, &result_args, *predicate);
+            let predicate = instantiate_value_0::<_, false>(self.tcx, &result_args, *predicate);
             self.register_outlives_constraint(predicate, cause);
         }
 
         let user_result: R =
-            query_response.instantiate_projected(self.tcx, &result_args, |q_r| q_r.value.clone());
+            query_response
+                .instantiate_projected::<_, false>(self.tcx, &result_args, |q_r| q_r.value.clone());
 
         Ok(InferOk { value: user_result, obligations })
     }
@@ -242,9 +243,10 @@ impl<'tcx> InferCtxt<'tcx> {
 
         for (index, original_value) in original_values.var_values.iter().enumerate() {
             // ...with the value `v_r` of that variable from the query.
-            let result_value = query_response.instantiate_projected(self.tcx, &result_args, |v| {
-                v.var_values[BoundVar::new(index)]
-            });
+            let result_value =
+                query_response.instantiate_projected::<_, false>(self.tcx, &result_args, |v| {
+                    v.var_values[BoundVar::new(index)]
+                });
             match (original_value.kind(), result_value.kind()) {
                 (GenericArgKind::Lifetime(re1), GenericArgKind::Lifetime(re2))
                     if re1.is_erased() && re2.is_erased() =>
@@ -289,7 +291,7 @@ impl<'tcx> InferCtxt<'tcx> {
         // ...also include the other query region constraints from the query.
         output_query_region_constraints.outlives.extend(
             query_response.value.region_constraints.outlives.iter().filter_map(|&r_c| {
-                let r_c = instantiate_value(self.tcx, &result_args, r_c);
+                let r_c = instantiate_value_0::<_, false>(self.tcx, &result_args, r_c);
 
                 // Screen out `'a: 'a` cases.
                 let ty::OutlivesPredicate(k1, r2) = r_c.0;
@@ -298,7 +300,8 @@ impl<'tcx> InferCtxt<'tcx> {
         );
 
         let user_result: R =
-            query_response.instantiate_projected(self.tcx, &result_args, |q_r| q_r.value.clone());
+            query_response
+                .instantiate_projected::<_, false>(self.tcx, &result_args, |q_r| q_r.value.clone());
 
         Ok(InferOk { value: user_result, obligations })
     }
@@ -469,8 +472,8 @@ impl<'tcx> InferCtxt<'tcx> {
 
         // Carry all newly resolved opaque types to the caller's scope
         for &(a, b) in &query_response.value.opaque_types {
-            let a = instantiate_value(self.tcx, &result_args, a);
-            let b = instantiate_value(self.tcx, &result_args, b);
+            let a = instantiate_value_0::<_, false>(self.tcx, &result_args, a);
+            let b = instantiate_value_0::<_, false>(self.tcx, &result_args, b);
             debug!(?a, ?b, "constrain opaque type");
             // We use equate here instead of, for example, just registering the
             // opaque type's hidden value directly, because the hidden type may have been an inference
@@ -512,7 +515,8 @@ impl<'tcx> InferCtxt<'tcx> {
         // `query_response.var_values` after applying the instantiation
         // by `result_args`.
         let instantiated_query_response = |index: BoundVar| -> GenericArg<'tcx> {
-            query_response.instantiate_projected(self.tcx, result_args, |v| v.var_values[index])
+            query_response
+                .instantiate_projected::<_, false>(self.tcx, result_args, |v| v.var_values[index])
         };
 
         // Unify the original value for each variable with the value

compiler/rustc_middle/src/ty/context.rs

@@ -1460,6 +1460,10 @@ pub struct GlobalCtxt<'tcx> {
 
     pub canonical_param_env_cache: CanonicalParamEnvCache<'tcx>,
 
+    pub highest_var_in_clauses_cache: Lock<FxHashMap<ty::Clauses<'tcx>, usize>>,
+    pub clauses_cache:
+        Lock<FxHashMap<(ty::Clauses<'tcx>, &'tcx [ty::GenericArg<'tcx>]), ty::Clauses<'tcx>>>,
+
     /// Data layout specification for the current target.
     pub data_layout: TargetDataLayout,
 
@@ -1707,6 +1711,8 @@ impl<'tcx> TyCtxt<'tcx> {
             new_solver_evaluation_cache: Default::default(),
             new_solver_canonical_param_env_cache: Default::default(),
             canonical_param_env_cache: Default::default(),
+            highest_var_in_clauses_cache: Default::default(),
+            clauses_cache: Default::default(),
             data_layout,
             alloc_map: interpret::AllocMap::new(),
             current_gcx,

compiler/rustc_trait_selection/src/solve/delegate.rs

@@ -222,7 +222,7 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
     where
         V: TypeFoldable<TyCtxt<'tcx>>,
     {
-        canonical.instantiate(self.tcx, &values)
+        canonical.instantiate::<false>(self.tcx, &values)
     }
 
     fn instantiate_canonical_var_with_infer(