charon_lib/transform/
duplicate_return.rs

1//! The MIR uses a unique `return` node, which can be an issue when reconstructing
2//! the control-flow.
3//!
4//! For instance, it often leads to code of the following shape:
5//! ```text
6//! if b {
7//!   ...
8//!   x = 0;
9//! }
10//! else {
11//!   ...
12//!   x = 1;
13//! }
14//! return x;
15//! ```
16//!
17//! while a more natural reconstruction would be:
18//! ```text
19//! if b {
20//!   ...
21//!   return 0;
22//! }
23//! else {
24//!   ...
25//!   return 1;
26//! }
27//! ```
28
29use crate::ids::Generator;
30use crate::transform::TransformCtx;
31use crate::ullbc_ast::*;
32use std::collections::HashMap;
33
34use super::ctx::UllbcPass;
35
36pub struct Transform;
37impl UllbcPass for Transform {
38    fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
39        // Find the return block id (there should be one).
40        let returns: HashMap<BlockId, Span> = b
41            .body
42            .iter_indexed()
43            .filter_map(|(bid, block)| {
44                if block.statements.is_empty() && block.terminator.content.is_return() {
45                    Some((bid, block.terminator.span))
46                } else {
47                    None
48                }
49            })
50            .collect();
51
52        // Whenever we find a goto the return block, introduce an auxiliary block
53        // for this (remark: in the end, it makes the return block dangling).
54        // We do this in two steps.
55        // First, introduce fresh ids.
56        let mut generator = Generator::new_with_init_value(b.body.next_id());
57        let mut new_spans = Vec::new();
58        b.body.dyn_visit_in_body_mut(|bid: &mut BlockId| {
59            if let Some(span) = returns.get(bid) {
60                *bid = generator.fresh_id();
61                new_spans.push(*span);
62            }
63        });
64
65        // Then introduce the new blocks
66        for span in new_spans {
67            let _ = b.body.push(BlockData {
68                statements: Vec::new(),
69                terminator: Terminator::new(span, RawTerminator::Return),
70            });
71        }
72    }
73}