charon_lib/transform/
index_to_function_calls.rs

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

charon_lib/transform/
index_to_function_calls.rs
