rustc_mir_transform/
ctfe_limit.rs

1//! A pass that inserts the `ConstEvalCounter` instruction into any blocks that have a back edge
2//! (thus indicating there is a loop in the CFG), or whose terminator is a function call.
3
4use rustc_data_structures::graph::dominators::Dominators;
5use rustc_middle::mir::{
6    BasicBlock, BasicBlockData, Body, Statement, StatementKind, TerminatorKind,
7};
8use rustc_middle::ty::TyCtxt;
9use tracing::instrument;
10
11pub(super) struct CtfeLimit;
12
13impl<'tcx> crate::MirPass<'tcx> for CtfeLimit {
14    #[instrument(skip(self, _tcx, body))]
15    fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
16        let doms = body.basic_blocks.dominators();
17        let indices: Vec<BasicBlock> = body
18            .basic_blocks
19            .iter_enumerated()
20            .filter_map(|(node, node_data)| {
21                if matches!(node_data.terminator().kind, TerminatorKind::Call { .. })
22                    // Back edges in a CFG indicate loops
23                    || has_back_edge(doms, node, node_data)
24                {
25                    Some(node)
26                } else {
27                    None
28                }
29            })
30            .collect();
31        for index in indices {
32            insert_counter(
33                body.basic_blocks_mut()
34                    .get_mut(index)
35                    .expect("basic_blocks index {index} should exist"),
36            );
37        }
38    }
39
40    fn is_required(&self) -> bool {
41        true
42    }
43}
44
45fn has_back_edge(
46    doms: &Dominators<BasicBlock>,
47    node: BasicBlock,
48    node_data: &BasicBlockData<'_>,
49) -> bool {
50    if !doms.is_reachable(node) {
51        return false;
52    }
53    // Check if any of the dominators of the node are also the node's successor.
54    node_data.terminator().successors().any(|succ| doms.dominates(succ, node))
55}
56
57fn insert_counter(basic_block_data: &mut BasicBlockData<'_>) {
58    basic_block_data.statements.push(Statement {
59        source_info: basic_block_data.terminator().source_info,
60        kind: StatementKind::ConstEvalCounter,
61    });
62}