deduplicate code and address review

Author: fee1-dead

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -401,6 +401,8 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeMachine<'tcx> {
 
     const PANIC_ON_ALLOC_FAIL: bool = false; // will be raised as a proper error
 
+    const SHOULD_RESPECT_CONST_BOUNDS_WHEN_RESOLVING_INSTANCES: bool = true;
+
     #[inline(always)]
     fn enforce_alignment(ecx: &InterpCx<'tcx, Self>) -> bool {
         matches!(ecx.machine.check_alignment, CheckAlignment::Error)
@@ -952,8 +954,6 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeMachine<'tcx> {
 
     fn get_default_alloc_params(&self) -> <Self::Bytes as mir::interpret::AllocBytes>::AllocParams {
     }
-
-    const SHOULD_RESPECT_CONST_BOUNDS_WHEN_RESOLVING_INSTANCES: bool = true;
 }
 
 // Please do not add any code below the above `Machine` trait impl. I (oli-obk) plan more cleanups

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -161,6 +161,10 @@ pub trait Machine<'tcx>: Sized {
     /// already been checked before.
     const ALL_CONSTS_ARE_PRECHECKED: bool = true;
 
+    /// Should we resolve instances as if we are in a const context? This means specialized impls whose
+    /// `T: [const] Trait` bounds that are not satisfied will not be selected.
+    const SHOULD_RESPECT_CONST_BOUNDS_WHEN_RESOLVING_INSTANCES: bool = false;
+
     /// Whether memory accesses should be alignment-checked.
     fn enforce_alignment(ecx: &InterpCx<'tcx, Self>) -> bool;
 
@@ -638,8 +642,6 @@ pub trait Machine<'tcx>: Sized {
     fn enter_trace_span(_span: impl FnOnce() -> tracing::Span) -> impl EnteredTraceSpan {
         ()
     }
-
-    const SHOULD_RESPECT_CONST_BOUNDS_WHEN_RESOLVING_INSTANCES: bool = false;
 }
 
 /// A lot of the flexibility above is just needed for `Miri`, but all "compile-time" machines

compiler/rustc_middle/src/queries.rs

@@ -1606,6 +1606,13 @@ rustc_queries! {
         desc { "computing candidate for `{}`", key.value }
     }
 
+    query codegen_select_candidate_for_ctfe(
+        key: PseudoCanonicalInput<'tcx, ty::TraitRef<'tcx>>
+    ) -> Result<&'tcx ImplSource<'tcx, ()>, CodegenObligationError> {
+        cache_on_disk
+        desc { "computing const candidate for `{}`", key.value }
+    }
+
     /// Return all `impl` blocks in the current crate.
     query all_local_trait_impls(_: ()) -> &'tcx rustc_data_structures::fx::FxIndexMap<DefId, Vec<LocalDefId>> {
         desc { "finding local trait impls" }

compiler/rustc_traits/src/codegen.rs

@@ -8,6 +8,7 @@ use rustc_middle::bug;
 use rustc_middle::traits::CodegenObligationError;
 use rustc_middle::ty::{self, PseudoCanonicalInput, TyCtxt, TypeVisitableExt, Unnormalized};
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
+use rustc_trait_selection::solve::InferCtxtSelectExt as _;
 use rustc_trait_selection::traits::{
     ImplSource, Obligation, ObligationCause, ObligationCtxt, ScrubbedTraitError, SelectionContext,
     SelectionError,
@@ -21,11 +22,15 @@ use tracing::debug;
 /// obligations *could be* resolved if we wanted to.
 ///
 /// This also expects that `trait_ref` is fully normalized.
-pub(crate) fn codegen_select_candidate<'tcx>(
+///
+/// When `use_const` is set, we assume the new trait solver is enabled and use
+/// `HostEffectPredicate` with `constness` set to `Const`.
+pub(crate) fn codegen_select_candidate_inner<'tcx>(
     tcx: TyCtxt<'tcx>,
-    key: PseudoCanonicalInput<'tcx, ty::TraitRef<'tcx>>,
+    typing_env: ty::TypingEnv<'tcx>,
+    trait_ref: ty::TraitRef<'tcx>,
+    use_const: bool,
 ) -> Result<&'tcx ImplSource<'tcx, ()>, CodegenObligationError> {
-    let PseudoCanonicalInput { typing_env, value: trait_ref } = key;
     // We expect the input to be fully normalized.
     debug_assert_eq!(
         trait_ref,
@@ -38,14 +43,26 @@ pub(crate) fn codegen_select_candidate<'tcx>(
     let mut selcx = SelectionContext::new(&infcx);
 
     let obligation_cause = ObligationCause::dummy();
-    let obligation = Obligation::new(tcx, obligation_cause, param_env, trait_ref);
 
-    let selection = match selcx.select(&obligation) {
+    let selection = if use_const {
+        let host_predicate =
+            ty::HostEffectPredicate { trait_ref, constness: ty::BoundConstness::Const };
+        let obligation = Obligation::new(tcx, obligation_cause, param_env, host_predicate);
+
+        infcx.select_host_effect_predicate_in_new_trait_solver(&obligation)
+    } else {
+        let obligation = Obligation::new(tcx, obligation_cause, param_env, trait_ref);
+        selcx.select(&obligation)
+    };
+
+    let selection = match selection {
         Ok(Some(selection)) => selection,
         Ok(None) => return Err(CodegenObligationError::Ambiguity),
         Err(SelectionError::Unimplemented) => return Err(CodegenObligationError::Unimplemented),
         Err(e) => {
-            bug!("Encountered error `{:?}` selecting `{:?}` during codegen", e, trait_ref)
+            bug!(
+                "Encountered error `{e:?}` selecting `{trait_ref:?}` during codegen (use_const: {use_const})"
+            )
         }
     };
 
@@ -94,3 +111,17 @@ pub(crate) fn codegen_select_candidate<'tcx>(
 
     Ok(&*tcx.arena.alloc(impl_source))
 }
+
+pub(crate) fn codegen_select_candidate<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: PseudoCanonicalInput<'tcx, ty::TraitRef<'tcx>>,
+) -> Result<&'tcx ImplSource<'tcx, ()>, CodegenObligationError> {
+    codegen_select_candidate_inner(tcx, key.typing_env, key.value, false)
+}
+
+pub(crate) fn codegen_select_candidate_for_ctfe<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: PseudoCanonicalInput<'tcx, ty::TraitRef<'tcx>>,
+) -> Result<&'tcx ImplSource<'tcx, ()>, CodegenObligationError> {
+    codegen_select_candidate_inner(tcx, key.typing_env, key.value, true)
+}

compiler/rustc_traits/src/lib.rs

@@ -25,5 +25,6 @@ pub fn provide(p: &mut Providers) {
     normalize_erasing_regions::provide(p);
     type_op::provide(p);
     p.codegen_select_candidate = codegen::codegen_select_candidate;
+    p.codegen_select_candidate_for_ctfe = codegen::codegen_select_candidate_for_ctfe;
     p.coroutine_hidden_types = coroutine_witnesses::coroutine_hidden_types;
 }

compiler/rustc_ty_utils/src/instance.rs

@@ -14,7 +14,6 @@ use rustc_middle::ty::{
 };
 use rustc_span::sym;
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt as _;
-use rustc_trait_selection::solve::InferCtxtSelectExt as _;
 use rustc_trait_selection::traits::{self, ObligationCtxt};
 use tracing::debug;
 use traits::translate_args;
@@ -119,81 +118,23 @@ fn resolve_associated_item<'tcx>(
     let trait_ref = ty::TraitRef::from_assoc(tcx, trait_id, rcvr_args);
 
     let input = typing_env.as_query_input(trait_ref);
-    let vtbl = if constness == Constness::Const && tcx.next_trait_solver_globally() {
-        let (infcx, param_env) =
-            tcx.infer_ctxt().ignoring_regions().build_with_typing_env(typing_env);
-
-        let obligation_cause = ObligationCause::dummy();
-        let host_predicate =
-            ty::HostEffectPredicate { trait_ref, constness: ty::BoundConstness::Const };
-        let obligation = Obligation::new(tcx, obligation_cause, param_env, host_predicate);
-
-        let selection = match infcx.select_host_effect_predicate_in_new_trait_solver(&obligation) {
-            Ok(Some(selection)) => selection,
-            Ok(None) => return Ok(None),
-            Err(SelectionError::Unimplemented) => return Ok(None),
-            Err(e) => {
-                bug!("Encountered error `{:?}` selecting `{:?}` during codegen", e, host_predicate)
-            }
-        };
-
-        debug!(?selection);
-
-        // Currently, we use a fulfillment context to completely resolve
-        // all nested obligations. This is because they can inform the
-        // inference of the impl's type parameters.
-        let ocx = ObligationCtxt::new(&infcx);
-        let impl_source = selection.map(|obligation| {
-            ocx.register_obligation(obligation);
-        });
+    let candidate = if constness == Constness::Const && tcx.next_trait_solver_globally() {
+        tcx.codegen_select_candidate_for_ctfe(input)
+    } else {
+        tcx.codegen_select_candidate(input)
+    };
 
-        // In principle, we only need to do this so long as `impl_source`
-        // contains unbound type parameters. It could be a slight
-        // optimization to stop iterating early.
-        let errors = ocx.evaluate_obligations_error_on_ambiguity();
-        if !errors.is_empty() {
-            // `rustc_monomorphize::collector` assumes there are no type errors.
-            // Cycle errors are the only post-monomorphization errors possible; emit them now so
-            // `rustc_ty_utils::resolve_associated_item` doesn't return `None` post-monomorphization.
-            for err in errors {
-                if let ScrubbedTraitError::Cycle(cycle) = err {
-                    infcx.err_ctxt().report_overflow_obligation_cycle(&cycle);
-                }
-            }
+    let impl_src = match candidate {
+        Ok(impl_src) => impl_src,
+        Err(CodegenObligationError::Ambiguity | CodegenObligationError::Unimplemented) => {
             return Ok(None);
         }
-
-        let impl_source = infcx.resolve_vars_if_possible(impl_source);
-        let impl_source = tcx.erase_and_anonymize_regions(impl_source);
-        if impl_source.has_non_region_infer() {
-            // Unused generic types or consts on an impl get replaced with inference vars,
-            // but never resolved, causing the return value of a query to contain inference
-            // vars. We do not have a concept for this and will in fact ICE in stable hashing
-            // of the return value. So bail out instead.
-            let guar = match impl_source {
-                ImplSource::UserDefined(impl_) => tcx.dcx().span_delayed_bug(
-                    tcx.def_span(impl_.impl_def_id),
-                    "this impl has unconstrained generic parameters",
-                ),
-                _ => unreachable!(),
-            };
-            return Err(guar);
-        }
-
-        &*tcx.arena.alloc(impl_source)
-    } else {
-        match tcx.codegen_select_candidate(input) {
-            Ok(vtbl) => vtbl,
-            Err(CodegenObligationError::Ambiguity | CodegenObligationError::Unimplemented) => {
-                return Ok(None);
-            }
-            Err(CodegenObligationError::UnconstrainedParam(guar)) => return Err(guar),
-        }
+        Err(CodegenObligationError::UnconstrainedParam(guar)) => return Err(guar),
     };
 
     // Now that we know which impl is being used, we can dispatch to
     // the actual function:
-    Ok(match vtbl {
+    Ok(match impl_src {
         traits::ImplSource::UserDefined(impl_data) => {
             debug!(
                 "resolving ImplSource::UserDefined: {:?}, {:?}, {:?}, {:?}",