Auto merge of rust-lang#121557 - RalfJung:const-fn-call-promotion, r=<try>

restrict promotion of `const fn` calls

We only promote them in `const`/`static` initializers, but even that is still unfortunate -- we still cannot add promoteds to required_consts. But we should add them there to make sure it's always okay to evaluate every const we encounter in a MIR body.  That effort of not promoting things that can fail to evaluate is tracked in rust-lang#80619. These `const fn` calls are the last missing piece.

So I propose that we do not promote const-fn calls in const when that may fail without the entire const failing, thereby completing rust-lang#80619. Unfortunately we can't just reject promoting these functions outright due to backwards compatibility. So let's see if we can find a hack that makes crater happy...

For the record, this is the [crater analysis](rust-lang#80243 (comment)) from when I tried to entirely forbid this kind of promotion. It's a tiny amount of breakage and if we had a nice alternative for code like that, we could conceivably push it through... but sadly, inline const expressions are still blocked on t-lang concerns about post-monomorphization errors and we haven't yet figured out an implementation that can resolve those concerns. So we're forced to make progress via other means, such as terrible hacks like this.

Attempt one: only promote calls on the "safe path" at the beginning of a MIR block. This is the path that starts at the start block and continues via gotos and calls, but stops at the first branch. If we had imposed this restriction before stabilizing `if` and `match` in `const`, this would have definitely been sufficient...

EDIT: Turns out that works. :)
**Here's the t-lang [nomination comment](rust-lang#121557 (comment)

r? `@oli-obk`

Author: bors

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -375,6 +375,20 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
 
     const PANIC_ON_ALLOC_FAIL: bool = false; // will be raised as a proper error
 
+    #[inline]
+    fn all_required_consts_are_checked(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool {
+        // Generally we expect required_consts to be properly filled, except for promoteds where
+        // storing these consts shows up negatively in benchmarks. A promoted can only be relevant
+        // if its parent MIR is relevant, and the consts in the promoted will be in the parent's
+        // `required_consts`, so we are still sure to catch any const-eval bugs, just a bit less
+        // directly.
+        if ecx.frame_idx() == 0 && ecx.frame().body.source.promoted.is_some() {
+            false
+        } else {
+            true
+        }
+    }
+
     #[inline(always)]
     fn enforce_alignment(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool {
         matches!(ecx.machine.check_alignment, CheckAlignment::Error)

compiler/rustc_const_eval/src/interpret/eval_context.rs

@@ -796,15 +796,13 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         self.stack_mut().push(frame);
 
         // Make sure all the constants required by this frame evaluate successfully (post-monomorphization check).
-        if M::POST_MONO_CHECKS {
-            for &const_ in &body.required_consts {
-                let c = self
-                    .instantiate_from_current_frame_and_normalize_erasing_regions(const_.const_)?;
-                c.eval(*self.tcx, self.param_env, Some(const_.span)).map_err(|err| {
-                    err.emit_note(*self.tcx);
-                    err
-                })?;
-            }
+        for &const_ in &body.required_consts {
+            let c =
+                self.instantiate_from_current_frame_and_normalize_erasing_regions(const_.const_)?;
+            c.eval(*self.tcx, self.param_env, Some(const_.span)).map_err(|err| {
+                err.emit_note(*self.tcx);
+                err
+            })?;
         }
 
         // done
@@ -1153,8 +1151,12 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ) -> InterpResult<'tcx, OpTy<'tcx, M::Provenance>> {
         M::eval_mir_constant(self, *val, span, layout, |ecx, val, span, layout| {
             let const_val = val.eval(*ecx.tcx, ecx.param_env, span).map_err(|err| {
-                // FIXME: somehow this is reachable even when POST_MONO_CHECKS is on.
-                // Are we not always populating `required_consts`?
+                if M::all_required_consts_are_checked(self)
+                    && !matches!(err, ErrorHandled::TooGeneric(..))
+                {
+                    // Looks like the const is not captued by `required_consts`, that's bad.
+                    bug!("interpret const eval failure of {val:?} which is not in required_consts");
+                }
                 err.emit_note(*ecx.tcx);
                 err
             })?;

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -140,8 +140,9 @@ pub trait Machine<'mir, 'tcx: 'mir>: Sized {
     /// Should the machine panic on allocation failures?
     const PANIC_ON_ALLOC_FAIL: bool;
 
-    /// Should post-monomorphization checks be run when a stack frame is pushed?
-    const POST_MONO_CHECKS: bool = true;
+    /// Determines whether `eval_mir_constant` can never fail because all required consts have
+    /// already been checked before.
+    fn all_required_consts_are_checked(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool;
 
     /// Whether memory accesses should be alignment-checked.
     fn enforce_alignment(ecx: &InterpCx<'mir, 'tcx, Self>) -> bool;

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

@@ -900,6 +900,12 @@ impl<'mir, 'tcx: 'mir> rustc_const_eval::interpret::Machine<'mir, 'tcx> for Dumm
     type MemoryKind = !;
     const PANIC_ON_ALLOC_FAIL: bool = true;
 
+    #[inline(always)]
+    fn all_required_consts_are_checked(_ecx: &InterpCx<'mir, 'tcx, Self>) -> bool {
+        // We want to just eval random consts in the program, so `eval_mir_const` can fail.
+        false
+    }
+
     #[inline(always)]
     fn enforce_alignment(_ecx: &InterpCx<'mir, 'tcx, Self>) -> bool {
         false // no reason to enforce alignment

compiler/rustc_mir_transform/src/inline.rs

@@ -706,17 +706,7 @@ impl<'tcx> Inliner<'tcx> {
         });
 
         // Copy only unevaluated constants from the callee_body into the caller_body.
-        // Although we are only pushing `ConstKind::Unevaluated` consts to
-        // `required_consts`, here we may not only have `ConstKind::Unevaluated`
-        // because we are calling `instantiate_and_normalize_erasing_regions`.
-        caller_body.required_consts.extend(callee_body.required_consts.iter().copied().filter(
-            |&ct| match ct.const_ {
-                Const::Ty(_) => {
-                    bug!("should never encounter ty::UnevaluatedConst in `required_consts`")
-                }
-                Const::Val(..) | Const::Unevaluated(..) => true,
-            },
-        ));
+        caller_body.required_consts.extend(callee_body.required_consts);
     }
 
     fn make_call_args(

compiler/rustc_mir_transform/src/lib.rs

@@ -343,6 +343,8 @@ fn mir_promoted(
         body.tainted_by_errors = Some(error_reported);
     }
 
+    // Collect `required_consts` *before* promotion, so if there are any consts being promoted
+    // we still add them to the list in the outer MIR body.
     let mut required_consts = Vec::new();
     let mut required_consts_visitor = RequiredConstsVisitor::new(&mut required_consts);
     for (bb, bb_data) in traversal::reverse_postorder(&body) {

compiler/rustc_mir_transform/src/promote_consts.rs

@@ -13,6 +13,7 @@
 //! move analysis runs after promotion on broken MIR.
 
 use either::{Left, Right};
+use rustc_data_structures::fx::FxHashSet;
 use rustc_hir as hir;
 use rustc_middle::mir;
 use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
@@ -175,6 +176,12 @@ fn collect_temps_and_candidates<'tcx>(
 struct Validator<'a, 'tcx> {
     ccx: &'a ConstCx<'a, 'tcx>,
     temps: &'a mut IndexSlice<Local, TempState>,
+    /// For backwards compatibility, we are promoting function calls in `const`/`static`
+    /// initializers. But we want to avoid evaluating code that might panic and that otherwise would
+    /// not have been evaluated, so we only promote such calls in basic blocks that are guaranteed
+    /// to execute. In other words, we only promote such calls in basic blocks that are definitely
+    /// not dead code. Here we cache the result of computing that set of basic blocks.
+    promotion_safe_blocks: Option<FxHashSet<BasicBlock>>,
 }
 
 impl<'a, 'tcx> std::ops::Deref for Validator<'a, 'tcx> {
@@ -260,7 +267,9 @@ impl<'tcx> Validator<'_, 'tcx> {
                     self.validate_rvalue(rhs)
                 }
                 Right(terminator) => match &terminator.kind {
-                    TerminatorKind::Call { func, args, .. } => self.validate_call(func, args),
+                    TerminatorKind::Call { func, args, .. } => {
+                        self.validate_call(func, args, loc.block)
+                    }
                     TerminatorKind::Yield { .. } => Err(Unpromotable),
                     kind => {
                         span_bug!(terminator.source_info.span, "{:?} not promotable", kind);
@@ -587,53 +596,103 @@ impl<'tcx> Validator<'_, 'tcx> {
         Ok(())
     }
 
+    /// Computes the sets of blocks of this MIR that are definitely going to be executed
+    /// if the function returns successfully. That makes it safe to promote calls in them
+    /// that might fail.
+    fn promotion_safe_blocks(body: &mir::Body<'tcx>) -> FxHashSet<BasicBlock> {
+        let mut safe_blocks = FxHashSet::default();
+        let mut safe_block = START_BLOCK;
+        loop {
+            safe_blocks.insert(safe_block);
+            // Let's see if we can find another safe block.
+            safe_block = match body.basic_blocks[safe_block].terminator().kind {
+                TerminatorKind::Goto { target } => target,
+                TerminatorKind::Call { target: Some(target), .. }
+                | TerminatorKind::Drop { target, .. } => {
+                    // This calls a function or the destructor. `target` does not get executed if
+                    // the callee loops or panics. But in both cases the const already fails to
+                    // evaluate, so we are fine considering `target` a safe block for promotion.
+                    target
+                }
+                TerminatorKind::Assert { target, .. } => {
+                    // Similar to above, we only consider successful execution.
+                    target
+                }
+                _ => {
+                    // No next safe block.
+                    break;
+                }
+            };
+        }
+        safe_blocks
+    }
+
+    /// Returns whether the block is "safe" for promotion, which means it cannot be dead code.
+    /// We use this to avoid promoting operations that can fail in dead code.
+    fn is_promotion_safe_block(&mut self, block: BasicBlock) -> bool {
+        let body = self.body;
+        let safe_blocks =
+            self.promotion_safe_blocks.get_or_insert_with(|| Self::promotion_safe_blocks(body));
+        safe_blocks.contains(&block)
+    }
+
     fn validate_call(
         &mut self,
         callee: &Operand<'tcx>,
         args: &[Spanned<Operand<'tcx>>],
+        block: BasicBlock,
     ) -> Result<(), Unpromotable> {
+        // Validate the operands. If they fail, there's no question -- we cannot promote.
+        self.validate_operand(callee)?;
+        for arg in args {
+            self.validate_operand(&arg.node)?;
+        }
+
+        // Functions marked `#[rustc_promotable]` are explicitly allowed to be promoted, so we can
+        // accept them at this point.
         let fn_ty = callee.ty(self.body, self.tcx);
+        if let ty::FnDef(def_id, _) = *fn_ty.kind() {
+            if self.tcx.is_promotable_const_fn(def_id) {
+                return Ok(());
+            }
+        }
 
-        // Inside const/static items, we promote all (eligible) function calls.
-        // Everywhere else, we require `#[rustc_promotable]` on the callee.
-        let promote_all_const_fn = matches!(
+        // Ideally, we'd stop here and reject the rest.
+        // But for backward compatibility, we have to accept some promotion in const/static
+        // initializers. Inline consts are explicitly excluded, they are more recent so we have no
+        // backwards compatibility reason to allow more promotion inside of them.
+        let promote_all_fn = matches!(
             self.const_kind,
             Some(hir::ConstContext::Static(_) | hir::ConstContext::Const { inline: false })
         );
-        if !promote_all_const_fn {
-            if let ty::FnDef(def_id, _) = *fn_ty.kind() {
-                // Never promote runtime `const fn` calls of
-                // functions without `#[rustc_promotable]`.
-                if !self.tcx.is_promotable_const_fn(def_id) {
-                    return Err(Unpromotable);
-                }
-            }
+        if !promote_all_fn {
+            return Err(Unpromotable);
         }
-
+        // Make sure the callee is a `const fn`.
         let is_const_fn = match *fn_ty.kind() {
             ty::FnDef(def_id, _) => self.tcx.is_const_fn_raw(def_id),
             _ => false,
         };
         if !is_const_fn {
             return Err(Unpromotable);
         }
-
-        self.validate_operand(callee)?;
-        for arg in args {
-            self.validate_operand(&arg.node)?;
+        // The problem is, this may promote calls to functions that panic.
+        // We don't want to introduce compilation errors if there's a panic in a call in dead code.
+        // So we ensure that this is not dead code.
+        if !self.is_promotion_safe_block(block) {
+            return Err(Unpromotable);
         }
-
+        // This passed all checks, so let's accept.
         Ok(())
     }
 }
 
-// FIXME(eddyb) remove the differences for promotability in `static`, `const`, `const fn`.
 fn validate_candidates(
     ccx: &ConstCx<'_, '_>,
     temps: &mut IndexSlice<Local, TempState>,
     candidates: &[Candidate],
 ) -> Vec<Candidate> {
-    let mut validator = Validator { ccx, temps };
+    let mut validator = Validator { ccx, temps, promotion_safe_blocks: None };
 
     candidates
         .iter()
@@ -652,6 +711,10 @@ struct Promoter<'a, 'tcx> {
     /// If true, all nested temps are also kept in the
     /// source MIR, not moved to the promoted MIR.
     keep_original: bool,
+
+    /// If true, add the new const (the promoted) to the required_consts of the parent MIR.
+    /// This is initially false and then set by the visitor when it encounters a `Call` terminator.
+    add_to_required: bool,
 }
 
 impl<'a, 'tcx> Promoter<'a, 'tcx> {
@@ -754,6 +817,10 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
                 TerminatorKind::Call {
                     mut func, mut args, call_source: desugar, fn_span, ..
                 } => {
+                    // This promoted involves a function call, so it may fail to evaluate.
+                    // Let's make sure it is added to `required_consts` so that that failure cannot get lost.
+                    self.add_to_required = true;
+
                     self.visit_operand(&mut func, loc);
                     for arg in &mut args {
                         self.visit_operand(&mut arg.node, loc);
@@ -788,7 +855,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
 
     fn promote_candidate(mut self, candidate: Candidate, next_promoted_id: usize) -> Body<'tcx> {
         let def = self.source.source.def_id();
-        let mut rvalue = {
+        let (mut rvalue, promoted_op) = {
             let promoted = &mut self.promoted;
             let promoted_id = Promoted::new(next_promoted_id);
             let tcx = self.tcx;
@@ -798,11 +865,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
                 let args = tcx.erase_regions(GenericArgs::identity_for_item(tcx, def));
                 let uneval = mir::UnevaluatedConst { def, args, promoted: Some(promoted_id) };
 
-                Operand::Constant(Box::new(ConstOperand {
-                    span,
-                    user_ty: None,
-                    const_: Const::Unevaluated(uneval, ty),
-                }))
+                ConstOperand { span, user_ty: None, const_: Const::Unevaluated(uneval, ty) }
             };
 
             let blocks = self.source.basic_blocks.as_mut();
@@ -838,22 +901,26 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
             let promoted_ref = local_decls.push(promoted_ref);
             assert_eq!(self.temps.push(TempState::Unpromotable), promoted_ref);
 
+            let promoted_operand = promoted_operand(ref_ty, span);
             let promoted_ref_statement = Statement {
                 source_info: statement.source_info,
                 kind: StatementKind::Assign(Box::new((
                     Place::from(promoted_ref),
-                    Rvalue::Use(promoted_operand(ref_ty, span)),
+                    Rvalue::Use(Operand::Constant(Box::new(promoted_operand))),
                 ))),
             };
             self.extra_statements.push((loc, promoted_ref_statement));
 
-            Rvalue::Ref(
-                tcx.lifetimes.re_erased,
-                *borrow_kind,
-                Place {
-                    local: mem::replace(&mut place.local, promoted_ref),
-                    projection: List::empty(),
-                },
+            (
+                Rvalue::Ref(
+                    tcx.lifetimes.re_erased,
+                    *borrow_kind,
+                    Place {
+                        local: mem::replace(&mut place.local, promoted_ref),
+                        projection: List::empty(),
+                    },
+                ),
+                promoted_operand,
             )
         };
 
@@ -865,6 +932,12 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
 
         let span = self.promoted.span;
         self.assign(RETURN_PLACE, rvalue, span);
+
+        // Now that we did promotion, we know whether we'll want to add this to `required_consts`.
+        if self.add_to_required {
+            self.source.required_consts.push(promoted_op);
+        }
+
         self.promoted
     }
 }
@@ -924,6 +997,11 @@ fn promote_candidates<'tcx>(
             None,
             body.tainted_by_errors,
         );
+        // We keep `required_consts` of the new MIR body empty. All consts mentioned here have
+        // already been added to the parent MIR's `required_consts` (that is computed before
+        // promotion), and no matter where this promoted const ends up, our parent MIR must be
+        // somewhere in the reachable dependency chain so we can rely on its required consts being
+        // evaluated.
         promoted.phase = MirPhase::Analysis(AnalysisPhase::Initial);
 
         let promoter = Promoter {
@@ -933,6 +1011,7 @@ fn promote_candidates<'tcx>(
             temps: &mut temps,
             extra_statements: &mut extra_statements,
             keep_original: false,
+            add_to_required: false,
         };
 
         let mut promoted = promoter.promote_candidate(candidate, promotions.len());