Auto merge of #89247 - fee1-dead:const-eval-select, r=oli-obk

Add `const_eval_select` intrinsic

Adds an intrinsic that calls a given function when evaluated at compiler time, but generates a call to another function when called at runtime.

See rust-lang/const-eval#7 for previous discussion.

r? `@oli-obk.`

compiler/rustc_codegen_cranelift/src/abi/mod.rs

@@ -309,13 +309,13 @@ pub(crate) fn codegen_terminator_call<'tcx>(
     span: Span,
     func: &Operand<'tcx>,
     args: &[Operand<'tcx>],
-    destination: Option<(Place<'tcx>, BasicBlock)>,
+    mir_dest: Option<(Place<'tcx>, BasicBlock)>,
 ) {
     let fn_ty = fx.monomorphize(func.ty(fx.mir, fx.tcx));
     let fn_sig =
         fx.tcx.normalize_erasing_late_bound_regions(ParamEnv::reveal_all(), fn_ty.fn_sig(fx.tcx));
 
-    let destination = destination.map(|(place, bb)| (codegen_place(fx, place), bb));
+    let destination = mir_dest.map(|(place, bb)| (codegen_place(fx, place), bb));
 
     // Handle special calls like instrinsics and empty drop glue.
     let instance = if let ty::FnDef(def_id, substs) = *fn_ty.kind() {

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -407,11 +407,9 @@ pub(crate) fn codegen_intrinsic_call<'tcx>(
     destination: Option<(CPlace<'tcx>, BasicBlock)>,
     span: Span,
 ) {
-    let def_id = instance.def_id();
+    let intrinsic = fx.tcx.item_name(instance.def_id());
     let substs = instance.substs;
 
-    let intrinsic = fx.tcx.item_name(def_id);
-
     let ret = match destination {
         Some((place, _)) => place,
         None => {

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -26,14 +26,35 @@ impl<'mir, 'tcx> InterpCx<'mir, 'tcx, CompileTimeInterpreter<'mir, 'tcx>> {
     /// "Intercept" a function call to a panic-related function
     /// because we have something special to do for it.
     /// If this returns successfully (`Ok`), the function should just be evaluated normally.
-    fn hook_panic_fn(
+    fn hook_special_const_fn(
         &mut self,
         instance: ty::Instance<'tcx>,
         args: &[OpTy<'tcx>],
+        is_const_fn: bool,
     ) -> InterpResult<'tcx, Option<ty::Instance<'tcx>>> {
         // The list of functions we handle here must be in sync with
-        // `is_lang_panic_fn` in `transform/check_consts/mod.rs`.
+        // `is_lang_special_const_fn` in `transform/check_consts/mod.rs`.
         let def_id = instance.def_id();
+
+        if is_const_fn {
+            if Some(def_id) == self.tcx.lang_items().const_eval_select() {
+                // redirect to const_eval_select_ct
+                if let Some(const_eval_select) = self.tcx.lang_items().const_eval_select_ct() {
+                    return Ok(Some(
+                        ty::Instance::resolve(
+                            *self.tcx,
+                            ty::ParamEnv::reveal_all(),
+                            const_eval_select,
+                            instance.substs,
+                        )
+                        .unwrap()
+                        .unwrap(),
+                    ));
+                }
+            }
+            return Ok(None);
+        }
+
         if Some(def_id) == self.tcx.lang_items().panic_fn()
             || Some(def_id) == self.tcx.lang_items().panic_str()
             || Some(def_id) == self.tcx.lang_items().panic_display()
@@ -255,31 +276,31 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
 
         // Only check non-glue functions
         if let ty::InstanceDef::Item(def) = instance.def {
+            let mut is_const_fn = true;
+
             // Execution might have wandered off into other crates, so we cannot do a stability-
             // sensitive check here.  But we can at least rule out functions that are not const
             // at all.
             if !ecx.tcx.is_const_fn_raw(def.did) {
                 // allow calling functions marked with #[default_method_body_is_const].
                 if !ecx.tcx.has_attr(def.did, sym::default_method_body_is_const) {
-                    // Some functions we support even if they are non-const -- but avoid testing
-                    // that for const fn!
-                    if let Some(new_instance) = ecx.hook_panic_fn(instance, args)? {
-                        // We call another const fn instead.
-                        return Self::find_mir_or_eval_fn(
-                            ecx,
-                            new_instance,
-                            _abi,
-                            args,
-                            _ret,
-                            _unwind,
-                        );
-                    } else {
-                        // We certainly do *not* want to actually call the fn
-                        // though, so be sure we return here.
-                        throw_unsup_format!("calling non-const function `{}`", instance)
-                    }
+                    is_const_fn = false;
                 }
             }
+
+            // Some functions we support even if they are non-const -- but avoid testing
+            // that for const fn!
+            // `const_eval_select` is a const fn because it must use const trait bounds.
+            if let Some(new_instance) = ecx.hook_special_const_fn(instance, args, is_const_fn)? {
+                // We call another const fn instead.
+                return Self::find_mir_or_eval_fn(ecx, new_instance, _abi, args, _ret, _unwind);
+            }
+
+            if !is_const_fn {
+                // We certainly do *not* want to actually call the fn
+                // though, so be sure we return here.
+                throw_unsup_format!("calling non-const function `{}`", instance)
+            }
         }
         // This is a const fn. Call it.
         Ok(Some(ecx.load_mir(instance.def, None)?))

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -231,7 +231,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     }
 
     /// Call this function -- pushing the stack frame and initializing the arguments.
-    fn eval_fn_call(
+    pub(crate) fn eval_fn_call(
         &mut self,
         fn_val: FnVal<'tcx, M::ExtraFnVal>,
         caller_abi: Abi,

compiler/rustc_const_eval/src/transform/check_consts/check.rs

@@ -24,7 +24,7 @@ use std::ops::Deref;
 use super::ops::{self, NonConstOp, Status};
 use super::qualifs::{self, CustomEq, HasMutInterior, NeedsNonConstDrop};
 use super::resolver::FlowSensitiveAnalysis;
-use super::{is_lang_panic_fn, ConstCx, Qualif};
+use super::{is_lang_panic_fn, is_lang_special_const_fn, ConstCx, Qualif};
 use crate::const_eval::is_unstable_const_fn;
 
 // We are using `MaybeMutBorrowedLocals` as a proxy for whether an item may have been mutated
@@ -259,7 +259,9 @@ impl Checker<'mir, 'tcx> {
             self.check_local_or_return_ty(return_ty.skip_binder(), RETURN_PLACE);
         }
 
-        self.visit_body(&body);
+        if !tcx.has_attr(def_id.to_def_id(), sym::rustc_do_not_const_check) {
+            self.visit_body(&body);
+        }
 
         // Ensure that the end result is `Sync` in a non-thread local `static`.
         let should_check_for_sync = self.const_kind()
@@ -886,7 +888,7 @@ impl Visitor<'tcx> for Checker<'mir, 'tcx> {
                 }
 
                 // At this point, we are calling a function, `callee`, whose `DefId` is known...
-                if is_lang_panic_fn(tcx, callee) {
+                if is_lang_special_const_fn(tcx, callee) {
                     // `begin_panic` and `panic_display` are generic functions that accept
                     // types other than str. Check to enforce that only str can be used in
                     // const-eval.
@@ -908,7 +910,10 @@ impl Visitor<'tcx> for Checker<'mir, 'tcx> {
                         }
                     }
 
-                    return;
+                    if is_lang_panic_fn(tcx, callee) {
+                        // run stability check on non-panic special const fns.
+                        return;
+                    }
                 }
 
                 if Some(callee) == tcx.lang_items().exchange_malloc_fn() {

compiler/rustc_const_eval/src/transform/check_consts/mod.rs

@@ -74,9 +74,6 @@ impl ConstCx<'mir, 'tcx> {
 
 /// Returns `true` if this `DefId` points to one of the official `panic` lang items.
 pub fn is_lang_panic_fn(tcx: TyCtxt<'tcx>, def_id: DefId) -> bool {
-    // We can allow calls to these functions because `hook_panic_fn` in
-    // `const_eval/machine.rs` ensures the calls are handled specially.
-    // Keep in sync with what that function handles!
     Some(def_id) == tcx.lang_items().panic_fn()
         || Some(def_id) == tcx.lang_items().panic_str()
         || Some(def_id) == tcx.lang_items().panic_display()
@@ -85,6 +82,15 @@ pub fn is_lang_panic_fn(tcx: TyCtxt<'tcx>, def_id: DefId) -> bool {
         || Some(def_id) == tcx.lang_items().begin_panic_fmt()
 }
 
+/// Returns `true` if this `DefId` points to one of the lang items that will be handled differently
+/// in const_eval.
+pub fn is_lang_special_const_fn(tcx: TyCtxt<'tcx>, def_id: DefId) -> bool {
+    // We can allow calls to these functions because `hook_special_const_fn` in
+    // `const_eval/machine.rs` ensures the calls are handled specially.
+    // Keep in sync with what that function handles!
+    is_lang_panic_fn(tcx, def_id) || Some(def_id) == tcx.lang_items().const_eval_select()
+}
+
 pub fn rustc_allow_const_fn_unstable(
     tcx: TyCtxt<'tcx>,
     def_id: DefId,

compiler/rustc_const_eval/src/transform/check_consts/post_drop_elaboration.rs

@@ -1,7 +1,7 @@
 use rustc_middle::mir::visit::Visitor;
 use rustc_middle::mir::{self, BasicBlock, Location};
 use rustc_middle::ty::TyCtxt;
-use rustc_span::Span;
+use rustc_span::{symbol::sym, Span};
 
 use super::check::Qualifs;
 use super::ops::{self, NonConstOp};
@@ -30,6 +30,10 @@ pub fn check_live_drops(tcx: TyCtxt<'tcx>, body: &mir::Body<'tcx>) {
         return;
     }
 
+    if tcx.has_attr(def_id.to_def_id(), sym::rustc_do_not_const_check) {
+        return;
+    }
+
     let ccx = ConstCx { body, tcx, const_kind, param_env: tcx.param_env(def_id) };
     if !checking_enabled(&ccx) {
         return;

compiler/rustc_const_eval/src/transform/promote_consts.rs

@@ -26,7 +26,7 @@ use rustc_index::vec::{Idx, IndexVec};
 use std::cell::Cell;
 use std::{cmp, iter, mem};
 
-use crate::transform::check_consts::{is_lang_panic_fn, qualifs, ConstCx};
+use crate::transform::check_consts::{is_lang_special_const_fn, qualifs, ConstCx};
 use crate::transform::MirPass;
 
 /// A `MirPass` for promotion.
@@ -657,7 +657,7 @@ impl<'tcx> Validator<'_, 'tcx> {
 
         let is_const_fn = match *fn_ty.kind() {
             ty::FnDef(def_id, _) => {
-                self.tcx.is_const_fn_raw(def_id) || is_lang_panic_fn(self.tcx, def_id)
+                self.tcx.is_const_fn_raw(def_id) || is_lang_special_const_fn(self.tcx, def_id)
             }
             _ => false,
         };

compiler/rustc_feature/src/builtin_attrs.rs

@@ -467,6 +467,8 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
 
     rustc_attr!(rustc_promotable, Normal, template!(Word), IMPL_DETAIL),
     rustc_attr!(rustc_legacy_const_generics, Normal, template!(List: "N"), INTERNAL_UNSTABLE),
+    // Do not const-check this function's body. It will always get replaced during CTFE.
+    rustc_attr!(rustc_do_not_const_check, Normal, template!(Word), INTERNAL_UNSTABLE),
 
     // ==========================================================================
     // Internal attributes, Layout related:

compiler/rustc_hir/src/lang_items.rs

@@ -299,6 +299,8 @@ language_item_table! {
     DropInPlace,             sym::drop_in_place,       drop_in_place_fn,           Target::Fn,             GenericRequirement::Minimum(1);
     Oom,                     sym::oom,                 oom,                        Target::Fn,             GenericRequirement::None;
     AllocLayout,             sym::alloc_layout,        alloc_layout,               Target::Struct,         GenericRequirement::None;
+    ConstEvalSelect,         sym::const_eval_select,   const_eval_select,          Target::Fn,             GenericRequirement::Exact(4);
+    ConstConstEvalSelect,    sym::const_eval_select_ct,const_eval_select_ct,       Target::Fn,             GenericRequirement::Exact(4);
 
     Start,                   sym::start,               start_fn,                   Target::Fn,             GenericRequirement::Exact(1);
 

compiler/rustc_span/src/symbol.rs

@@ -441,6 +441,8 @@ symbols! {
         const_compare_raw_pointers,
         const_constructor,
         const_eval_limit,
+        const_eval_select,
+        const_eval_select_ct,
         const_evaluatable_checked,
         const_extern_fn,
         const_fn,
@@ -1097,6 +1099,7 @@ symbols! {
         rustc_diagnostic_item,
         rustc_diagnostic_macros,
         rustc_dirty,
+        rustc_do_not_const_check,
         rustc_dummy,
         rustc_dump_env_program_clauses,
         rustc_dump_program_clauses,

compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs

@@ -973,12 +973,16 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 ty::Tuple(_) => stack.extend(ty.tuple_fields().map(|t| (t, depth + 1))),
 
                 ty::Closure(_, substs) => {
-                    stack.extend(substs.as_closure().upvar_tys().map(|t| (t, depth + 1)))
+                    let substs = substs.as_closure();
+                    let ty = self.infcx.shallow_resolve(substs.tupled_upvars_ty());
+                    stack.push((ty, depth + 1));
                 }
 
                 ty::Generator(_, substs, _) => {
                     let substs = substs.as_generator();
-                    stack.extend(substs.upvar_tys().map(|t| (t, depth + 1)));
+                    let ty = self.infcx.shallow_resolve(substs.tupled_upvars_ty());
+
+                    stack.push((ty, depth + 1));
                     stack.push((substs.witness(), depth + 1));
                 }
 

compiler/rustc_typeck/src/check/callee.rs

@@ -17,7 +17,7 @@ use rustc_infer::{
 use rustc_middle::ty::adjustment::{
     Adjust, Adjustment, AllowTwoPhase, AutoBorrow, AutoBorrowMutability,
 };
-use rustc_middle::ty::subst::SubstsRef;
+use rustc_middle::ty::subst::{Subst, SubstsRef};
 use rustc_middle::ty::{self, Ty, TyCtxt, TypeFoldable};
 use rustc_span::symbol::{sym, Ident};
 use rustc_span::Span;
@@ -317,6 +317,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     ) -> Ty<'tcx> {
         let (fn_sig, def_id) = match *callee_ty.kind() {
             ty::FnDef(def_id, subst) => {
+                let fn_sig = self.tcx.fn_sig(def_id).subst(self.tcx, subst);
+
                 // Unit testing: function items annotated with
                 // `#[rustc_evaluate_where_clauses]` trigger special output
                 // to let us test the trait evaluation system.
@@ -342,7 +344,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                             .emit();
                     }
                 }
-                (callee_ty.fn_sig(self.tcx), Some(def_id))
+                (fn_sig, Some(def_id))
             }
             ty::FnPtr(sig) => (sig, None),
             ref t => {

compiler/rustc_typeck/src/check/intrinsic.rs

@@ -390,6 +390,8 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
 
             sym::black_box => (1, vec![param(0)], param(0)),
 
+            sym::const_eval_select => (4, vec![param(0), param(1), param(2)], param(3)),
+
             other => {
                 tcx.sess.emit_err(UnrecognizedIntrinsicFunction { span: it.span, name: other });
                 return;

library/core/src/intrinsics.rs

@@ -2221,3 +2221,74 @@ pub unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize) {
     // SAFETY: the safety contract for `write_bytes` must be upheld by the caller.
     unsafe { write_bytes(dst, val, count) }
 }
+
+/// Selects which function to call depending on the context.
+///
+/// If this function is evaluated at compile-time, then a call to this
+/// intrinsic will be replaced with a call to `called_in_const`. It gets
+/// replaced with a call to `called_at_rt` otherwise.
+///
+/// # Type Requirements
+///
+/// The two functions must be both function items. They cannot be function
+/// pointers or closures.
+///
+/// `arg` will be the arguments that will be passed to either one of the
+/// two functions, therefore, both functions must accept the same type of
+/// arguments. Both functions must return RET.
+///
+/// # Safety
+///
+/// This intrinsic allows breaking [referential transparency] in `const fn`
+/// and is therefore `unsafe`.
+///
+/// Code that uses this intrinsic must be extremely careful to ensure that
+/// `const fn`s remain referentially-transparent independently of when they
+/// are evaluated.
+///
+/// The Rust compiler assumes that it is sound to replace a call to a `const
+/// fn` with the result produced by evaluating it at compile-time. If
+/// evaluating the function at run-time were to produce a different result,
+/// or have any other observable side-effects, the behavior is undefined.
+///
+/// [referential transparency]: https://en.wikipedia.org/wiki/Referential_transparency
+#[cfg(not(bootstrap))]
+#[unstable(
+    feature = "const_eval_select",
+    issue = "none",
+    reason = "const_eval_select will never be stable"
+)]
+#[rustc_const_unstable(feature = "const_eval_select", issue = "none")]
+#[lang = "const_eval_select"]
+#[rustc_do_not_const_check]
+pub const unsafe fn const_eval_select<ARG, F, G, RET>(
+    arg: ARG,
+    _called_in_const: F,
+    called_at_rt: G,
+) -> RET
+where
+    F: ~const FnOnce<ARG, Output = RET>,
+    G: FnOnce<ARG, Output = RET> + ~const Drop,
+{
+    called_at_rt.call_once(arg)
+}
+
+#[cfg(not(bootstrap))]
+#[unstable(
+    feature = "const_eval_select",
+    issue = "none",
+    reason = "const_eval_select will never be stable"
+)]
+#[rustc_const_unstable(feature = "const_eval_select", issue = "none")]
+#[lang = "const_eval_select_ct"]
+pub const unsafe fn const_eval_select_ct<ARG, F, G, RET>(
+    arg: ARG,
+    called_in_const: F,
+    _called_at_rt: G,
+) -> RET
+where
+    F: ~const FnOnce<ARG, Output = RET>,
+    G: FnOnce<ARG, Output = RET> + ~const Drop,
+{
+    called_in_const.call_once(arg)
+}