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 at_5_ty;
57 let box_generics;
58 let value_to_write;
59 let old_assign_idx;
60 let assign_span;
61 let unwind_target;
62
63 if let Some(candidate_block) = b.body.get(candidate_block_idx)
64 // If the terminator is a call
65 && let RawTerminator::Call {
66 target: target_block_idx,
67 call:
68 Call {
69 args: malloc_args,
70 func: _, // TODO: once we have a system to recognize intrinsics, check the call is to exchange_malloc.
71 dest: malloc_dest,
72 },
73 on_unwind,
74 } = &candidate_block.terminator.content
75 // The call has two move arguments
76 && let [Operand::Move(arg0), Operand::Move(arg1)] = malloc_args.as_slice()
77 && let [ .., Statement {
78 content: RawStatement::Assign(size, Rvalue::NullaryOp(NullOp::SizeOf, _)),
79 ..
80 }, Statement {
81 content: RawStatement::Assign(align, Rvalue::NullaryOp(NullOp::AlignOf, _)),
82 ..
83 }] = candidate_block.statements.as_slice()
84 && arg0 == size && arg1 == align
85 && let Some(target_block) = b.body.get(*target_block_idx)
86 && let [Statement {
87 content: RawStatement::StorageLive(target_var),
88 ..
89 }, Statement {
90 content:
91 RawStatement::Assign(box_make, Rvalue::ShallowInitBox(Operand::Move(alloc_use), _)),
92 ..
93 }, rest @ ..] = target_block.statements.as_slice()
94 && alloc_use == malloc_dest
95 && let Some(local_id) = box_make.as_local()
96 && local_id == *target_var
97 && let TyKind::Adt(ty_ref) = b.locals[*target_var].ty.kind()
98 && let TypeId::Builtin(BuiltinTy::Box) = ty_ref.id
99 && let Some((assign_idx_in_rest, val, span)) = rest.iter().enumerate().find_map(|(idx, st)| {
100 if let Statement {
101 content: RawStatement::Assign(box_deref, val),
102 span,
103 ..
104 } = st
105 && let Some((sub, ProjectionElem::Deref)) = box_deref.as_projection()
106 && sub == box_make
107 {
108 Some((idx, val, span))
109 } else {
110 None
111 }
112 })
113 {
114 at_5 = local_id;
115 at_5_ty = box_make.ty().clone();
116 old_assign_idx = assign_idx_in_rest + 2; // +2 because rest skips the first two statements
117 value_to_write = val.clone();
118 box_generics = ty_ref.generics.clone();
119 second_block = *target_block_idx;
120 assign_span = *span;
121 unwind_target = *on_unwind;
122 } else {
123 continue;
124 }
125
126 let first_block = b.body.get_mut(candidate_block_idx).unwrap();
127 let number_statements = first_block.statements.len();
128 let value_to_write = match value_to_write {
129 Rvalue::Use(op) => {
130 first_block
131 .statements
132 .get_mut(number_statements - 2)
133 .unwrap()
134 .content = RawStatement::Nop;
135 op
136 }
137 _ => {
138 // We need to create a new variable to store the value.
139 let name = b.locals[at_5].name.clone();
140 let ty = box_generics.types[0].clone();
141 let var = b.locals.new_var(name, ty);
142 let st = Statement::new(
143 assign_span,
144 RawStatement::Assign(var.clone(), value_to_write),
145 );
146 // We overide the @2 := size_of<i32> statement with the rvalue assignment
147 *first_block
148 .statements
149 .get_mut(number_statements - 2)
150 .unwrap() = st;
151 Operand::Move(var)
152 }
153 };
154 first_block
155 .statements
156 .get_mut(number_statements - 1)
157 .unwrap()
158 .content = RawStatement::StorageLive(at_5);
159 first_block.terminator.content = RawTerminator::Call {
160 call: Call {
161 func: FnOperand::Regular(FnPtr {
162 func: Box::new(FunIdOrTraitMethodRef::Fun(FunId::Builtin(
163 BuiltinFunId::BoxNew,
164 ))),
165 generics: box_generics,
166 }),
167 args: vec![value_to_write],
168 dest: Place::new(at_5, at_5_ty),
169 },
170 target: second_block,
171 on_unwind: unwind_target,
172 };
173
174 // We now update the statements in the second block.
175 let second_block = b.body.get_mut(second_block).unwrap();
176 second_block.statements.get_mut(0).unwrap().content = RawStatement::Nop;
177 second_block.statements.get_mut(1).unwrap().content = RawStatement::Nop;
178 second_block
179 .statements
180 .get_mut(old_assign_idx)
181 .unwrap()
182 .content = RawStatement::Nop;
183 }
184
185 // Make sure we got all the `ShallowInitBox`es.
186 b.body.dyn_visit_in_body(|rvalue: &Rvalue| {
187 if rvalue.is_shallow_init_box() {
188 register_error!(
189 ctx,
190 b.span,
191 "Could not reconstruct `Box` initialization; \
192 branching during `Box` initialization is not supported."
193 );
194 }
195 });
196 }
197}