resolve instances in ctfe based on current const context

Author: fee1-dead

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -935,6 +935,8 @@ 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/eval_context.rs

@@ -370,7 +370,12 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         trace!("resolve: {:?}, {:#?}", def, args);
         trace!("typing_env: {:#?}", self.typing_env);
         trace!("args: {:#?}", args);
-        match ty::Instance::try_resolve(*self.tcx, self.typing_env, def, args) {
+        let resolve = if M::SHOULD_RESPECT_CONST_BOUNDS_WHEN_RESOLVING_INSTANCES {
+            ty::Instance::try_resolve_for_ctfe
+        } else {
+            ty::Instance::try_resolve
+        };
+        match resolve(*self.tcx, self.typing_env, def, args) {
             Ok(Some(instance)) => interp_ok(instance),
             Ok(None) => throw_inval!(TooGeneric),
 

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -638,6 +638,8 @@ 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_hir/src/hir.rs

@@ -4521,7 +4521,7 @@ impl fmt::Display for Safety {
     }
 }
 
-#[derive(Copy, Clone, PartialEq, Eq, Debug, Encodable, Decodable, StableHash)]
+#[derive(Copy, Clone, PartialEq, Eq, Debug, Encodable, Decodable, StableHash, Hash)]
 #[derive(Default)]
 pub enum Constness {
     #[default]

compiler/rustc_middle/src/queries.rs

@@ -2622,7 +2622,7 @@ rustc_queries! {
     ///    from `Ok(None)` to avoid misleading diagnostics when an error
     ///    has already been/will be emitted, for the original cause.
     query resolve_instance_raw(
-        key: ty::PseudoCanonicalInput<'tcx, (DefId, GenericArgsRef<'tcx>)>
+        key: ty::PseudoCanonicalInput<'tcx, (DefId, GenericArgsRef<'tcx>, hir::Constness)>
     ) -> Result<Option<ty::Instance<'tcx>>, ErrorGuaranteed> {
         desc { "resolving instance `{}`", ty::Instance::new_raw(key.value.0, key.value.1) }
     }

compiler/rustc_middle/src/query/keys.rs

@@ -6,6 +6,7 @@ use std::hash::Hash;
 
 use rustc_ast::tokenstream::TokenStream;
 use rustc_data_structures::stable_hasher::StableHash;
+use rustc_hir as hir;
 use rustc_hir::def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId, LocalModDefId};
 use rustc_hir::hir_id::OwnerId;
 use rustc_span::{DUMMY_SP, Ident, LocalExpnId, Span, Symbol};
@@ -243,6 +244,12 @@ impl<'tcx> QueryKey for (DefId, GenericArgsRef<'tcx>) {
     }
 }
 
+impl<'tcx> QueryKey for (DefId, GenericArgsRef<'tcx>, hir::Constness) {
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        self.0.default_span(tcx)
+    }
+}
+
 impl<'tcx> QueryKey for ty::TraitRef<'tcx> {
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
         tcx.def_span(self.def_id)

compiler/rustc_middle/src/ty/instance.rs

@@ -494,15 +494,16 @@ impl<'tcx> Instance<'tcx> {
     /// couldn't complete due to errors elsewhere - this is distinct
     /// from `Ok(None)` to avoid misleading diagnostics when an error
     /// has already been/will be emitted, for the original cause
-    #[instrument(level = "debug", skip(tcx), ret)]
-    pub fn try_resolve(
+    fn try_resolve_inner(
         tcx: TyCtxt<'tcx>,
         typing_env: ty::TypingEnv<'tcx>,
         def_id: DefId,
         args: GenericArgsRef<'tcx>,
+        mut constness: hir::Constness,
     ) -> Result<Option<Instance<'tcx>>, ErrorGuaranteed> {
+        let def_kind = tcx.def_kind(def_id);
         assert_matches!(
-            tcx.def_kind(def_id),
+            def_kind,
             DefKind::Fn
                 | DefKind::AssocFn
                 | DefKind::Const { .. }
@@ -518,6 +519,10 @@ impl<'tcx> Instance<'tcx> {
             `try_normalize_erasing_regions`."
         );
 
+        if !def_kind.is_assoc() {
+            constness = hir::Constness::NotConst;
+        }
+
         // Rust code can easily create exponentially-long types using only a
         // polynomial recursion depth. Even with the default recursion
         // depth, you can easily get cases that take >2^60 steps to run,
@@ -530,11 +535,36 @@ impl<'tcx> Instance<'tcx> {
             return Ok(None);
         }
 
+        let input = tcx.erase_and_anonymize_regions(typing_env.as_query_input((def_id, args)));
+
         // All regions in the result of this query are erased, so it's
         // fine to erase all of the input regions.
-        tcx.resolve_instance_raw(
-            tcx.erase_and_anonymize_regions(typing_env.as_query_input((def_id, args))),
-        )
+        tcx.resolve_instance_raw(ty::PseudoCanonicalInput {
+            typing_env: input.typing_env,
+            value: (input.value.0, input.value.1, constness),
+        })
+    }
+
+    /// See `try_resolve_inner`.
+    #[instrument(level = "debug", skip(tcx), ret)]
+    pub fn try_resolve(
+        tcx: TyCtxt<'tcx>,
+        typing_env: ty::TypingEnv<'tcx>,
+        def_id: DefId,
+        args: GenericArgsRef<'tcx>,
+    ) -> Result<Option<Instance<'tcx>>, ErrorGuaranteed> {
+        Self::try_resolve_inner(tcx, typing_env, def_id, args, hir::Constness::NotConst)
+    }
+
+    /// See `try_resolve_inner`.
+    #[instrument(level = "debug", skip(tcx), ret)]
+    pub fn try_resolve_for_ctfe(
+        tcx: TyCtxt<'tcx>,
+        typing_env: ty::TypingEnv<'tcx>,
+        def_id: DefId,
+        args: GenericArgsRef<'tcx>,
+    ) -> Result<Option<Instance<'tcx>>, ErrorGuaranteed> {
+        Self::try_resolve_inner(tcx, typing_env, def_id, args, hir::Constness::Const)
     }
 
     pub fn expect_resolve(

compiler/rustc_trait_selection/src/solve/select.rs

@@ -8,7 +8,7 @@ use rustc_infer::traits::{
     Selection, SelectionError, SelectionResult, TraitObligation,
 };
 use rustc_macros::extension;
-use rustc_middle::{bug, span_bug};
+use rustc_middle::{bug, span_bug, ty};
 use rustc_span::Span;
 use thin_vec::thin_vec;
 
@@ -30,6 +30,19 @@ impl<'tcx> InferCtxt<'tcx> {
         .break_value()
         .unwrap()
     }
+    fn select_host_effect_predicate_in_new_trait_solver(
+        &self,
+        obligation: &Obligation<'tcx, ty::HostEffectPredicate<'tcx>>,
+    ) -> SelectionResult<'tcx, Selection<'tcx>> {
+        assert!(self.next_trait_solver());
+
+        self.visit_proof_tree(
+            Goal::new(self.tcx, obligation.param_env, ty::Binder::dummy(obligation.predicate)),
+            &mut Select { span: obligation.cause.span },
+        )
+        .break_value()
+        .unwrap()
+    }
 }
 
 struct Select {

compiler/rustc_ty_utils/src/instance.rs

@@ -1,7 +1,10 @@
 use rustc_errors::ErrorGuaranteed;
-use rustc_hir::LangItem;
 use rustc_hir::def_id::DefId;
+use rustc_hir::{Constness, LangItem};
 use rustc_infer::infer::TyCtxtInferExt;
+use rustc_infer::traits::{
+    ImplSource, Obligation, ObligationCause, ScrubbedTraitError, SelectionError,
+};
 use rustc_middle::bug;
 use rustc_middle::query::Providers;
 use rustc_middle::traits::{BuiltinImplSource, CodegenObligationError};
@@ -10,17 +13,19 @@ use rustc_middle::ty::{
     Unnormalized,
 };
 use rustc_span::sym;
-use rustc_trait_selection::traits;
+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;
 
 use crate::errors::UnexpectedFnPtrAssociatedItem;
 
 fn resolve_instance_raw<'tcx>(
     tcx: TyCtxt<'tcx>,
-    key: ty::PseudoCanonicalInput<'tcx, (DefId, GenericArgsRef<'tcx>)>,
+    key: ty::PseudoCanonicalInput<'tcx, (DefId, GenericArgsRef<'tcx>, Constness)>,
 ) -> Result<Option<Instance<'tcx>>, ErrorGuaranteed> {
-    let PseudoCanonicalInput { typing_env, value: (def_id, args) } = key;
+    let PseudoCanonicalInput { typing_env, value: (def_id, args, constness) } = key;
 
     let result = if let Some(trait_def_id) = tcx.trait_of_assoc(def_id) {
         debug!(" => associated item, attempting to find impl in typing_env {:#?}", typing_env);
@@ -30,6 +35,7 @@ fn resolve_instance_raw<'tcx>(
             typing_env,
             trait_def_id,
             tcx.normalize_erasing_regions(typing_env, Unnormalized::new_wip(args)),
+            constness,
         )
     } else {
         let def = if tcx.intrinsic(def_id).is_some() {
@@ -101,24 +107,88 @@ fn resolve_instance_raw<'tcx>(
     result
 }
 
+#[tracing::instrument(level = "debug", skip(tcx))]
 fn resolve_associated_item<'tcx>(
     tcx: TyCtxt<'tcx>,
     trait_item_id: DefId,
     typing_env: ty::TypingEnv<'tcx>,
     trait_id: DefId,
     rcvr_args: GenericArgsRef<'tcx>,
+    constness: Constness,
 ) -> Result<Option<Instance<'tcx>>, ErrorGuaranteed> {
-    debug!(?trait_item_id, ?typing_env, ?trait_id, ?rcvr_args, "resolve_associated_item");
-
     let trait_ref = ty::TraitRef::from_assoc(tcx, trait_id, rcvr_args);
 
     let input = typing_env.as_query_input(trait_ref);
-    let vtbl = match tcx.codegen_select_candidate(input) {
-        Ok(vtbl) => vtbl,
-        Err(CodegenObligationError::Ambiguity | CodegenObligationError::Unimplemented) => {
+    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);
+        });
+
+        // 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);
+                }
+            }
             return Ok(None);
         }
-        Err(CodegenObligationError::UnconstrainedParam(guar)) => return Err(guar),
+
+        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),
+        }
     };
 
     // Now that we know which impl is being used, we can dispatch to

tests/ui/traits/const-traits/specialization/dont-run-runtime-code.rs

@@ -0,0 +1,43 @@
+//@ compile-flags: -Znext-solver
+//@ check-pass
+
+#![feature(min_specialization, const_trait_impl)]
+
+struct Dummy;
+
+const trait DummyTrait {
+    fn dummy_fn() -> u32;
+}
+impl DummyTrait for Wrap<i32> {
+    fn dummy_fn() -> u32 {
+        println!("wut");
+        0
+    }
+}
+
+const trait Trait {
+    fn trait_fn() -> u32;
+}
+impl<T> const Trait for T where T: DummyTrait {
+    default fn trait_fn() -> u32 {
+        42
+    }
+}
+
+struct Wrap<T>(T);
+
+impl<T> const Trait for Wrap<T> where Self: [const] DummyTrait {
+    fn trait_fn() -> u32 {
+        <Wrap<T>>::dummy_fn()
+    }
+}
+
+const fn indirect<T: DummyTrait>() -> u32 {
+    T::trait_fn()
+}
+
+const A: u32 = indirect::<Wrap<i32>>();
+
+const B: () = { assert!(A == 42); };
+
+fn main() {}