charon_lib/transform/simplify_output/
index_to_function_calls.rs1use crate::llbc_ast::*;
4use crate::transform::TransformCtx;
5use crate::transform::ctx::BodyTransformCtx;
6use derive_generic_visitor::*;
7
8use crate::transform::ctx::LlbcPass;
9
10#[derive(Visitor)]
16struct IndexVisitor<'a, 'b> {
17    locals: &'a mut Locals,
18    statements: Vec<Statement>,
20    place_mutability_stack: Vec<bool>,
25    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    fn transform_place(&mut self, mut_access: bool, place: &mut Place) {
61        use ProjectionElem::*;
62        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        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            let mut generics = tref.generics.clone();
79            generics.regions = [Region::Erased].into();
80            FnOperand::Regular(FnPtr::new(FnPtrKind::mk_builtin(builtin_fun), generics))
81        };
82
83        let input_ty = TyKind::Ref(
84            Region::Erased,
85            subplace.ty().clone(),
86            RefKind::mutable(mut_access),
87        )
88        .into_ty();
89
90        let elem_ty = tref.generics.types[0].clone();
91        let output_inner_ty = if matches!(pe, Index { .. }) {
92            elem_ty
93        } else {
94            TyKind::Adt(TypeDeclRef {
95                id: TypeId::Builtin(BuiltinTy::Slice),
96                generics: Box::new(GenericArgs::new_for_builtin(vec![elem_ty].into())),
97            })
98            .into_ty()
99        };
100        let output_ty = {
101            TyKind::Ref(
102                Region::Erased,
103                output_inner_ty.clone(),
104                RefKind::mutable(mut_access),
105            )
106            .into_ty()
107        };
108
109        let ptr_metadata = self.compute_place_metadata(subplace);
111
112        let input_var = {
116            let input_var = self.fresh_var(None, input_ty);
117            let kind = StatementKind::Assign(
118                input_var.clone(),
119                Rvalue::Ref {
120                    place: subplace.clone(),
121                    kind: BorrowKind::mutable(mut_access),
122                    ptr_metadata,
123                },
124            );
125            self.statements.push(Statement::new(self.span, kind));
126            input_var
127        };
128
129        let mut args = vec![Operand::Move(input_var)];
131        if let Subslice { from, .. } = &pe {
132            args.push(from.as_ref().clone());
133        }
134        let (last_arg, from_end) = match &pe {
135            Index {
136                offset: x,
137                from_end,
138                ..
139            }
140            | Subslice {
141                to: x, from_end, ..
142            } => (x.as_ref().clone(), *from_end),
143            _ => unreachable!(),
144        };
145        let to_idx = self.compute_subslice_end_idx(subplace, last_arg, from_end);
146        args.push(to_idx);
147
148        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 = StatementKind::Call(index_call);
159            self.statements.push(Statement::new(self.span, kind));
160            output_var
161        };
162
163        *place = output_var.project(ProjectionElem::Deref, output_inner_ty);
165    }
166
167    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>> + BodyVisitable,
175    {
176        self.place_mutability_stack.push(mutability);
177        self.visit_inner(x)?;
178        self.place_mutability_stack.pop();
179        Continue(())
180    }
181}
182
183impl VisitBodyMut for IndexVisitor<'_, '_> {
185    fn exit_place(&mut self, place: &mut Place) {
187        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            RawPtr {
218                kind: RefKind::Mut, ..
219            }
220            | Ref {
221                kind: BorrowKind::Mut | BorrowKind::TwoPhaseMut | BorrowKind::UniqueImmutable,
222                ..
223            } => self.visit_inner_with_mutability(x, true),
224            RawPtr {
225                kind: RefKind::Shared,
226                ..
227            }
228            | Ref {
229                kind: BorrowKind::Shared | BorrowKind::Shallow,
230                ..
231            }
232            | Discriminant(..)
233            | Len(..) => self.visit_inner_with_mutability(x, false),
234
235            Use(_) | NullaryOp(..) | UnaryOp(..) | BinaryOp(..) | Aggregate(..) | Repeat(..)
236            | ShallowInitBox(..) => self.visit_inner(x),
237        }
238    }
239
240    fn visit_llbc_block(&mut self, _: &mut llbc_ast::Block) -> ControlFlow<Infallible> {
241        ControlFlow::Continue(())
242    }
243}
244
245pub struct Transform;
246
247impl Transform {
248    fn transform_body_with_param(
249        &self,
250        ctx: &mut TransformCtx,
251        b: &mut ExprBody,
252        params: &GenericParams,
253    ) {
254        b.body.transform(|st: &mut Statement| {
255            let mut visitor = IndexVisitor {
256                locals: &mut b.locals,
257                statements: Vec::new(),
258                place_mutability_stack: Vec::new(),
259                span: st.span,
260                ctx: &ctx,
261                params,
262            };
263            use StatementKind::*;
264            match &mut st.kind {
265                Assign(..)
266                | SetDiscriminant(..)
267                | CopyNonOverlapping(_)
268                | Drop(..)
269                | Deinit(..)
270                | Call(..) => {
271                    let _ = visitor.visit_inner_with_mutability(st, true);
272                }
273                Switch(..) => {
274                    let _ = visitor.visit_inner_with_mutability(st, false);
275                }
276                Nop | Error(..) | Assert(..) | Abort(..) | StorageDead(..) | StorageLive(..)
277                | Return | Break(..) | Continue(..) | Loop(..) => {
278                    let _ = st.drive_body_mut(&mut visitor);
279                }
280            }
281            visitor.statements
282        });
283    }
284}
285
286impl LlbcPass for Transform {
345    fn transform_function(&self, ctx: &mut TransformCtx, decl: &mut FunDecl) {
346        if let Ok(body) = &mut decl.body {
347            self.transform_body_with_param(
348                ctx,
349                body.as_structured_mut().unwrap(),
350                &decl.signature.generics,
351            )
352        }
353    }
354}