charon_lib/transform/control_flow/
merge_goto_chains.rs

1//! # Micro-pass: merge single-origin gotos into their parent to reduce CFG graph size.
2use std::mem;
3
4use crate::ids::IndexVec;
5use crate::transform::TransformCtx;
6use crate::ullbc_ast::*;
7
8use crate::transform::ctx::UllbcPass;
9
10/// Set of antecedents of a given block. We only care about block ids if there's a single
11/// antecedent.
12enum Antecedents {
13    Zero,
14    One {
15        id: BlockId,
16        /// Whether the antecedent is a goto.
17        is_goto: bool,
18    },
19    Many,
20}
21
22impl Antecedents {
23    fn add(&mut self, id: BlockId, is_goto: bool) {
24        match self {
25            Antecedents::Zero => *self = Antecedents::One { id, is_goto },
26            Antecedents::One { .. } => *self = Antecedents::Many,
27            Antecedents::Many => {}
28        }
29    }
30}
31
32pub struct Transform;
33impl UllbcPass for Transform {
34    fn transform_body(&self, _ctx: &mut TransformCtx, body: &mut ExprBody) {
35        // Compute for each block the set of blocks that points to it.
36        let mut antecedents: IndexVec<BlockId, Antecedents> =
37            body.body.map_ref(|_| Antecedents::Zero);
38        for (block_id, block) in body.body.iter_indexed() {
39            let is_goto = block.terminator.kind.is_goto();
40            for target in block.targets() {
41                antecedents.get_mut(target).unwrap().add(block_id, is_goto);
42            }
43        }
44        // Merge blocks with a single antecedent into their antecedent.
45        for mut id in body.body.all_indices() {
46            // Go up the chain to find the first parent with zero (the start block) or multiple
47            // antecedents. This avoids quadratic behavior where we repeatedly copy a growing list
48            // of statements, since blocks may not be sorted..
49            while let Antecedents::One {
50                id: antecedent_id,
51                is_goto: true,
52            } = antecedents[id]
53            {
54                id = antecedent_id;
55            }
56            // While the current block is a straight goto, merge the target block back into this
57            // one.
58            while let Some(source) = body.body.get(id)
59                && let TerminatorKind::Goto { target } = source.terminator.kind
60                && let Antecedents::One { .. } = antecedents[target]
61            {
62                let mut target = mem::replace(&mut body.body[target], BlockData::new_unreachable());
63                let source = &mut body.body[id];
64                source.statements.append(&mut target.statements);
65                source.terminator = target.terminator;
66            }
67        }
68    }
69}