Skip to main content

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