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