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_borrowck/src/dataflow.rs

@@ -393,6 +393,7 @@ impl<'tcx> rustc_mir_dataflow::GenKillAnalysis<'tcx> for Borrows<'_, 'tcx> {
             | mir::StatementKind::AscribeUserType(..)
             | mir::StatementKind::Coverage(..)
             | mir::StatementKind::Intrinsic(..)
+            | mir::StatementKind::ConstEvalCounter
             | mir::StatementKind::Nop => {}
         }
     }

compiler/rustc_borrowck/src/invalidation.rs

@@ -91,7 +91,8 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                     LocalMutationIsAllowed::Yes,
                 );
             }
-            StatementKind::Nop
+            StatementKind::ConstEvalCounter
+            | StatementKind::Nop
             | StatementKind::Retag { .. }
             | StatementKind::Deinit(..)
             | StatementKind::SetDiscriminant { .. } => {

compiler/rustc_borrowck/src/lib.rs

@@ -609,7 +609,8 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
             StatementKind::AscribeUserType(..)
             // Doesn't have any language semantics
             | StatementKind::Coverage(..)
-            // Does not actually affect borrowck
+            // These do not actually affect borrowck
+            | StatementKind::ConstEvalCounter
             | StatementKind::StorageLive(..) => {}
             StatementKind::StorageDead(local) => {
                 self.access_place(

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1258,6 +1258,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             | StatementKind::StorageDead(..)
             | StatementKind::Retag { .. }
             | StatementKind::Coverage(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
             StatementKind::Deinit(..) | StatementKind::SetDiscriminant { .. } => {
                 bug!("Statement not allowed in this MIR phase")

compiler/rustc_codegen_cranelift/src/base.rs

@@ -794,6 +794,7 @@ fn codegen_stmt<'tcx>(
         StatementKind::StorageLive(_)
         | StatementKind::StorageDead(_)
         | StatementKind::Deinit(_)
+        | StatementKind::ConstEvalCounter
         | StatementKind::Nop
         | StatementKind::FakeRead(..)
         | StatementKind::Retag { .. }

compiler/rustc_codegen_cranelift/src/constant.rs

@@ -530,6 +530,7 @@ pub(crate) fn mir_operand_get_const_val<'tcx>(
                         | StatementKind::Retag(_, _)
                         | StatementKind::AscribeUserType(_, _)
                         | StatementKind::Coverage(_)
+                        | StatementKind::ConstEvalCounter
                         | StatementKind::Nop => {}
                     }
                 }

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -91,6 +91,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::StatementKind::FakeRead(..)
             | mir::StatementKind::Retag { .. }
             | mir::StatementKind::AscribeUserType(..)
+            | mir::StatementKind::ConstEvalCounter
             | mir::StatementKind::Nop => {}
         }
     }

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -561,8 +561,8 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
         throw_unsup_format!("pointer arithmetic or comparison is not supported at compile-time");
     }
 
-    fn before_terminator(ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx> {
-        // The step limit has already been hit in a previous call to `before_terminator`.
+    fn increment_const_eval_counter(ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx> {
+        // The step limit has already been hit in a previous call to `increment_const_eval_counter`.
         if ecx.machine.steps_remaining == 0 {
             return Ok(());
         }

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -244,12 +244,18 @@ pub trait Machine<'mir, 'tcx>: Sized {
     }
 
     /// Called before a basic block terminator is executed.
-    /// You can use this to detect endlessly running programs.
     #[inline]
     fn before_terminator(_ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx> {
         Ok(())
     }
 
+    /// Called when the interpreter encounters a `StatementKind::ConstEvalCounter` instruction.
+    /// You can use this to detect long or endlessly running programs.
+    #[inline]
+    fn increment_const_eval_counter(_ecx: &mut InterpCx<'mir, 'tcx, Self>) -> InterpResult<'tcx> {
+        Ok(())
+    }
+
     /// Called before a global allocation is accessed.
     /// `def_id` is `Some` if this is the "lazy" allocation of a static.
     #[inline]

compiler/rustc_const_eval/src/interpret/step.rs

@@ -129,6 +129,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             // FIXME(#73156): Handle source code coverage in const eval
             Coverage(..) => {}
 
+            ConstEvalCounter => {
+                M::increment_const_eval_counter(self)?;
+            }
+
             // Defined to do nothing. These are added by optimization passes, to avoid changing the
             // size of MIR constantly.
             Nop => {}

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

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

compiler/rustc_const_eval/src/transform/validate.rs

@@ -761,6 +761,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
             StatementKind::StorageLive(..)
             | StatementKind::StorageDead(..)
             | StatementKind::Coverage(_)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
         }
 

compiler/rustc_data_structures/src/graph/dominators/mod.rs

@@ -135,7 +135,47 @@ pub fn dominators<G: ControlFlowGraph>(graph: G) -> Dominators<G::Node> {
         // This loop computes the semi[w] for w.
         semi[w] = w;
         for v in graph.predecessors(pre_order_to_real[w]) {
-            // Reachable vertices may have unreachable predecessors, so ignore any of them
+            // TL;DR: Reachable vertices may have unreachable predecessors, so ignore any of them.
+            //
+            // Ignore blocks which are not connected to the entry block.
+            //
+            // The algorithm that was used to traverse the graph and build the
+            // `pre_order_to_real` and `real_to_pre_order` vectors does so by
+            // starting from the entry block and following the successors.
+            // Therefore, any blocks not reachable from the entry block will be
+            // set to `None` in the `pre_order_to_real` vector.
+            //
+            // For example, in this graph, A and B should be skipped:
+            //
+            //           ┌─────┐
+            //           │     │
+            //           └──┬──┘
+            //              │
+            //           ┌──▼──┐              ┌─────┐
+            //           │     │              │  A  │
+            //           └──┬──┘              └──┬──┘
+            //              │                    │
+            //      ┌───────┴───────┐            │
+            //      │               │            │
+            //   ┌──▼──┐         ┌──▼──┐      ┌──▼──┐
+            //   │     │         │     │      │  B  │
+            //   └──┬──┘         └──┬──┘      └──┬──┘
+            //      │               └──────┬─────┘
+            //   ┌──▼──┐                   │
+            //   │     │                   │
+            //   └──┬──┘                ┌──▼──┐
+            //      │                   │     │
+            //      │                   └─────┘
+            //   ┌──▼──┐
+            //   │     │
+            //   └──┬──┘
+            //      │
+            //   ┌──▼──┐
+            //   │     │
+            //   └─────┘
+            //
+            // ...this may be the case if a MirPass modifies the CFG to remove
+            // or rearrange certain blocks/edges.
             let Some(v) = real_to_pre_order[v] else {
                 continue
             };
@@ -264,13 +304,18 @@ fn compress(
     }
 }
 
+/// Tracks the list of dominators for each node.
 #[derive(Clone, Debug)]
 pub struct Dominators<N: Idx> {
     post_order_rank: IndexVec<N, usize>,
+    // Even though we track only the immediate dominator of each node, it's
+    // possible to get its full list of dominators by looking up the dominator
+    // of each dominator. (See the `impl Iterator for Iter` definition).
     immediate_dominators: IndexVec<N, Option<N>>,
 }
 
 impl<Node: Idx> Dominators<Node> {
+    /// Whether the given Node has an immediate dominator.
     pub fn is_reachable(&self, node: Node) -> bool {
         self.immediate_dominators[node].is_some()
     }
@@ -280,6 +325,8 @@ impl<Node: Idx> Dominators<Node> {
         self.immediate_dominators[node].unwrap()
     }
 
+    /// Provides an iterator over each dominator up the CFG, for the given Node.
+    /// See the `impl Iterator for Iter` definition to understand how this works.
     pub fn dominators(&self, node: Node) -> Iter<'_, Node> {
         assert!(self.is_reachable(node), "node {node:?} is not reachable");
         Iter { dominators: self, node: Some(node) }

compiler/rustc_interface/src/tests.rs

@@ -802,6 +802,7 @@ fn test_unstable_options_tracking_hash() {
     tracked!(teach, true);
     tracked!(thinlto, Some(true));
     tracked!(thir_unsafeck, true);
+    tracked!(tiny_const_eval_limit, true);
     tracked!(tls_model, Some(TlsModel::GeneralDynamic));
     tracked!(trait_solver, TraitSolver::Chalk);
     tracked!(translate_remapped_path_to_local_path, false);

compiler/rustc_middle/src/mir/mod.rs

@@ -1463,6 +1463,7 @@ impl Debug for Statement<'_> {
             }
             Coverage(box ref coverage) => write!(fmt, "Coverage::{:?}", coverage.kind),
             Intrinsic(box ref intrinsic) => write!(fmt, "{intrinsic}"),
+            ConstEvalCounter => write!(fmt, "ConstEvalCounter"),
             Nop => write!(fmt, "nop"),
         }
     }

compiler/rustc_middle/src/mir/spanview.rs

@@ -250,6 +250,7 @@ pub fn statement_kind_name(statement: &Statement<'_>) -> &'static str {
         AscribeUserType(..) => "AscribeUserType",
         Coverage(..) => "Coverage",
         Intrinsic(..) => "Intrinsic",
+        ConstEvalCounter => "ConstEvalCounter",
         Nop => "Nop",
     }
 }

compiler/rustc_middle/src/mir/syntax.rs

@@ -355,6 +355,12 @@ pub enum StatementKind<'tcx> {
     /// This avoids adding a new block and a terminator for simple intrinsics.
     Intrinsic(Box<NonDivergingIntrinsic<'tcx>>),
 
+    /// Instructs the const eval interpreter to increment a counter; this counter is used to track
+    /// how many steps the interpreter has taken. It is used to prevent the user from writing const
+    /// code that runs for too long or infinitely. Other than in the const eval interpreter, this
+    /// is a no-op.
+    ConstEvalCounter,
+
     /// No-op. Useful for deleting instructions without affecting statement indices.
     Nop,
 }

compiler/rustc_middle/src/mir/visit.rs

@@ -427,6 +427,7 @@ macro_rules! make_mir_visitor {
                             }
                         }
                     }
+                    StatementKind::ConstEvalCounter => {}
                     StatementKind::Nop => {}
                 }
             }

compiler/rustc_middle/src/ty/context.rs

@@ -77,6 +77,8 @@ use std::iter;
 use std::mem;
 use std::ops::{Bound, Deref};
 
+const TINY_CONST_EVAL_LIMIT: Limit = Limit(20);
+
 pub trait OnDiskCache<'tcx>: rustc_data_structures::sync::Sync {
     /// Creates a new `OnDiskCache` instance from the serialized data in `data`.
     fn new(sess: &'tcx Session, data: Mmap, start_pos: usize) -> Self
@@ -1104,7 +1106,11 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     pub fn const_eval_limit(self) -> Limit {
-        self.limits(()).const_eval_limit
+        if self.sess.opts.unstable_opts.tiny_const_eval_limit {
+            TINY_CONST_EVAL_LIMIT
+        } else {
+            self.limits(()).const_eval_limit
+        }
     }
 
     pub fn all_traits(self) -> impl Iterator<Item = DefId> + 'tcx {

compiler/rustc_mir_dataflow/src/impls/liveness.rs

@@ -271,6 +271,7 @@ impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'a> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::Intrinsic(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => None,
         };
         if let Some(destination) = destination {

compiler/rustc_mir_dataflow/src/impls/storage_liveness.rs

@@ -141,6 +141,7 @@ impl<'mir, 'tcx> crate::GenKillAnalysis<'tcx> for MaybeRequiresStorage<'mir, 'tc
             StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::FakeRead(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop
             | StatementKind::Retag(..)
             | StatementKind::Intrinsic(..)

compiler/rustc_mir_dataflow/src/move_paths/builder.rs

@@ -331,6 +331,7 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::Intrinsic(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
         }
     }

compiler/rustc_mir_dataflow/src/value_analysis.rs

@@ -84,7 +84,8 @@ pub trait ValueAnalysis<'tcx> {
             StatementKind::Retag(..) => {
                 // We don't track references.
             }
-            StatementKind::Nop
+            StatementKind::ConstEvalCounter
+            | StatementKind::Nop
             | StatementKind::FakeRead(..)
             | StatementKind::Coverage(..)
             | StatementKind::AscribeUserType(..) => (),

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 => {}
         }
     }