charon_lib/transform/
reconstruct_boxes.rs

1//! # Micro-pass: reconstruct piecewise box allocations using `malloc` and `ShallowInitBox`.
2
3use crate::register_error;
4use crate::transform::TransformCtx;
5use crate::ullbc_ast::*;
6
7use super::ctx::UllbcPass;
8
9pub struct Transform;
10
11/// The special `alloc::boxed::box_new(x)` intrinsic becomes the following:
12///
13/// ```text
14/// @2 := size_of<i32>
15/// @3 := align_of<i32>
16/// @4 := alloc::alloc::exchange_malloc(move (@2), move (@3))
17/// storage_live(@5)
18/// @5 := shallow_init_box::<i32>(move (@4))
19/// // possibly some intermediate statements
20/// *(@5) := x
21/// ```
22///
23/// We reconstruct this into a call to `Box::new(x)`.
24impl UllbcPass for Transform {
25    fn transform_body(&self, ctx: &mut TransformCtx, b: &mut ExprBody) {
26        if ctx.options.raw_boxes {
27            return;
28        }
29
30        // We need to find a block that has exchange_malloc as the following terminator:
31        // ```text
32        // @4 := alloc::alloc::exchange_malloc(move (@2), move (@3))
33        // ```
34        // We then chekc that that this block ends with two assignments:
35        // ```text
36        // @2 := size_of<i32>
37        // @3 := align_of<i32>
38        // ```
39        // If that is the case, we look at the target block and check that it starts with`
40        // ```text
41        // storage_live(@5)
42        // @5 := shallow_init_box::<i32>(move (@4))
43        // ```
44        // We then look for the assignment into the box and take a not of its index.
45        // ```text
46        // *(@5) := x
47        // ```
48        // Finally, we replace all these assignments with a call to `@5 = Box::new(x)`
49        // We do so by replacing the terminator (exchange_malloc) with the correct call
50        // and replacing the assignment @3 := align_of<i32> with the storage live.
51        // Everything else becomes Nop.
52
53        for candidate_block_idx in b.body.all_indices() {
54            let second_block;
55            let at_5;
56            let box_generics;
57            let value_to_write;
58            let old_assign_idx;
59            let assign_span;
60            let unwind_target;
61
62            if let Some(candidate_block) = b.body.get(candidate_block_idx)
63                // If the terminator is a call
64                && let RawTerminator::Call {
65                    target: target_block_idx,
66                    call:
67                        Call {
68                            args: malloc_args,
69                            func: _, // TODO: once we have a system to recognize intrinsics, check the call is to exchange_malloc.
70                            dest: malloc_dest,
71                        },
72                        on_unwind,
73                } = &candidate_block.terminator.content
74                // The call has two move arguments
75                && let [Operand::Move(arg0), Operand::Move(arg1)] = malloc_args.as_slice()
76                && let [ .., Statement {
77                            content: RawStatement::Assign(size, Rvalue::NullaryOp(NullOp::SizeOf, _)),
78                            ..
79                        }, Statement {
80                            content: RawStatement::Assign(align, Rvalue::NullaryOp(NullOp::AlignOf, _)),
81                            ..
82                        }] = candidate_block.statements.as_slice()
83                && arg0 == size && arg1 == align
84                && let Some(target_block) = b.body.get(*target_block_idx)
85                && let [Statement {
86                            content: RawStatement::StorageLive(target_var),
87                            ..
88                        }, Statement {
89                            content:
90                                RawStatement::Assign(box_make, Rvalue::ShallowInitBox(Operand::Move(alloc_use), _)),
91                            ..
92                        }, rest @ ..] = target_block.statements.as_slice()
93                && alloc_use == malloc_dest
94                && box_make.is_local()
95                && box_make.local_id() == *target_var
96                && let TyKind::Adt(ty_ref) = b.locals[*target_var].ty.kind()
97                && let TypeId::Builtin(BuiltinTy::Box) = ty_ref.id
98                && let Some((assign_idx_in_rest, val, span)) = rest.iter().enumerate().find_map(|(idx, st)| {
99                    if let Statement {
100                            content: RawStatement::Assign(box_deref, val),
101                            span,
102                            ..
103                        } = st
104                        && let Some((sub, ProjectionElem::Deref)) = box_deref.as_projection()
105                        && sub == box_make
106                    {
107                        Some((idx, val, span))
108                    } else {
109                        None
110                    }
111                })
112            {
113                at_5 = box_make.clone();
114                old_assign_idx = assign_idx_in_rest + 2; // +2 because rest skips the first two statements
115                value_to_write = val.clone();
116                box_generics = ty_ref.generics.clone();
117                second_block = *target_block_idx;
118                assign_span = *span;
119                unwind_target = *on_unwind;
120            } else {
121                continue;
122            }
123
124            let first_block = b.body.get_mut(candidate_block_idx).unwrap();
125            let number_statements = first_block.statements.len();
126            let value_to_write = match value_to_write {
127                Rvalue::Use(op) => {
128                    first_block
129                        .statements
130                        .get_mut(number_statements - 2)
131                        .unwrap()
132                        .content = RawStatement::Nop;
133                    op
134                }
135                _ => {
136                    // We need to create a new variable to store the value.
137                    let name = b.locals[at_5.local_id()].name.clone();
138                    let ty = box_generics.types[0].clone();
139                    let var = b.locals.new_var(name, ty);
140                    let st = Statement::new(
141                        assign_span,
142                        RawStatement::Assign(var.clone(), value_to_write),
143                    );
144                    // We overide the @2 := size_of<i32> statement with the rvalue assignment
145                    *first_block
146                        .statements
147                        .get_mut(number_statements - 2)
148                        .unwrap() = st;
149                    Operand::Move(var)
150                }
151            };
152            first_block
153                .statements
154                .get_mut(number_statements - 1)
155                .unwrap()
156                .content = RawStatement::StorageLive(at_5.local_id());
157            first_block.terminator.content = RawTerminator::Call {
158                call: Call {
159                    func: FnOperand::Regular(FnPtr {
160                        func: Box::new(FunIdOrTraitMethodRef::Fun(FunId::Builtin(
161                            BuiltinFunId::BoxNew,
162                        ))),
163                        generics: box_generics,
164                    }),
165                    args: vec![value_to_write],
166                    dest: at_5,
167                },
168                target: second_block,
169                on_unwind: unwind_target,
170            };
171
172            // We now update the statements in the second block.
173            let second_block = b.body.get_mut(second_block).unwrap();
174            second_block.statements.get_mut(0).unwrap().content = RawStatement::Nop;
175            second_block.statements.get_mut(1).unwrap().content = RawStatement::Nop;
176            second_block
177                .statements
178                .get_mut(old_assign_idx)
179                .unwrap()
180                .content = RawStatement::Nop;
181        }
182
183        // Make sure we got all the `ShallowInitBox`es.
184        b.body.dyn_visit_in_body(|rvalue: &Rvalue| {
185            if rvalue.is_shallow_init_box() {
186                register_error!(
187                    ctx,
188                    b.span,
189                    "Could not reconstruct `Box` initialization; \
190                    branching during `Box` initialization is not supported."
191                );
192            }
193        });
194    }
195}