Skip to main content

charon_lib/transform/control_flow/
merge_goto_chains.rs

1//! # Micro-pass: merge single-origin gotos into their parent and skip over blocks that consist of
2//! only a goto.
3use rustc_hash::FxHashSet as HashSet;
4
5use crate::ids::IndexVec;
6use crate::transform::TransformCtx;
7use crate::ullbc_ast::*;
8
9use crate::transform::ctx::UllbcPass;
10
11/// Set of antecedents of a given block. We only care about block ids if there's a single
12/// antecedent.
13enum Antecedents {
14    Zero,
15    One {
16        id: BlockId,
17        /// Whether the antecedent is a goto.
18        is_goto: bool,
19    },
20    Many,
21}
22
23impl Antecedents {
24    fn add(&mut self, id: BlockId, is_goto: bool) {
25        match self {
26            Antecedents::Zero => *self = Antecedents::One { id, is_goto },
27            Antecedents::One { .. } => *self = Antecedents::Many,
28            Antecedents::Many => {}
29        }
30    }
31}
32
33pub struct Transform;
34impl UllbcPass for Transform {
35    fn transform_body(&self, _ctx: &mut TransformCtx, body: &mut ExprBody) {
36        // Compute for each block the set of blocks that points to it.
37        let mut antecedents: IndexVec<BlockId, Antecedents> =
38            body.body.map_ref(|_| Antecedents::Zero);
39        for (block_id, block) in body.body.iter_enumerated() {
40            let is_goto = block.terminator.kind.is_goto();
41            for target in block.targets() {
42                antecedents.get_mut(target).unwrap().add(block_id, is_goto);
43            }
44        }
45        // Merge blocks with a single antecedent into their antecedent.
46        for mut id in body.body.indices() {
47            // Go up the chain to find the first parent into which we can merge.
48            while let Antecedents::One {
49                id: antecedent_id,
50                is_goto: true,
51            } = antecedents[id]
52            {
53                id = antecedent_id;
54            }
55            // While the current block is a straight goto, merge the target block back into this
56            // one.
57            while let source = &body.body[id]
58                && let TerminatorKind::Goto { target } = source.terminator.kind
59                && let Antecedents::One { .. } = antecedents[target]
60                && source.kind == body.body[target].kind
61            {
62                antecedents[target] = Antecedents::Zero;
63                let mut target = body.body[target].take();
64                let source = &mut body.body[id];
65                source.statements.append(&mut target.statements);
66                source.terminator = target.terminator;
67            }
68        }
69        // Skip over trivial gotos.
70        let mut visited = HashSet::default(); // detect and skip loops
71        for id in body.body.indices() {
72            if body.body[id]
73                .targets()
74                .into_iter()
75                .any(|t| body.body[t].as_trivial_goto().is_some())
76            {
77                // Merge any forward goto chains that start here.
78                let mut source = body.body[id].take();
79                for target_id in source.terminator.targets_mut() {
80                    visited.clear();
81                    visited.insert(id);
82                    while let Some(b) = body.body[*target_id].as_trivial_goto()
83                        && source.kind == body.body[*target_id].kind
84                        && visited.insert(*target_id)
85                    {
86                        *target_id = b
87                    }
88                }
89                body.body[id] = source;
90            }
91        }
92    }
93}