Auto merge of #106227 - bryangarza:ctfe-limit, r=oli-obk

Use stable metric for const eval limit instead of current terminator-based logic

This patch adds a `MirPass` that inserts a new MIR instruction `ConstEvalCounter` to any loops and function calls in the CFG. This instruction is used during Const Eval to count against the `const_eval_limit`, and emit the `StepLimitReached` error, replacing the current logic which uses Terminators only.

The new method of counting loops and function calls should be more stable across compiler versions (i.e., not cause crates that compiled successfully before, to no longer compile when changes to the MIR generation/optimization are made).

Also see: #103877

compiler/rustc_mir_transform/src/check_unsafety.rs

@@ -104,6 +104,7 @@ impl<'tcx> Visitor<'tcx> for UnsafetyChecker<'_, 'tcx> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::Intrinsic(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {
                 // safe (at least as emitted during MIR construction)
             }

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -802,6 +802,8 @@ pub(super) fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span>
         | StatementKind::StorageDead(_)
         // Coverage should not be encountered, but don't inject coverage coverage
         | StatementKind::Coverage(_)
+        // Ignore `ConstEvalCounter`s
+        | StatementKind::ConstEvalCounter
         // Ignore `Nop`s
         | StatementKind::Nop => None,
 

compiler/rustc_mir_transform/src/ctfe_limit.rs

@@ -0,0 +1,59 @@
+//! A pass that inserts the `ConstEvalCounter` instruction into any blocks that have a back edge
+//! (thus indicating there is a loop in the CFG), or whose terminator is a function call.
+use crate::MirPass;
+
+use rustc_data_structures::graph::dominators::Dominators;
+use rustc_middle::mir::{
+    BasicBlock, BasicBlockData, Body, Statement, StatementKind, TerminatorKind,
+};
+use rustc_middle::ty::TyCtxt;
+
+pub struct CtfeLimit;
+
+impl<'tcx> MirPass<'tcx> for CtfeLimit {
+    #[instrument(skip(self, _tcx, body))]
+    fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let doms = body.basic_blocks.dominators();
+        let indices: Vec<BasicBlock> = body
+            .basic_blocks
+            .iter_enumerated()
+            .filter_map(|(node, node_data)| {
+                if matches!(node_data.terminator().kind, TerminatorKind::Call { .. })
+                    // Back edges in a CFG indicate loops
+                    || has_back_edge(&doms, node, &node_data)
+                {
+                    Some(node)
+                } else {
+                    None
+                }
+            })
+            .collect();
+        for index in indices {
+            insert_counter(
+                body.basic_blocks_mut()
+                    .get_mut(index)
+                    .expect("basic_blocks index {index} should exist"),
+            );
+        }
+    }
+}
+
+fn has_back_edge(
+    doms: &Dominators<BasicBlock>,
+    node: BasicBlock,
+    node_data: &BasicBlockData<'_>,
+) -> bool {
+    if !doms.is_reachable(node) {
+        return false;
+    }
+    // Check if any of the dominators of the node are also the node's successor.
+    doms.dominators(node)
+        .any(|dom| node_data.terminator().successors().into_iter().any(|succ| succ == dom))
+}
+
+fn insert_counter(basic_block_data: &mut BasicBlockData<'_>) {
+    basic_block_data.statements.push(Statement {
+        source_info: basic_block_data.terminator().source_info,
+        kind: StatementKind::ConstEvalCounter,
+    });
+}

compiler/rustc_mir_transform/src/dead_store_elimination.rs

@@ -53,6 +53,7 @@ pub fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, borrowed: &BitS
                 | StatementKind::StorageDead(_)
                 | StatementKind::Coverage(_)
                 | StatementKind::Intrinsic(_)
+                | StatementKind::ConstEvalCounter
                 | StatementKind::Nop => (),
 
                 StatementKind::FakeRead(_) | StatementKind::AscribeUserType(_, _) => {

compiler/rustc_mir_transform/src/dest_prop.rs

@@ -577,6 +577,7 @@ impl WriteInfo {
                 self.add_place(**place);
             }
             StatementKind::Intrinsic(_)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop
             | StatementKind::Coverage(_)
             | StatementKind::StorageLive(_)

compiler/rustc_mir_transform/src/generator.rs

@@ -1657,6 +1657,7 @@ impl<'tcx> Visitor<'tcx> for EnsureGeneratorFieldAssignmentsNeverAlias<'_> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::Intrinsic(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
         }
     }

compiler/rustc_mir_transform/src/lib.rs

@@ -55,6 +55,7 @@ mod const_goto;
 mod const_prop;
 mod const_prop_lint;
 mod coverage;
+mod ctfe_limit;
 mod dataflow_const_prop;
 mod dead_store_elimination;
 mod deaggregator;
@@ -410,6 +411,8 @@ fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -
         }
     }
 
+    pm::run_passes(tcx, &mut body, &[&ctfe_limit::CtfeLimit], None);
+
     debug_assert!(!body.has_free_regions(), "Free regions in MIR for CTFE");
 
     body

compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs

@@ -35,6 +35,7 @@ impl RemoveNoopLandingPads {
                 | StatementKind::StorageDead(_)
                 | StatementKind::AscribeUserType(..)
                 | StatementKind::Coverage(..)
+                | StatementKind::ConstEvalCounter
                 | StatementKind::Nop => {
                     // These are all noops in a landing pad
                 }

compiler/rustc_mir_transform/src/separate_const_switch.rs

@@ -250,6 +250,7 @@ fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<
             | StatementKind::Coverage(_)
             | StatementKind::StorageDead(_)
             | StatementKind::Intrinsic(_)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
         }
     }
@@ -318,6 +319,7 @@ fn find_determining_place<'tcx>(
             | StatementKind::AscribeUserType(_, _)
             | StatementKind::Coverage(_)
             | StatementKind::Intrinsic(_)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
 
             // If the discriminant is set, it is always set

compiler/rustc_mir_transform/src/simplify.rs

@@ -517,7 +517,7 @@ impl<'tcx> Visitor<'tcx> for UsedLocals {
                 self.super_statement(statement, location);
             }
 
-            StatementKind::Nop => {}
+            StatementKind::ConstEvalCounter | StatementKind::Nop => {}
 
             StatementKind::StorageLive(_local) | StatementKind::StorageDead(_local) => {}