charon_lib/transform/
index_to_function_calls.rs

1//! Desugar array/slice index operations to function calls.
2use crate::llbc_ast::*;
3use crate::transform::TransformCtx;
4use derive_generic_visitor::*;
5
6use super::ctx::LlbcPass;
7
8/// We replace some place constructors with function calls. To do that, we explore all the places
9/// in a body and deconstruct a given place access into intermediate assignments.
10///
11/// We accumulate the new assignments as statements in the visitor, and at the end we insert these
12/// statements before the one that was just explored.
13#[derive(Visitor)]
14struct IndexVisitor<'a> {
15    locals: &'a mut Locals,
16    /// Statements to prepend to the statement currently being explored.
17    statements: Vec<Statement>,
18    // When we visit a place, we need to know if it is being accessed mutably or not. Whenever we
19    // visit something that contains a place we push the relevant mutability on this stack.
20    // Unfortunately this requires us to be very careful to catch all the cases where we see
21    // places.
22    place_mutability_stack: Vec<bool>,
23    // Span of the statement.
24    span: Span,
25}
26
27impl<'a> IndexVisitor<'a> {
28    fn fresh_var(&mut self, name: Option<String>, ty: Ty) -> Place {
29        let var = self.locals.new_var(name, ty);
30        let live_kind = RawStatement::StorageLive(var.local_id());
31        self.statements.push(Statement::new(self.span, live_kind));
32        var
33    }
34
35    fn transform_place(&mut self, mut_access: bool, place: &mut Place) {
36        use ProjectionElem::*;
37        let Some((subplace, pe @ (Index { .. } | Subslice { .. }))) = place.as_projection() else {
38            return;
39        };
40        let tref = subplace.ty.as_adt().unwrap();
41        let builtin_ty = tref.id.as_builtin().unwrap();
42
43        // The built-in function to call.
44        let indexing_function = {
45            let builtin_fun = BuiltinFunId::Index(BuiltinIndexOp {
46                is_array: matches!(builtin_ty, BuiltinTy::Array),
47                mutability: RefKind::mutable(mut_access),
48                is_range: matches!(pe, Subslice { .. }),
49            });
50            // Same generics as the array/slice type, except for the extra lifetime.
51            let mut generics = tref.generics.clone();
52            generics.regions = [Region::Erased].into();
53            FnOperand::Regular(FnPtr {
54                func: Box::new(FunIdOrTraitMethodRef::mk_builtin(builtin_fun)),
55                generics,
56            })
57        };
58
59        let input_ty = TyKind::Ref(
60            Region::Erased,
61            subplace.ty().clone(),
62            RefKind::mutable(mut_access),
63        )
64        .into_ty();
65
66        let elem_ty = tref.generics.types[0].clone();
67        let output_inner_ty = if matches!(pe, Index { .. }) {
68            elem_ty
69        } else {
70            TyKind::Adt(TypeDeclRef {
71                id: TypeId::Builtin(BuiltinTy::Slice),
72                generics: Box::new(GenericArgs::new_for_builtin(vec![elem_ty].into())),
73            })
74            .into_ty()
75        };
76        let output_ty = {
77            TyKind::Ref(
78                Region::Erased,
79                output_inner_ty.clone(),
80                RefKind::mutable(mut_access),
81            )
82            .into_ty()
83        };
84
85        // Push the statements:
86        // `storage_live(tmp0)`
87        // `tmp0 = &{mut}p`
88        let input_var = {
89            let input_var = self.fresh_var(None, input_ty);
90            let kind = RawStatement::Assign(
91                input_var.clone(),
92                Rvalue::Ref(subplace.clone(), BorrowKind::mutable(mut_access)),
93            );
94            self.statements.push(Statement::new(self.span, kind));
95            input_var
96        };
97
98        // Construct the arguments to pass to the indexing function.
99        let mut args = vec![Operand::Move(input_var)];
100        if let Subslice { from, .. } = &pe {
101            args.push(from.as_ref().clone());
102        }
103        let (last_arg, from_end) = match &pe {
104            Index {
105                offset: x,
106                from_end,
107                ..
108            }
109            | Subslice {
110                to: x, from_end, ..
111            } => (x.as_ref().clone(), *from_end),
112            _ => unreachable!(),
113        };
114        if from_end {
115            // `storage_live(len_var)`
116            // `len_var = len(p)`
117            let usize_ty = TyKind::Literal(LiteralTy::Integer(IntegerTy::Usize)).into_ty();
118            let len_var = self.fresh_var(None, usize_ty.clone());
119            let kind = RawStatement::Assign(
120                len_var.clone(),
121                Rvalue::Len(
122                    subplace.clone(),
123                    subplace.ty().clone(),
124                    tref.generics.const_generics.get(0.into()).cloned(),
125                ),
126            );
127            self.statements.push(Statement::new(self.span, kind));
128
129            // `storage_live(index_var)`
130            // `index_var = len_var - last_arg`
131            // `storage_dead(len_var)`
132            let index_var = self.fresh_var(None, usize_ty);
133            let kind = RawStatement::Assign(
134                index_var.clone(),
135                Rvalue::BinaryOp(
136                    BinOp::Sub(OverflowMode::UB),
137                    Operand::Copy(len_var.clone()),
138                    last_arg,
139                ),
140            );
141            self.statements.push(Statement::new(self.span, kind));
142            let dead_kind = RawStatement::StorageDead(len_var.local_id());
143            self.statements.push(Statement::new(self.span, dead_kind));
144            args.push(Operand::Copy(index_var));
145        } else {
146            args.push(last_arg);
147        }
148
149        // Call the indexing function:
150        // `storage_live(tmp1)`
151        // `tmp1 = {Array,Slice}{Mut,Shared}{Index,SubSlice}(move tmp0, <other args>)`
152        let output_var = {
153            let output_var = self.fresh_var(None, output_ty);
154            let index_call = Call {
155                func: indexing_function,
156                args,
157                dest: output_var.clone(),
158            };
159            let kind = RawStatement::Call(index_call);
160            self.statements.push(Statement::new(self.span, kind));
161            output_var
162        };
163
164        // Update the place.
165        *place = output_var.project(ProjectionElem::Deref, output_inner_ty);
166    }
167
168    /// Calls `self.visit_inner()` with `mutability` pushed on the stack.
169    fn visit_inner_with_mutability<T>(
170        &mut self,
171        x: &mut T,
172        mutability: bool,
173    ) -> ControlFlow<Infallible>
174    where
175        T: for<'s> DriveMut<'s, BodyVisitableWrapper<Self>>,
176    {
177        self.place_mutability_stack.push(mutability);
178        self.visit_inner(x)?;
179        self.place_mutability_stack.pop();
180        Continue(())
181    }
182}
183
184/// The visitor methods.
185impl VisitBodyMut for IndexVisitor<'_> {
186    /// We explore places from the inside-out.
187    fn exit_place(&mut self, place: &mut Place) {
188        // We have intercepted every traversal that would reach a place and pushed the correct
189        // mutability on the stack.
190        let mut_access = *self.place_mutability_stack.last().unwrap();
191        self.transform_place(mut_access, place);
192    }
193
194    fn visit_operand(&mut self, x: &mut Operand) -> ControlFlow<Infallible> {
195        match x {
196            Operand::Move(_) => self.visit_inner_with_mutability(x, true),
197            Operand::Copy(_) => self.visit_inner_with_mutability(x, false),
198            Operand::Const(..) => self.visit_inner(x),
199        }
200    }
201
202    fn visit_call(&mut self, x: &mut Call) -> ControlFlow<Infallible> {
203        self.visit_inner_with_mutability(x, true)
204    }
205
206    fn visit_fn_operand(&mut self, x: &mut FnOperand) -> ControlFlow<Infallible> {
207        match x {
208            FnOperand::Regular(_) => self.visit_inner(x),
209            FnOperand::Move(_) => self.visit_inner_with_mutability(x, true),
210        }
211    }
212
213    fn visit_rvalue(&mut self, x: &mut Rvalue) -> ControlFlow<Infallible> {
214        use Rvalue::*;
215        match x {
216            // `UniqueImmutable` de facto gives mutable access and only shows up if there is nested
217            // mutable access.
218            RawPtr(_, RefKind::Mut)
219            | Ref(_, BorrowKind::Mut | BorrowKind::TwoPhaseMut | BorrowKind::UniqueImmutable) => {
220                self.visit_inner_with_mutability(x, true)
221            }
222            RawPtr(_, RefKind::Shared)
223            | Ref(_, BorrowKind::Shared | BorrowKind::Shallow)
224            | Discriminant(..)
225            | Len(..) => self.visit_inner_with_mutability(x, false),
226
227            Use(_) | NullaryOp(..) | UnaryOp(..) | BinaryOp(..) | Aggregate(..) | Global(..)
228            | GlobalRef(..) | Repeat(..) | ShallowInitBox(..) => self.visit_inner(x),
229        }
230    }
231
232    fn visit_llbc_block(&mut self, _: &mut llbc_ast::Block) -> ControlFlow<Infallible> {
233        ControlFlow::Continue(())
234    }
235}
236
237pub struct Transform;
238
239/// We do the following.
240///
241/// If `p` is a projection (for instance: `var`, `*var`, `var.f`, etc.), we
242/// detect:
243/// - whether it operates on a slice or an array (we keep track of the types)
244/// - whether the access might mutate the value or not (it is
245///   the case if it is in a `move`, `&mut` or at the lhs of an assignment),
246///   and do the following transformations
247///
248/// ```text
249///   // If array and mutable access:
250///   ... p[i] ...
251///      ~~>
252///   tmp0 = &mut p
253///   tmp1 = ArrayIndexMut(move p, i)
254///   ... *tmp1 ...
255///
256///   // If array and non-mutable access:
257///   ... p[i] ...
258///      ~~>
259///   tmp0 := & p
260///   tmp1 := ArrayIndexShared(move tmp0, i)
261///   ... *tmp1 ...
262///
263///   // Omitting the slice cases, which are similar
264/// ```
265///
266/// For instance, it leads to the following transformations:
267/// ```text
268///   // x : [u32; N]
269///   y : u32 = copy x[i]
270///      ~~>
271///   tmp0 : & [u32; N] := &x
272///   tmp1 : &u32 = ArrayIndexShared(move tmp0, i)
273///   y : u32 = copy (*tmp1)
274///
275///   // x : &[T; N]
276///   y : &T = & (*x)[i]
277///      ~~>
278///   tmp0 : & [T; N] := & (*x)
279///   tmp1 : &T = ArrayIndexShared(move tmp0, i)
280///   y : &T = & (*tmp1)
281///
282///   // x : [u32; N]
283///   y = &mut x[i]
284///      ~~>
285///   tmp0 : &mut [u32; N] := &mut x
286///   tmp1 : &mut u32 := ArrayIndexMut(move tmp0, i)
287///   y = &mut (*tmp)
288///
289///   // When using an index on the lhs:
290///   // y : [T; N]
291///   y[i] = x
292///      ~~>
293///   tmp0 : &mut [T; N] := &mut y;
294///   tmp1 : &mut T = ArrayIndexMut(move y, i)
295///   *tmp1 = x
296/// ```
297impl LlbcPass for Transform {
298    fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
299        b.body.transform(|st: &mut Statement| {
300            let mut visitor = IndexVisitor {
301                locals: &mut b.locals,
302                statements: Vec::new(),
303                place_mutability_stack: Vec::new(),
304                span: st.span,
305            };
306            use RawStatement::*;
307            match &mut st.content {
308                Assign(..)
309                | SetDiscriminant(..)
310                | CopyNonOverlapping(_)
311                | Drop(..)
312                | Deinit(..)
313                | Call(..) => {
314                    let _ = visitor.visit_inner_with_mutability(st, true);
315                }
316                Switch(..) => {
317                    let _ = visitor.visit_inner_with_mutability(st, false);
318                }
319                Nop | Error(..) | Assert(..) | Abort(..) | StorageDead(..) | StorageLive(..)
320                | Return | Break(..) | Continue(..) | Loop(..) => {
321                    let _ = st.drive_body_mut(&mut visitor);
322                }
323            }
324            visitor.statements
325        });
326    }
327}