Create stable metric to measure long computation in Const Eval

This patch adds a `MirPass` that tracks the number of back-edges and
function calls in the CFG, adds a new MIR instruction to increment a
counter every time they are encountered during Const Eval, and emit a
warning if a configured limit is breached.

10 files changed, 111 insertions, 3 deletions

diff --git a/compiler/rustc_mir_transform/src/check_unsafety.rs b/compiler/rustc_mir_transform/src/check_unsafety.rs
index adf6ae4c727..837233953e7 100644
--- a/compiler/rustc_mir_transform/src/check_unsafety.rs
+++ b/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)
             }
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index 31d5541a31b..f973c1ed28f 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/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,
 
diff --git a/compiler/rustc_mir_transform/src/ctfe_limit.rs b/compiler/rustc_mir_transform/src/ctfe_limit.rs
new file mode 100644
index 00000000000..dc54b983c0e
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/ctfe_limit.rs
@@ -0,0 +1,92 @@
+use crate::MirPass;
+
+use rustc_middle::mir::{BasicBlock, Body, Statement, StatementKind, TerminatorKind};
+use rustc_middle::ty::TyCtxt;
+
+use tracing::{info, instrument};
+
+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();
+        //info!("Got body with {} basic blocks: {:#?}", body.basic_blocks.len(), body.basic_blocks);
+        //info!("With doms: {doms:?}");
+
+        /*
+        for (index, basic_block) in body.basic_blocks.iter().enumerate() {
+            info!("bb{index}: {basic_block:#?}")
+        }*/
+        for (index, basic_block) in body.basic_blocks.iter().enumerate() {
+            info!(
+                "bb{index} -> successors = {:?}",
+                basic_block.terminator().successors().collect::<Vec<BasicBlock>>()
+            );
+        }
+        for (index, basic_block) in body.basic_blocks.iter().enumerate() {
+            info!("bb{index} -> unwind = {:?}", basic_block.terminator().unwind())
+        }
+
+        let mut dominators = Vec::new();
+        for idom in 0..body.basic_blocks.len() {
+            let mut nodes = Vec::new();
+            for inode in 0..body.basic_blocks.len() {
+                let dom = BasicBlock::from_usize(idom);
+                let node = BasicBlock::from_usize(inode);
+                if doms.is_reachable(dom)
+                    && doms.is_reachable(node)
+                    && doms.is_dominated_by(node, dom)
+                {
+                    //info!("{idom} dominates {inode}");
+                    nodes.push(true);
+                } else {
+                    nodes.push(false);
+                }
+            }
+            dominators.push(nodes);
+        }
+        /*
+        for idom in 0..body.basic_blocks.len() {
+            print!("{idom} | dom | ");
+            for inode in 0..body.basic_blocks.len() {
+                if dominators[idom][inode] {
+                    print!("{inode} | ");
+                } else {
+                    print!("  | ");
+                }
+            }
+            print!("\n");
+        }
+        */
+
+        for (index, basic_block) in body.basic_blocks_mut().iter_mut().enumerate() {
+            // info!("bb{index}: {basic_block:#?}");
+            //info!("bb{index} -> successors = {:?}", basic_block.terminator().successors().collect::<Vec<BasicBlock>>());
+            let is_back_edge_or_fn_call = 'label: {
+                match basic_block.terminator().kind {
+                    TerminatorKind::Call { .. } => {
+                        break 'label true;
+                    }
+                    _ => (),
+                }
+                for successor in basic_block.terminator().successors() {
+                    let s_index = successor.as_usize();
+                    if dominators[s_index][index] {
+                        info!("{s_index} to {index} is a loop");
+                        break 'label true;
+                    }
+                }
+                false
+            };
+            if is_back_edge_or_fn_call {
+                basic_block.statements.push(Statement {
+                    source_info: basic_block.terminator().source_info,
+                    kind: StatementKind::ConstEvalCounter,
+                });
+                info!("New basic block statements vector: {:?}", basic_block.statements);
+            }
+        }
+        info!("With doms: {doms:?}");
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/dead_store_elimination.rs b/compiler/rustc_mir_transform/src/dead_store_elimination.rs
index 09546330cec..9dbfb089dc6 100644
--- a/compiler/rustc_mir_transform/src/dead_store_elimination.rs
+++ b/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(_, _) => {
diff --git a/compiler/rustc_mir_transform/src/dest_prop.rs b/compiler/rustc_mir_transform/src/dest_prop.rs
index 08e296a8371..20ffb0ab334 100644
--- a/compiler/rustc_mir_transform/src/dest_prop.rs
+++ b/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(_)
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
index 39c61a34afc..0df732aa22b 100644
--- a/compiler/rustc_mir_transform/src/generator.rs
+++ b/compiler/rustc_mir_transform/src/generator.rs
@@ -1583,6 +1583,7 @@ impl<'tcx> Visitor<'tcx> for EnsureGeneratorFieldAssignmentsNeverAlias<'_> {
             | StatementKind::AscribeUserType(..)
             | StatementKind::Coverage(..)
             | StatementKind::Intrinsic(..)
+            | StatementKind::ConstEvalCounter
             | StatementKind::Nop => {}
         }
     }
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 20b7fdcfe6d..e5c8127bea1 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/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;
@@ -349,11 +350,14 @@ fn mir_promoted(
 /// Compute the MIR that is used during CTFE (and thus has no optimizations run on it)
 fn mir_for_ctfe(tcx: TyCtxt<'_>, def_id: DefId) -> &Body<'_> {
     let did = def_id.expect_local();
-    if let Some(def) = ty::WithOptConstParam::try_lookup(did, tcx) {
+    let body = if let Some(def) = ty::WithOptConstParam::try_lookup(did, tcx) {
         tcx.mir_for_ctfe_of_const_arg(def)
     } else {
         tcx.arena.alloc(inner_mir_for_ctfe(tcx, ty::WithOptConstParam::unknown(did)))
-    }
+    };
+    //info!("MIR_FOR_CTFE (DefId = {def_id:?}) body res: {:#?}", body);
+    info!("MIR_FOR_CTFE (DefId = {def_id:?})");
+    body
 }
 
 /// Same as `mir_for_ctfe`, but used to get the MIR of a const generic parameter.
@@ -447,6 +451,7 @@ fn mir_drops_elaborated_and_const_checked(
 
     run_analysis_to_runtime_passes(tcx, &mut body);
 
+    //info!("MIR after runtime passes: {:#?}", body);
     tcx.alloc_steal_mir(body)
 }
 
@@ -517,6 +522,7 @@ fn run_runtime_lowering_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         // CTFE support for aggregates.
         &deaggregator::Deaggregator,
         &Lint(const_prop_lint::ConstProp),
+        &ctfe_limit::CtfeLimit,
     ];
     pm::run_passes_no_validate(tcx, body, passes, Some(MirPhase::Runtime(RuntimePhase::Initial)));
 }
@@ -617,6 +623,7 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
     let mut body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::NotConst);
     debug!("body: {:#?}", body);
     run_optimization_passes(tcx, &mut body);
+    //info!("body after OPTIMIZATION: {:#?}", body);
 
     debug_assert!(!body.has_free_regions(), "Free regions in optimized MIR");
 
diff --git a/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs b/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs
index f1bbf2ea7e8..e3a03aa08af 100644
--- a/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs
+++ b/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
                 }
diff --git a/compiler/rustc_mir_transform/src/separate_const_switch.rs b/compiler/rustc_mir_transform/src/separate_const_switch.rs
index 2f116aaa958..a24d2d34d79 100644
--- a/compiler/rustc_mir_transform/src/separate_const_switch.rs
+++ b/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
diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs
index 8f6abe7a912..7b6fa2baf2f 100644
--- a/compiler/rustc_mir_transform/src/simplify.rs
+++ b/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) => {}