charon_lib/transform/control_flow/
merge_goto_chains.rs1use 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
11enum Antecedents {
14 Zero,
15 One {
16 id: BlockId,
17 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 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 for mut id in body.body.indices() {
47 while let Antecedents::One {
49 id: antecedent_id,
50 is_goto: true,
51 } = antecedents[id]
52 {
53 id = antecedent_id;
54 }
55 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 let mut visited = HashSet::default(); 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 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}