Skip to main content

charon_lib/transform/resugar/
resugar_drops.rs

1//! Reconstruct conditional drops using the drop flags identified by `detect_drop_flags`.
2use crate::options::TranslateOptions;
3use crate::transform::TransformCtx;
4use crate::transform::ctx::UllbcPass;
5use crate::ullbc_ast::*;
6
7struct ConditionalDrop {
8    /// The block that switched on the drop flag.
9    switch_block: BlockId,
10    /// The block that contains the unconditional drop.
11    drop_block: BlockId,
12}
13
14#[derive(Default)]
15struct FlagUses {
16    assignments: Vec<StmtLoc>,
17    if_count: usize,
18    resugared_if_count: usize,
19}
20
21pub struct Transform;
22impl UllbcPass for Transform {
23    fn should_run(&self, options: &TranslateOptions) -> bool {
24        options.resugar_drops
25    }
26
27    fn transform_body(&self, _ctx: &mut TransformCtx, body: &mut ExprBody) {
28        let mut predecessor_counts = body.body.map_ref(|_| 0usize);
29        for block in &body.body {
30            for target in block.terminator.targets() {
31                predecessor_counts[target] += 1;
32            }
33        }
34
35        let mut flag_uses = body.locals.locals.map_ref(|_| FlagUses::default());
36        let mut conditional_drops = Vec::new();
37
38        for (block_id, block) in body.body.iter_enumerated() {
39            for (statement_id, statement) in block.statements.iter().enumerate() {
40                if let StatementKind::Assign(place, _) = &statement.kind
41                    && let Some(flag) = place.as_local()
42                    && body.locals.locals[flag].drop_flag_for.is_some()
43                {
44                    flag_uses[flag]
45                        .assignments
46                        .push(StmtLoc::new(block_id, statement_id));
47                }
48            }
49
50            if let TerminatorKind::Switch { data, branches } = &block.terminator.kind
51                && let SwitchScrutinee::Value(Operand::Copy(flag_place) | Operand::Move(flag_place)) =
52                    &data.scrutinee
53                && let Some(flag) = flag_place.as_local()
54                && let Some(flagged_place) = &body.locals.locals[flag].drop_flag_for
55                && let Some((true_branch, false_branch)) = data.as_if()
56            {
57                flag_uses[flag].if_count += 1;
58                let true_target = branches[true_branch];
59                let false_target = branches[false_branch];
60                let drop_block = &body.body[true_target];
61
62                // Detect a simple pattern: `if drop_flag { unconditional_drop } else {}`.
63                if true_target != false_target
64                    && predecessor_counts[true_target] == 1
65                    && let TerminatorKind::Drop {
66                        kind: DropKind::Precise,
67                        place,
68                        target,
69                        ..
70                    } = &drop_block.terminator.kind
71                    && place == flagged_place
72                    && drop_block.statements.is_empty()
73                    && *target == false_target
74                {
75                    flag_uses[flag].resugared_if_count += 1;
76                    conditional_drops.push(ConditionalDrop {
77                        switch_block: block_id,
78                        drop_block: true_target,
79                    });
80                }
81            }
82        }
83
84        for drop in conditional_drops {
85            let mut terminator = body.body[drop.drop_block].terminator.kind.take();
86            let TerminatorKind::Drop { kind, .. } = &mut terminator else {
87                unreachable!()
88            };
89            *kind = DropKind::Conditional;
90            body.body[drop.switch_block].terminator.kind = terminator;
91        }
92
93        // If every use of a flag has been resugared, remove its assignments. The unused-locals
94        // pass will then remove the local entirely.
95        for uses in flag_uses {
96            if uses.if_count == uses.resugared_if_count {
97                for loc in uses.assignments {
98                    body[loc].kind = StatementKind::Nop;
99                }
100            }
101        }
102    }
103}