optimize `CanonicalVarValues::instantiate`

8 files changed, 62 insertions, 56 deletions

diff --git a/compiler/rustc_hir_typeck/src/method/probe.rs b/compiler/rustc_hir_typeck/src/method/probe.rs
index ab584eb7c90..52a8eff984b 100644
--- a/compiler/rustc_hir_typeck/src/method/probe.rs
+++ b/compiler/rustc_hir_typeck/src/method/probe.rs
@@ -403,15 +403,13 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 // special handling for this "trivial case" is a good idea.
 
                 let infcx = &self.infcx;
-                let (ParamEnvAnd { param_env: _, value: self_ty }, canonical_inference_vars) =
+                let (ParamEnvAnd { param_env: _, value: self_ty }, var_values) =
                     infcx.instantiate_canonical(span, &query_input.canonical);
                 debug!(?self_ty, ?query_input, "probe_op: Mode::Path");
                 MethodAutoderefStepsResult {
                     steps: infcx.tcx.arena.alloc_from_iter([CandidateStep {
-                        self_ty: self.make_query_response_ignoring_pending_obligations(
-                            canonical_inference_vars,
-                            self_ty,
-                        ),
+                        self_ty: self
+                            .make_query_response_ignoring_pending_obligations(var_values, self_ty),
                         autoderefs: 0,
                         from_unsafe_deref: false,
                         unsize: false,
diff --git a/compiler/rustc_infer/src/infer/canonical/mod.rs b/compiler/rustc_infer/src/infer/canonical/mod.rs
index 4e1ea52af4c..517331c3245 100644
--- a/compiler/rustc_infer/src/infer/canonical/mod.rs
+++ b/compiler/rustc_infer/src/infer/canonical/mod.rs
@@ -24,7 +24,7 @@
 pub use instantiate::CanonicalExt;
 use rustc_index::IndexVec;
 pub use rustc_middle::infer::canonical::*;
-use rustc_middle::ty::{self, GenericArg, List, Ty, TyCtxt, TypeFoldable};
+use rustc_middle::ty::{self, GenericArg, Ty, TyCtxt, TypeFoldable};
 use rustc_span::Span;
 
 use crate::infer::{InferCtxt, RegionVariableOrigin};
@@ -67,29 +67,12 @@ impl<'tcx> InferCtxt<'tcx> {
             .chain((1..=canonical.max_universe.as_u32()).map(|_| self.create_next_universe()))
             .collect();
 
-        let canonical_inference_vars =
-            self.instantiate_canonical_vars(span, canonical.variables, |ui| universes[ui]);
-        let result = canonical.instantiate(self.tcx, &canonical_inference_vars);
-        (result, canonical_inference_vars)
-    }
-
-    /// Given the "infos" about the canonical variables from some
-    /// canonical, creates fresh variables with the same
-    /// characteristics (see `instantiate_canonical_var` for
-    /// details). You can then use `instantiate` to instantiate the
-    /// canonical variable with these inference variables.
-    fn instantiate_canonical_vars(
-        &self,
-        span: Span,
-        variables: &List<CanonicalVarKind<'tcx>>,
-        universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
-    ) -> CanonicalVarValues<'tcx> {
-        let mut var_values = Vec::with_capacity(variables.len());
-        for info in variables.iter() {
-            let value = self.instantiate_canonical_var(span, info, &var_values, &universe_map);
-            var_values.push(value);
-        }
-        CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) }
+        let var_values =
+            CanonicalVarValues::instantiate(self.tcx, &canonical.variables, |var_values, info| {
+                self.instantiate_canonical_var(span, info, &var_values, |ui| universes[ui])
+            });
+        let result = canonical.instantiate(self.tcx, &var_values);
+        (result, var_values)
     }
 
     /// Given the "info" about a canonical variable, creates a fresh
diff --git a/compiler/rustc_infer/src/infer/canonical/query_response.rs b/compiler/rustc_infer/src/infer/canonical/query_response.rs
index 22c9c4edccc..5d1b4be9e57 100644
--- a/compiler/rustc_infer/src/infer/canonical/query_response.rs
+++ b/compiler/rustc_infer/src/infer/canonical/query_response.rs
@@ -453,16 +453,17 @@ impl<'tcx> InferCtxt<'tcx> {
         // Create result arguments: if we found a value for a
         // given variable in the loop above, use that. Otherwise, use
         // a fresh inference variable.
-        let mut var_values = Vec::with_capacity(query_response.variables.len());
-        for (index, kind) in query_response.variables.iter().enumerate() {
-            let value = if kind.universe() != ty::UniverseIndex::ROOT {
+        let tcx = self.tcx;
+        let variables = query_response.variables;
+        let var_values = CanonicalVarValues::instantiate(tcx, variables, |var_values, kind| {
+            if kind.universe() != ty::UniverseIndex::ROOT {
                 // A variable from inside a binder of the query. While ideally these shouldn't
                 // exist at all, we have to deal with them for now.
                 self.instantiate_canonical_var(cause.span, kind, &var_values, |u| {
                     universe_map[u.as_usize()]
                 })
             } else if kind.is_existential() {
-                match opt_values[BoundVar::new(index)] {
+                match opt_values[BoundVar::new(var_values.len())] {
                     Some(k) => k,
                     None => self.instantiate_canonical_var(cause.span, kind, &var_values, |u| {
                         universe_map[u.as_usize()]
@@ -471,20 +472,17 @@ impl<'tcx> InferCtxt<'tcx> {
             } else {
                 // For placeholders which were already part of the input, we simply map this
                 // universal bound variable back the placeholder of the input.
-                opt_values[BoundVar::new(index)]
+                opt_values[BoundVar::new(var_values.len())]
                     .expect("expected placeholder to be unified with itself during response")
-            };
-            var_values.push(value);
-        }
-
-        let result_args = CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) };
+            }
+        });
 
         let mut obligations = PredicateObligations::new();
 
         // 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(self.tcx, &var_values, a);
+            let b = instantiate_value(self.tcx, &var_values, 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
@@ -501,7 +499,7 @@ impl<'tcx> InferCtxt<'tcx> {
             );
         }
 
-        Ok(InferOk { value: result_args, obligations })
+        Ok(InferOk { value: var_values, obligations })
     }
 
     /// Given a "guess" at the values for the canonical variables in
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
index 8d7a55c55be..169832ca5fb 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
@@ -365,10 +365,8 @@ where
                 }
             }
         }
-
-        let mut var_values = Vec::with_capacity(response.variables.len());
-        for (index, kind) in response.variables.iter().enumerate() {
-            let value = if kind.universe() != ty::UniverseIndex::ROOT {
+        CanonicalVarValues::instantiate(delegate.cx(), response.variables, |var_values, kind| {
+            if kind.universe() != ty::UniverseIndex::ROOT {
                 // A variable from inside a binder of the query. While ideally these shouldn't
                 // exist at all (see the FIXME at the start of this method), we have to deal with
                 // them for now.
@@ -383,7 +381,7 @@ where
                 // more placeholders then they should be able to. However the inference variables have
                 // to "come from somewhere", so by equating them with the original values of the caller
                 // later on, we pull them down into their correct universe again.
-                if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
+                if let Some(v) = opt_values[ty::BoundVar::from_usize(var_values.len())] {
                     v
                 } else {
                     delegate.instantiate_canonical_var(kind, span, &var_values, |_| prev_universe)
@@ -392,11 +390,8 @@ where
                 // For placeholders which were already part of the input, we simply map this
                 // universal bound variable back the placeholder of the input.
                 original_values[kind.expect_placeholder_index()]
-            };
-            var_values.push(value)
-        }
-
-        CanonicalVarValues { var_values: delegate.cx().mk_args(&var_values) }
+            }
+        })
     }
 
     /// Unify the `original_values` with the `var_values` returned by the canonical query..
diff --git a/compiler/rustc_trait_selection/src/infer.rs b/compiler/rustc_trait_selection/src/infer.rs
index 7c6b7b14ecb..4c50c44b841 100644
--- a/compiler/rustc_trait_selection/src/infer.rs
+++ b/compiler/rustc_trait_selection/src/infer.rs
@@ -184,10 +184,9 @@ impl<'tcx> InferCtxtBuilder<'tcx> {
         R: Debug + TypeFoldable<TyCtxt<'tcx>>,
         Canonical<'tcx, QueryResponse<'tcx, R>>: ArenaAllocatable<'tcx>,
     {
-        let (infcx, key, canonical_inference_vars) =
-            self.build_with_canonical(DUMMY_SP, canonical_key);
+        let (infcx, key, var_values) = self.build_with_canonical(DUMMY_SP, canonical_key);
         let ocx = ObligationCtxt::new(&infcx);
         let value = operation(&ocx, key)?;
-        ocx.make_canonicalized_query_response(canonical_inference_vars, value)
+        ocx.make_canonicalized_query_response(var_values, value)
     }
 }
diff --git a/compiler/rustc_traits/src/evaluate_obligation.rs b/compiler/rustc_traits/src/evaluate_obligation.rs
index 819b8e3231c..8f72bdf0972 100644
--- a/compiler/rustc_traits/src/evaluate_obligation.rs
+++ b/compiler/rustc_traits/src/evaluate_obligation.rs
@@ -19,7 +19,7 @@ fn evaluate_obligation<'tcx>(
 ) -> Result<EvaluationResult, OverflowError> {
     assert!(!tcx.next_trait_solver_globally());
     debug!("evaluate_obligation(canonical_goal={:#?})", canonical_goal);
-    let (ref infcx, goal, _canonical_inference_vars) =
+    let (ref infcx, goal, _var_values) =
         tcx.infer_ctxt().build_with_canonical(DUMMY_SP, &canonical_goal);
     debug!("evaluate_obligation: goal={:#?}", goal);
     let ParamEnvAnd { param_env, value: predicate } = goal;
diff --git a/compiler/rustc_type_ir/Cargo.toml b/compiler/rustc_type_ir/Cargo.toml
index d55e9b3b1be..6ccddb17ff2 100644
--- a/compiler/rustc_type_ir/Cargo.toml
+++ b/compiler/rustc_type_ir/Cargo.toml
@@ -5,6 +5,7 @@ edition = "2024"
 
 [dependencies]
 # tidy-alphabetical-start
+arrayvec = { version = "0.7", default-features = false }
 bitflags = "2.4.1"
 derive-where = "1.2.7"
 ena = "0.14.3"
diff --git a/compiler/rustc_type_ir/src/canonical.rs b/compiler/rustc_type_ir/src/canonical.rs
index 04bf9008035..ecf3ae4f8b2 100644
--- a/compiler/rustc_type_ir/src/canonical.rs
+++ b/compiler/rustc_type_ir/src/canonical.rs
@@ -1,6 +1,7 @@
 use std::fmt;
 use std::ops::Index;
 
+use arrayvec::ArrayVec;
 use derive_where::derive_where;
 #[cfg(feature = "nightly")]
 use rustc_macros::{Decodable_NoContext, Encodable_NoContext, HashStable_NoContext};
@@ -306,6 +307,37 @@ impl<I: Interner> CanonicalVarValues<I> {
         CanonicalVarValues { var_values: Default::default() }
     }
 
+    pub fn instantiate(
+        cx: I,
+        variables: I::CanonicalVarKinds,
+        mut f: impl FnMut(&[I::GenericArg], CanonicalVarKind<I>) -> I::GenericArg,
+    ) -> CanonicalVarValues<I> {
+        // Instantiating `CanonicalVarValues` is really hot, but limited to less than
+        // 4 most of the time. Avoid creating a `Vec` here.
+        if variables.len() <= 4 {
+            let mut var_values = ArrayVec::<_, 4>::new();
+            for info in variables.iter() {
+                var_values.push(f(&var_values, info));
+            }
+            CanonicalVarValues { var_values: cx.mk_args(&var_values) }
+        } else {
+            CanonicalVarValues::instantiate_cold(cx, variables, f)
+        }
+    }
+
+    #[cold]
+    fn instantiate_cold(
+        cx: I,
+        variables: I::CanonicalVarKinds,
+        mut f: impl FnMut(&[I::GenericArg], CanonicalVarKind<I>) -> I::GenericArg,
+    ) -> CanonicalVarValues<I> {
+        let mut var_values = Vec::with_capacity(variables.len());
+        for info in variables.iter() {
+            var_values.push(f(&var_values, info));
+        }
+        CanonicalVarValues { var_values: cx.mk_args(&var_values) }
+    }
+
     #[inline]
     pub fn len(&self) -> usize {
         self.var_values.len()