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