1use crate::llbc_ast::*;
3use crate::transform::TransformCtx;
4use derive_generic_visitor::*;
5
6use super::ctx::LlbcPass;
7
8#[derive(Visitor)]
14struct IndexVisitor<'a> {
15 locals: &'a mut Locals,
16 statements: Vec<Statement>,
18 place_mutability_stack: Vec<bool>,
23 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 tref = subplace.ty.as_adt().unwrap();
41 let builtin_ty = tref.id.as_builtin().unwrap();
42
43 let indexing_function = {
45 let builtin_fun = BuiltinFunId::Index(BuiltinIndexOp {
46 is_array: matches!(builtin_ty, BuiltinTy::Array),
47 mutability: RefKind::mutable(mut_access),
48 is_range: matches!(pe, Subslice { .. }),
49 });
50 let mut generics = tref.generics.clone();
52 generics.regions = [Region::Erased].into();
53 FnOperand::Regular(FnPtr {
54 func: Box::new(FunIdOrTraitMethodRef::mk_builtin(builtin_fun)),
55 generics,
56 })
57 };
58
59 let input_ty = TyKind::Ref(
60 Region::Erased,
61 subplace.ty().clone(),
62 RefKind::mutable(mut_access),
63 )
64 .into_ty();
65
66 let elem_ty = tref.generics.types[0].clone();
67 let output_inner_ty = if matches!(pe, Index { .. }) {
68 elem_ty
69 } else {
70 TyKind::Adt(TypeDeclRef {
71 id: TypeId::Builtin(BuiltinTy::Slice),
72 generics: Box::new(GenericArgs::new_for_builtin(vec![elem_ty].into())),
73 })
74 .into_ty()
75 };
76 let output_ty = {
77 TyKind::Ref(
78 Region::Erased,
79 output_inner_ty.clone(),
80 RefKind::mutable(mut_access),
81 )
82 .into_ty()
83 };
84
85 let input_var = {
89 let input_var = self.fresh_var(None, input_ty);
90 let kind = RawStatement::Assign(
91 input_var.clone(),
92 Rvalue::Ref(subplace.clone(), BorrowKind::mutable(mut_access)),
93 );
94 self.statements.push(Statement::new(self.span, kind));
95 input_var
96 };
97
98 let mut args = vec![Operand::Move(input_var)];
100 if let Subslice { from, .. } = &pe {
101 args.push(from.as_ref().clone());
102 }
103 let (last_arg, from_end) = match &pe {
104 Index {
105 offset: x,
106 from_end,
107 ..
108 }
109 | Subslice {
110 to: x, from_end, ..
111 } => (x.as_ref().clone(), *from_end),
112 _ => unreachable!(),
113 };
114 if from_end {
115 let usize_ty = TyKind::Literal(LiteralTy::Integer(IntegerTy::Usize)).into_ty();
118 let len_var = self.fresh_var(None, usize_ty.clone());
119 let kind = RawStatement::Assign(
120 len_var.clone(),
121 Rvalue::Len(
122 subplace.clone(),
123 subplace.ty().clone(),
124 tref.generics.const_generics.get(0.into()).cloned(),
125 ),
126 );
127 self.statements.push(Statement::new(self.span, kind));
128
129 let index_var = self.fresh_var(None, usize_ty);
133 let kind = RawStatement::Assign(
134 index_var.clone(),
135 Rvalue::BinaryOp(BinOp::Sub, Operand::Copy(len_var.clone()), last_arg),
136 );
137 self.statements.push(Statement::new(self.span, kind));
138 let dead_kind = RawStatement::StorageDead(len_var.local_id());
139 self.statements.push(Statement::new(self.span, dead_kind));
140 args.push(Operand::Copy(index_var));
141 } else {
142 args.push(last_arg);
143 }
144
145 let output_var = {
149 let output_var = self.fresh_var(None, output_ty);
150 let index_call = Call {
151 func: indexing_function,
152 args,
153 dest: output_var.clone(),
154 };
155 let kind = RawStatement::Call(index_call);
156 self.statements.push(Statement::new(self.span, kind));
157 output_var
158 };
159
160 *place = output_var.project(ProjectionElem::Deref, output_inner_ty);
162 }
163
164 fn visit_inner_with_mutability<T>(
166 &mut self,
167 x: &mut T,
168 mutability: bool,
169 ) -> ControlFlow<Infallible>
170 where
171 T: for<'s> DriveMut<'s, BodyVisitableWrapper<Self>>,
172 {
173 self.place_mutability_stack.push(mutability);
174 self.visit_inner(x)?;
175 self.place_mutability_stack.pop();
176 Continue(())
177 }
178}
179
180impl VisitBodyMut for IndexVisitor<'_> {
182 fn exit_place(&mut self, place: &mut Place) {
184 let mut_access = *self.place_mutability_stack.last().unwrap();
187 self.transform_place(mut_access, place);
188 }
189
190 fn visit_operand(&mut self, x: &mut Operand) -> ControlFlow<Infallible> {
191 match x {
192 Operand::Move(_) => self.visit_inner_with_mutability(x, true),
193 Operand::Copy(_) => self.visit_inner_with_mutability(x, false),
194 Operand::Const(..) => self.visit_inner(x),
195 }
196 }
197
198 fn visit_call(&mut self, x: &mut Call) -> ControlFlow<Infallible> {
199 self.visit_inner_with_mutability(x, true)
200 }
201
202 fn visit_fn_operand(&mut self, x: &mut FnOperand) -> ControlFlow<Infallible> {
203 match x {
204 FnOperand::Regular(_) => self.visit_inner(x),
205 FnOperand::Move(_) => self.visit_inner_with_mutability(x, true),
206 }
207 }
208
209 fn visit_rvalue(&mut self, x: &mut Rvalue) -> ControlFlow<Infallible> {
210 use Rvalue::*;
211 match x {
212 RawPtr(_, RefKind::Mut)
215 | Ref(_, BorrowKind::Mut | BorrowKind::TwoPhaseMut | BorrowKind::UniqueImmutable) => {
216 self.visit_inner_with_mutability(x, true)
217 }
218 RawPtr(_, RefKind::Shared)
219 | Ref(_, BorrowKind::Shared | BorrowKind::Shallow)
220 | Discriminant(..)
221 | Len(..) => self.visit_inner_with_mutability(x, false),
222
223 Use(_) | NullaryOp(..) | UnaryOp(..) | BinaryOp(..) | Aggregate(..) | Global(..)
224 | GlobalRef(..) | Repeat(..) | ShallowInitBox(..) => self.visit_inner(x),
225 }
226 }
227
228 fn visit_llbc_block(&mut self, _: &mut llbc_ast::Block) -> ControlFlow<Infallible> {
229 ControlFlow::Continue(())
230 }
231}
232
233pub struct Transform;
234
235impl LlbcPass for Transform {
294 fn transform_body(&self, _ctx: &mut TransformCtx, b: &mut ExprBody) {
295 b.body.transform(|st: &mut Statement| {
296 let mut visitor = IndexVisitor {
297 locals: &mut b.locals,
298 statements: Vec::new(),
299 place_mutability_stack: Vec::new(),
300 span: st.span,
301 };
302 use RawStatement::*;
303 match &mut st.content {
304 Assign(..)
305 | SetDiscriminant(..)
306 | CopyNonOverlapping(_)
307 | Drop(..)
308 | Deinit(..)
309 | Call(..) => {
310 let _ = visitor.visit_inner_with_mutability(st, true);
311 }
312 Switch(..) => {
313 let _ = visitor.visit_inner_with_mutability(st, false);
314 }
315 Nop | Error(..) | Assert(..) | Abort(..) | StorageDead(..) | StorageLive(..)
316 | Return | Break(..) | Continue(..) | Loop(..) => {
317 let _ = st.drive_body_mut(&mut visitor);
318 }
319 }
320 visitor.statements
321 });
322 }
323}