charon_lib/transform/simplify_output/
index_to_function_calls.rs

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

charon_lib/transform/simplify_output/
index_to_function_calls.rs
