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}