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, LlbcStatementTransformCtx};
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    ctx: &'b mut LlbcStatementTransformCtx<'a>,
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}
24
25impl<'a, 'b> IndexVisitor<'a, 'b> {
26    /// transform `place: subplace[i]` into indexing function calls for `subplace` and `i`
27    fn transform_place(&mut self, mut_access: bool, place: &mut Place) {
28        use ProjectionElem::*;
29        // This function is naturally called recusively, so `subplace` cannot be another `Index` or `Subslice`.
30        // Hence, `subplace`, if still projecting, must be either a `Deref` or a `Field`.
31        let Some((subplace, pe @ (Index { .. } | Subslice { .. }))) = place.as_projection() else {
32            return;
33        };
34        let tref = subplace.ty.as_adt().unwrap();
35        let builtin_ty = tref.id.as_builtin().unwrap();
36
37        // The built-in function to call.
38        let indexing_function = {
39            let builtin_fun = BuiltinFunId::Index(BuiltinIndexOp {
40                is_array: matches!(builtin_ty, BuiltinTy::Array),
41                mutability: RefKind::mutable(mut_access),
42                is_range: matches!(pe, Subslice { .. }),
43            });
44            // Same generics as the array/slice type, except for the extra lifetime.
45            let mut generics = tref.generics.clone();
46            generics.regions = [Region::Erased].into();
47            FnOperand::Regular(FnPtr::new(FnPtrKind::mk_builtin(builtin_fun), generics))
48        };
49
50        let elem_ty = tref.generics.types[0].clone();
51        let output_inner_ty = if matches!(pe, Index { .. }) {
52            elem_ty
53        } else {
54            TyKind::Adt(TypeDeclRef {
55                id: TypeId::Builtin(BuiltinTy::Slice),
56                generics: Box::new(GenericArgs::new_for_builtin(vec![elem_ty].into())),
57            })
58            .into_ty()
59        };
60        let output_ty = {
61            TyKind::Ref(
62                Region::Erased,
63                output_inner_ty.clone(),
64                RefKind::mutable(mut_access),
65            )
66            .into_ty()
67        };
68
69        // Push the statements:
70        // `storage_live(tmp0)`
71        // `tmp0 = &{mut}p`
72        let input_var =
73            self.ctx
74                .borrow_to_new_var(subplace.clone(), BorrowKind::mutable(mut_access), None);
75
76        // Construct the arguments to pass to the indexing function.
77        let mut args = vec![Operand::Move(input_var)];
78        if let Subslice { from, .. } = &pe {
79            args.push(from.as_ref().clone());
80        }
81        let (last_arg, from_end) = match &pe {
82            Index {
83                offset: x,
84                from_end,
85                ..
86            }
87            | Subslice {
88                to: x, from_end, ..
89            } => (x.as_ref().clone(), *from_end),
90            _ => unreachable!(),
91        };
92        let to_idx = self
93            .ctx
94            .compute_subslice_end_idx(subplace, last_arg, from_end);
95        args.push(to_idx);
96
97        // Call the indexing function:
98        // `storage_live(tmp1)`
99        // `tmp1 = {Array,Slice}{Mut,Shared}{Index,SubSlice}(move tmp0, <other args>)`
100        let output_var = {
101            let output_var = self.ctx.fresh_var(None, output_ty);
102            let index_call = Call {
103                func: indexing_function,
104                args,
105                dest: output_var.clone(),
106            };
107            let kind = StatementKind::Call(index_call);
108            self.ctx
109                .statements
110                .push(Statement::new(self.ctx.span, kind));
111            output_var
112        };
113
114        // Update the place.
115        *place = output_var.project(ProjectionElem::Deref, output_inner_ty);
116    }
117
118    /// Calls `self.visit_inner()` with `mutability` pushed on the stack.
119    fn visit_inner_with_mutability<T>(
120        &mut self,
121        x: &mut T,
122        mutability: bool,
123    ) -> ControlFlow<Infallible>
124    where
125        T: for<'s> DriveMut<'s, BodyVisitableWrapper<Self>> + BodyVisitable,
126    {
127        self.place_mutability_stack.push(mutability);
128        self.visit_inner(x)?;
129        self.place_mutability_stack.pop();
130        Continue(())
131    }
132}
133
134/// The visitor methods.
135impl VisitBodyMut for IndexVisitor<'_, '_> {
136    /// We explore places from the inside-out --- recursion naturally happens here.
137    fn exit_place(&mut self, place: &mut Place) {
138        // We have intercepted every traversal that would reach a place and pushed the correct
139        // mutability on the stack.
140        let mut_access = *self.place_mutability_stack.last().unwrap();
141        self.transform_place(mut_access, place);
142    }
143
144    fn visit_operand(&mut self, x: &mut Operand) -> ControlFlow<Infallible> {
145        match x {
146            Operand::Move(_) => self.visit_inner_with_mutability(x, true),
147            Operand::Copy(_) => self.visit_inner_with_mutability(x, false),
148            Operand::Const(..) => self.visit_inner(x),
149        }
150    }
151
152    fn visit_call(&mut self, x: &mut Call) -> ControlFlow<Infallible> {
153        self.visit_inner_with_mutability(x, true)
154    }
155
156    fn visit_fn_operand(&mut self, x: &mut FnOperand) -> ControlFlow<Infallible> {
157        match x {
158            FnOperand::Regular(_) => self.visit_inner(x),
159            FnOperand::Dynamic(_) => self.visit_inner_with_mutability(x, true),
160        }
161    }
162
163    fn visit_rvalue(&mut self, x: &mut Rvalue) -> ControlFlow<Infallible> {
164        use Rvalue::*;
165        match x {
166            // `UniqueImmutable` de facto gives mutable access and only shows up if there is nested
167            // mutable access.
168            RawPtr {
169                kind: RefKind::Mut, ..
170            }
171            | Ref {
172                kind: BorrowKind::Mut | BorrowKind::TwoPhaseMut | BorrowKind::UniqueImmutable,
173                ..
174            } => self.visit_inner_with_mutability(x, true),
175            RawPtr {
176                kind: RefKind::Shared,
177                ..
178            }
179            | Ref {
180                kind: BorrowKind::Shared | BorrowKind::Shallow,
181                ..
182            }
183            | Discriminant(..)
184            | Len(..) => self.visit_inner_with_mutability(x, false),
185
186            Use(_) | NullaryOp(..) | UnaryOp(..) | BinaryOp(..) | Aggregate(..) | Repeat(..)
187            | ShallowInitBox(..) => self.visit_inner(x),
188        }
189    }
190
191    fn visit_llbc_block(&mut self, _: &mut llbc_ast::Block) -> ControlFlow<Infallible> {
192        ControlFlow::Continue(())
193    }
194}
195
196/// We do the following.
197///
198/// If `p` is a projection (for instance: `var`, `*var`, `var.f`, etc.), we
199/// detect:
200/// - whether it operates on a slice or an array (we keep track of the types)
201/// - whether the access might mutate the value or not (it is
202///   the case if it is in a `move`, `&mut` or at the lhs of an assignment),
203///   and do the following transformations
204///
205/// ```text
206///   // If array and mutable access:
207///   ... p[i] ...
208///      ~~>
209///   tmp0 = &mut p
210///   tmp1 = ArrayIndexMut(move p, i)
211///   ... *tmp1 ...
212///
213///   // If array and non-mutable access:
214///   ... p[i] ...
215///      ~~>
216///   tmp0 := & p
217///   tmp1 := ArrayIndexShared(move tmp0, i)
218///   ... *tmp1 ...
219///
220///   // Omitting the slice cases, which are similar
221/// ```
222///
223/// For instance, it leads to the following transformations:
224/// ```text
225///   // x : [u32; N]
226///   y : u32 = copy x[i]
227///      ~~>
228///   tmp0 : & [u32; N] := &x
229///   tmp1 : &u32 = ArrayIndexShared(move tmp0, i)
230///   y : u32 = copy (*tmp1)
231///
232///   // x : &[T; N]
233///   y : &T = & (*x)[i]
234///      ~~>
235///   tmp0 : & [T; N] := & (*x)
236///   tmp1 : &T = ArrayIndexShared(move tmp0, i)
237///   y : &T = & (*tmp1)
238///
239///   // x : [u32; N]
240///   y = &mut x[i]
241///      ~~>
242///   tmp0 : &mut [u32; N] := &mut x
243///   tmp1 : &mut u32 := ArrayIndexMut(move tmp0, i)
244///   y = &mut (*tmp)
245///
246///   // When using an index on the lhs:
247///   // y : [T; N]
248///   y[i] = x
249///      ~~>
250///   tmp0 : &mut [T; N] := &mut y;
251///   tmp1 : &mut T = ArrayIndexMut(move y, i)
252///   *tmp1 = x
253/// ```
254pub struct Transform;
255impl LlbcPass for Transform {
256    fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
257        options.index_to_function_calls
258    }
259
260    fn transform_function(&self, ctx: &mut TransformCtx, decl: &mut FunDecl) {
261        decl.transform_llbc_statements(ctx, |ctx, st: &mut Statement| {
262            let mut visitor = IndexVisitor {
263                ctx,
264                place_mutability_stack: Vec::new(),
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        })
285    }
286}