Auto merge of #66866 - oli-obk:const_fn_memoization, r=RalfJung

Only memoize const fn calls during const eval

Miri and other engines may want to execute the function in order to detect UB inside of them.

r? @RalfJung

src/librustc_mir/const_eval.rs

@@ -328,20 +328,32 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
         false // for now, we don't enforce validity
     }
 
-    fn find_fn(
+    fn find_mir_or_eval_fn(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
         instance: ty::Instance<'tcx>,
         args: &[OpTy<'tcx>],
         ret: Option<(PlaceTy<'tcx>, mir::BasicBlock)>,
         _unwind: Option<mir::BasicBlock> // unwinding is not supported in consts
     ) -> InterpResult<'tcx, Option<&'mir mir::Body<'tcx>>> {
-        debug!("eval_fn_call: {:?}", instance);
+        debug!("find_mir_or_eval_fn: {:?}", instance);
+
         // Only check non-glue functions
         if let ty::InstanceDef::Item(def_id) = instance.def {
             // 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_id) {
+            if ecx.tcx.is_const_fn_raw(def_id) {
+                // If this function is a `const fn` then as an optimization we can query this
+                // evaluation immediately.
+                //
+                // For the moment we only do this for functions which take no arguments
+                // (or all arguments are ZSTs) so that we don't memoize too much.
+                if args.iter().all(|a| a.layout.is_zst()) {
+                    let gid = GlobalId { instance, promoted: None };
+                    ecx.eval_const_fn_call(gid, ret)?;
+                    return Ok(None);
+                }
+            } else {
                 // Some functions we support even if they are non-const -- but avoid testing
                 // that for const fn!  We certainly do *not* want to actually call the fn
                 // though, so be sure we return here.

src/librustc_mir/interpret/machine.rs

@@ -146,7 +146,7 @@ pub trait Machine<'mir, 'tcx>: Sized {
     /// nor just jump to `ret`, but instead push their own stack frame.)
     /// Passing `dest`and `ret` in the same `Option` proved very annoying when only one of them
     /// was used.
-    fn find_fn(
+    fn find_mir_or_eval_fn(
         ecx: &mut InterpCx<'mir, 'tcx, Self>,
         instance: ty::Instance<'tcx>,
         args: &[OpTy<'tcx, Self::PointerTag>],

src/librustc_mir/interpret/terminator.rs

@@ -266,20 +266,8 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             ty::InstanceDef::DropGlue(..) |
             ty::InstanceDef::CloneShim(..) |
             ty::InstanceDef::Item(_) => {
-                // If this function is a `const fn` then as an optimization we can query this
-                // evaluation immediately.
-                //
-                // For the moment we only do this for functions which take no arguments
-                // (or all arguments are ZSTs) so that we don't memoize too much.
-                if self.tcx.is_const_fn_raw(instance.def.def_id()) &&
-                   args.iter().all(|a| a.layout.is_zst())
-                {
-                    let gid = GlobalId { instance, promoted: None };
-                    return self.eval_const_fn_call(gid, ret);
-                }
-
                 // We need MIR for this fn
-                let body = match M::find_fn(self, instance, args, ret, unwind)? {
+                let body = match M::find_mir_or_eval_fn(self, instance, args, ret, unwind)? {
                     Some(body) => body,
                     None => return Ok(()),
                 };
@@ -445,7 +433,8 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 
     /// Evaluate a const function where all arguments (if any) are zero-sized types.
     /// The evaluation is memoized thanks to the query system.
-    fn eval_const_fn_call(
+    // FIXME: Consider moving this to `const_eval.rs`.
+    pub (crate) fn eval_const_fn_call(
         &mut self,
         gid: GlobalId<'tcx>,
         ret: Option<(PlaceTy<'tcx, M::PointerTag>, mir::BasicBlock)>,

src/librustc_mir/transform/const_prop.rs

@@ -128,7 +128,7 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine {
         false
     }
 
-    fn find_fn(
+    fn find_mir_or_eval_fn(
         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
         _instance: ty::Instance<'tcx>,
         _args: &[OpTy<'tcx>],