New solver opts for diesel (part 2)

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

@@ -52,7 +52,7 @@ impl<'tcx, V> Canonical<'tcx, V> {
 /// 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 fn instantiate_value<'tcx, T>(
     tcx: TyCtxt<'tcx>,
     var_values: &CanonicalVarValues<'tcx>,
     value: T,

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

@@ -21,7 +21,7 @@
 //!
 //! [c]: https://rust-lang.github.io/chalk/book/canonical_queries/canonicalization.html
 
-pub use instantiate::CanonicalExt;
+pub use instantiate::{CanonicalExt, instantiate_value};
 use rustc_index::IndexVec;
 pub use rustc_middle::infer::canonical::*;
 use rustc_middle::ty::{self, GenericArg, Ty, TyCtxt, TypeFoldable};

compiler/rustc_middle/src/infer/canonical.rs

@@ -34,6 +34,7 @@ use crate::ty::{self, GenericArg, GenericArgsRef, List, Ty, TyCtxt, TypeFlags, T
 
 pub type CanonicalQueryInput<'tcx, V> = ir::CanonicalQueryInput<TyCtxt<'tcx>, V>;
 pub type Canonical<'tcx, V> = ir::Canonical<TyCtxt<'tcx>, V>;
+pub type Canonicalized<V> = ir::Canonicalized<V>;
 pub type CanonicalVarKind<'tcx> = ir::CanonicalVarKind<TyCtxt<'tcx>>;
 pub type CanonicalVarValues<'tcx> = ir::CanonicalVarValues<TyCtxt<'tcx>>;
 pub type CanonicalVarKinds<'tcx> = &'tcx List<CanonicalVarKind<'tcx>>;

compiler/rustc_next_trait_solver/src/canonical/mod.rs

@@ -18,16 +18,17 @@ use rustc_type_ir::relate::{
     self, Relate, RelateResult, TypeRelation, VarianceDiagInfo, relate_args_invariantly,
 };
 use rustc_type_ir::{
-    self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, Const, InferCtxtLike, Interner,
-    Region, TypeFoldable, TypingMode, TypingModeEqWrapper,
+    self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, Canonicalized, Const,
+    InferCtxtLike, Interner, Region, TypeFoldable, TypingMode, TypingModeEqWrapper,
 };
 use thin_vec::ThinVec;
 use tracing::instrument;
 
 use crate::delegate::SolverDelegate;
 use crate::solve::{
     CanonicalResponse, Certainty, ExternalConstraintsData, ExternalRegionConstraints, Goal,
-    NestedNormalizationGoals, QueryInput, Response, VisibleForLeakCheck, inspect,
+    IsIdentity, NestedNormalizationGoals, QueryInput, Response, VisibleForLeakCheck,
+    has_no_inference_or_external_constraints, inspect,
 };
 
 pub mod canonicalizer;
@@ -93,7 +94,7 @@ where
 /// After calling a canonical query, we apply the constraints returned
 /// by the query using this function.
 ///
-/// This happens in three steps:
+/// If a quick early-return check fails, this happens in three steps:
 /// - we instantiate the bound variables of the query response
 /// - we unify the `var_values` of the response with the `original_values`
 /// - we apply the `external_constraints` returned by the query, returning
@@ -102,19 +103,39 @@ pub(super) fn instantiate_and_apply_query_response<D, I>(
     delegate: &D,
     original_values: &[I::GenericArg],
     response: CanonicalResponse<I>,
+    canonical_goal_var_kinds: I::CanonicalVarKinds,
     span: I::Span,
 ) -> (NestedNormalizationGoals<I>, Certainty)
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
 {
+    let certainty = response.value.certainty;
+
+    // If the response is an identity, and there are no external constraints, and the response var
+    // kinds match the canonical goal var kinds, this function is a no-op and we can return early.
+    if has_no_inference_or_external_constraints(response)
+        && response.var_kinds == canonical_goal_var_kinds
+    {
+        return (NestedNormalizationGoals::empty(), certainty);
+    }
+
     let instantiation =
         compute_query_response_instantiation_values(delegate, &original_values, &response, span);
 
-    let Response { var_values, external_constraints, certainty } =
-        delegate.instantiate_canonical(response, instantiation);
-
-    unify_query_var_values(delegate, &original_values, var_values, span);
+    // We could call `instantiate_canonical` on `response`, but it's more efficient to instantiate
+    // the fields within `response` separately. That way we can instantiate and unify `var_values`
+    // in a single step, rather than instantiating first (which creates an intermediate list) and
+    // then unifying afterwards.
+    let external_constraints =
+        delegate.instantiate_value(response.project(|r| r.external_constraints), instantiation);
+    instantiate_and_unify_query_var_values(
+        delegate,
+        &original_values,
+        response.project(|r| r.var_values),
+        instantiation,
+        span,
+    );
 
     let ExternalConstraintsData { region_constraints, opaque_types, normalization_nested_goals } =
         &*external_constraints;
@@ -474,7 +495,8 @@ where
     }
 }
 
-/// Unify the `original_values` with the `var_values` returned by the canonical query..
+/// Instantiate `var_values` as returned by the canonical query and also unify it with
+/// `original_values`.
 ///
 /// This assumes that this unification will always succeed. This is the case when
 /// applying a query response right away. However, calling a canonical query, doing any
@@ -487,10 +509,11 @@ where
 /// from the solver we assume that the solver correctly handled aliases and therefore
 /// always relate them structurally here.
 #[instrument(level = "trace", skip(delegate))]
-fn unify_query_var_values<D, I>(
+fn instantiate_and_unify_query_var_values<D, I>(
     delegate: &D,
     original_values: &[I::GenericArg],
-    var_values: I::GenericArgs,
+    var_values: Canonicalized<I::GenericArgs>,
+    instantiation: CanonicalVarValues<I>,
     span: I::Span,
 ) where
     D: SolverDelegate<Interner = I>,
@@ -499,6 +522,7 @@ fn unify_query_var_values<D, I>(
     assert_eq!(original_values.len(), var_values.len());
 
     for (&orig, response) in iter::zip(original_values, var_values.iter()) {
+        let response = instantiate_generic_arg(delegate, response, instantiation);
         if orig == response {
             continue;
         }
@@ -508,6 +532,41 @@ fn unify_query_var_values<D, I>(
     }
 }
 
+fn instantiate_generic_arg<D, I>(
+    delegate: &D,
+    arg: Canonicalized<I::GenericArg>,
+    instantiation: CanonicalVarValues<I>,
+) -> I::GenericArg
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+{
+    // Canonical bound vars are very common here. `CanonicalInstantiator::fold_{ty,region,const}`'s
+    // handling of them is simple, and we replicate that handling here to avoid the cost of the
+    // surrounding machinery.
+    fn canonical_bound_var<I: Interner>(arg: Canonicalized<I::GenericArg>) -> Option<ty::BoundVar> {
+        match arg.into_uninstantiated().kind() {
+            ty::GenericArgKind::Type(t) => match t.kind() {
+                ty::Bound(ty::BoundVarIndexKind::Canonical, bt) => Some(bt.var()),
+                _ => None,
+            },
+            ty::GenericArgKind::Lifetime(r) => match r.kind() {
+                ty::ReBound(ty::BoundVarIndexKind::Canonical, br) => Some(br.var()),
+                _ => None,
+            },
+            ty::GenericArgKind::Const(c) => match c.kind() {
+                ty::ConstKind::Bound(ty::BoundVarIndexKind::Canonical, bc) => Some(bc.var()),
+                _ => None,
+            },
+        }
+    }
+
+    match canonical_bound_var::<I>(arg) {
+        Some(var) => instantiation[var],
+        None => delegate.instantiate_value(arg, instantiation),
+    }
+}
+
 fn register_region_constraints<D, I>(
     delegate: &D,
     constraints: impl IntoIterator<Item = (ty::RegionConstraint<I>, VisibleForLeakCheck)>,
@@ -586,7 +645,7 @@ pub fn instantiate_canonical_state<D, I, T>(
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
-    T: TypeFoldable<I>,
+    T: TypeFoldable<I> + Copy,
 {
     // In case any fresh inference variables have been created between `state`
     // and the previous instantiation, extend `orig_values` for it.
@@ -608,9 +667,17 @@ where
         prev_universe,
     );
 
-    let inspect::State { var_values, data } = delegate.instantiate_canonical(state, instantiation);
+    // We instantiate the fields within `state` separately. See the corresponding code in
+    // `instantiate_and_apply_query_response` for an explanation why.
+    let data = delegate.instantiate_value(state.project(|r| r.data), instantiation);
+    instantiate_and_unify_query_var_values(
+        delegate,
+        orig_values,
+        state.project(|r| r.var_values),
+        instantiation,
+        span,
+    );
 
-    unify_query_var_values(delegate, orig_values, var_values, span);
     data
 }
 
@@ -644,11 +711,12 @@ pub fn response_no_constraints_raw<I: Interner>(
         max_universe,
         var_kinds,
         value: Response {
+            certainty,
             var_values,
+            var_values_is_identity: IsIdentity::Yes,
             // FIXME: maybe we should store the "no response" version in cx, like
             // we do for cx.types and stuff.
             external_constraints: cx.mk_external_constraints(ExternalConstraintsData::new(cx)),
-            certainty,
         },
     }
 }

compiler/rustc_next_trait_solver/src/delegate.rs

@@ -79,6 +79,14 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
         &self,
     ) -> Vec<(ty::RegionConstraint<Self::Interner>, VisibleForLeakCheck)>;
 
+    fn instantiate_value<V>(
+        &self,
+        value: ty::Canonicalized<V>,
+        values: ty::CanonicalVarValues<Self::Interner>,
+    ) -> V
+    where
+        V: TypeFoldable<Self::Interner>;
+
     fn instantiate_canonical<V>(
         &self,
         canonical: ty::Canonical<Self::Interner, V>,

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs

@@ -43,9 +43,9 @@ use crate::solve::search_graph::SearchGraph;
 use crate::solve::ty::may_use_unstable_feature;
 use crate::solve::{
     CanonicalResponse, Certainty, ExternalConstraintsData, FIXPOINT_STEP_LIMIT, Goal,
-    GoalEvaluation, GoalSource, GoalStalledOn, GoalStalledOnOpaques, HasChanged, MaybeCause,
-    NestedNormalizationGoals, NoSolution, QueryInput, QueryResult, Response, SucceededInErased,
-    VisibleForLeakCheck, inspect,
+    GoalEvaluation, GoalSource, GoalStalledOn, GoalStalledOnOpaques, HasChanged, IsIdentity,
+    MaybeCause, NestedNormalizationGoals, NoSolution, QueryInput, QueryResult, Response,
+    SucceededInErased, VisibleForLeakCheck, inspect,
 };
 
 pub mod fast_path;
@@ -883,6 +883,7 @@ where
             self.delegate,
             &orig_values,
             response,
+            canonical_goal.canonical.var_kinds,
             self.origin_span,
         );
 
@@ -1711,15 +1712,17 @@ where
 
         filter_irrelevant_region_constraints(self.delegate, var_values, &mut external_constraints);
 
-        let canonical = canonicalize_response(
+        let mut canonical = canonicalize_response(
             self.delegate,
             self.max_input_universe,
             Response {
-                var_values,
                 certainty,
+                var_values,
+                var_values_is_identity: IsIdentity::No, // updated below, after canonicalization
                 external_constraints: self.cx().mk_external_constraints(external_constraints),
             },
         );
+        canonical.value.update_var_values_is_identity();
 
         Ok(canonical)
     }
@@ -2062,6 +2065,7 @@ pub(super) fn evaluate_root_goal_for_proof_tree<D: SolverDelegate<Interner = I>,
         delegate,
         &proof_tree.orig_values,
         response,
+        canonical_goal.canonical.var_kinds,
         origin_span,
     );
 

compiler/rustc_next_trait_solver/src/solve/mod.rs

@@ -57,15 +57,15 @@ pub enum HasChanged {
 
 // FIXME(trait-system-refactor-initiative#117): we don't detect whether a response
 // ended up pulling down any universes.
-fn has_no_inference_or_external_constraints<I: Interner>(
+pub(crate) fn has_no_inference_or_external_constraints<I: Interner>(
     response: ty::Canonical<I, Response<I>>,
 ) -> bool {
     let ExternalConstraintsData {
         ref region_constraints,
         ref opaque_types,
         ref normalization_nested_goals,
     } = *response.value.external_constraints;
-    response.value.is_identity()
+    response.value.var_values_is_identity == IsIdentity::Yes
         && region_constraints.is_empty()
         && opaque_types.is_empty()
         && normalization_nested_goals.is_empty()
@@ -77,7 +77,7 @@ fn has_only_region_constraints<I: Interner>(response: ty::Canonical<I, Response<
         ref opaque_types,
         ref normalization_nested_goals,
     } = *response.value.external_constraints;
-    response.value.is_identity_modulo_regions()
+    response.value.var_values_is_identity >= IsIdentity::ModuloRegions
         && opaque_types.is_empty()
         && normalization_nested_goals.is_empty()
 }
@@ -93,8 +93,9 @@ fn equal_response_modulo_region_constraints<I: Interner>(
         var_kinds: a_var_kinds,
         value:
             Response {
-                var_values: a_var_values,
                 certainty: a_certainty,
+                var_values: a_var_values,
+                var_values_is_identity: _, // ignore because it's purely derived from `var_values`
                 external_constraints: a_external_constraints,
             },
     } = a;

compiler/rustc_next_trait_solver/src/solve/trait_goals.rs

@@ -5,9 +5,9 @@ use rustc_type_ir::fast_reject::DeepRejectCtxt;
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::lang_items::SolverTraitLangItem;
 use rustc_type_ir::solve::{
-    AliasBoundKind, CandidatePreferenceMode, CanonicalResponse, ExternalConstraintsData, MaybeInfo,
-    NoSolutionOrRerunNonErased, OpaqueTypesJank, QueryResultOrRerunNonErased, RerunNonErased,
-    RerunReason, RerunResultExt, SizedTraitKind,
+    AliasBoundKind, CandidatePreferenceMode, CanonicalResponse, ExternalConstraintsData,
+    IsIdentity, MaybeInfo, NoSolutionOrRerunNonErased, OpaqueTypesJank,
+    QueryResultOrRerunNonErased, RerunNonErased, RerunReason, RerunResultExt, SizedTraitKind,
 };
 use rustc_type_ir::{
     self as ty, ClausePolarity, ExistentialPredicate, FieldInfo, Interner, MayBeErased, Movability,
@@ -1332,7 +1332,7 @@ where
 
         // Proving an auto trait for the hidden type must not constrain inference
         // variables, as that would leak the hidden type itself.
-        if !candidate.result.value.is_identity_modulo_regions() {
+        if candidate.result.value.var_values_is_identity == IsIdentity::No {
             return self.forced_ambiguity(MaybeInfo::AMBIGUOUS);
         }
 

compiler/rustc_trait_selection/src/solve/delegate.rs

@@ -11,7 +11,7 @@ use rustc_hir::def_id::{CRATE_DEF_ID, DefId, LOCAL_CRATE};
 use rustc_infer::infer::canonical::query_response::make_query_region_constraints;
 use rustc_infer::infer::canonical::{
     Canonical, CanonicalExt as _, CanonicalQueryInput, CanonicalVarKind, CanonicalVarValues,
-    QueryRegionConstraint,
+    Canonicalized, QueryRegionConstraint, instantiate_value,
 };
 use rustc_infer::infer::{InferCtxt, RegionVariableOrigin, SubregionOrigin, TyCtxtInferExt};
 use rustc_infer::traits::solve::{
@@ -406,6 +406,13 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
         constraints
     }
 
+    fn instantiate_value<V>(&self, value: Canonicalized<V>, values: CanonicalVarValues<'tcx>) -> V
+    where
+        V: TypeFoldable<TyCtxt<'tcx>>,
+    {
+        instantiate_value(self.tcx, &values, value.into_uninstantiated())
+    }
+
     fn instantiate_canonical<V>(
         &self,
         canonical: Canonical<'tcx, V>,