1use std::collections::{HashMap, HashSet};
4
5use crate::llbc_ast::*;
6use crate::name_matcher::NamePattern;
7use crate::transform::ctx::{BodyTransformCtx, LlbcStatementTransformCtx};
8use crate::transform::{CowBox, TransformCtx};
9use derive_generic_visitor::*;
10use itertools::Itertools;
11
12use crate::transform::ctx::LlbcPass;
13
14fn mk_fn_ptr(ctx: &TransformCtx, id: ItemId, mut generics: GenericArgs) -> FnPtr {
15 if ctx.options.add_destruct_bounds
16 && let Some(item) = ctx.translated.get_item(id)
17 {
18 let trait_decl_ref = item
20 .generic_params()
21 .trait_clauses
22 .last()
23 .unwrap()
24 .trait_
25 .clone()
26 .substitute(&generics);
27 let kind = TraitRefKind::BuiltinOrAuto {
28 builtin_data: BuiltinImplData::UntrackedDestruct,
29 parent_trait_refs: Default::default(),
30 types: Default::default(),
31 vtable: None,
32 };
33 generics
34 .trait_refs
35 .push(TraitRef::new(kind, trait_decl_ref));
36 }
37 let fun_id = *id.as_fun().unwrap();
38 FnPtr::new(FnPtrKind::Fun(fun_id), generics)
39}
40
41fn method_impl_trait_ref(
43 ctx: &TransformCtx,
44 id: ItemId,
45 fun_generics: &GenericArgs,
46) -> Option<TraitRef> {
47 let fun_id = *id.as_fun()?;
48 let fun = ctx.translated.fun_decls.get(fun_id)?;
49 let FunSource::TraitImpl { impl_ref, .. } = &fun.src else {
50 return None;
51 };
52 let impl_ref = impl_ref.clone().substitute(fun_generics);
53 let trait_impl = ctx.translated.trait_impls.get(impl_ref.id)?;
54 let trait_decl_ref =
55 RegionBinder::empty(trait_impl.impl_trait.clone().substitute(&impl_ref.generics));
56 Some(TraitRef::new(
57 TraitRefKind::TraitImpl(impl_ref),
58 trait_decl_ref,
59 ))
60}
61
62fn index_method_types(
63 params: &GenericParams,
64 elem_ty: &Ty,
65 index_ty: &Ty,
66 output_ty: &Ty,
67) -> Option<IndexVec<TypeVarId, Ty>> {
68 let mut types: IndexVec<TypeVarId, Ty> =
69 [elem_ty.clone(), index_ty.clone()].into_iter().collect();
70 match params.types.len() {
71 3 => {
73 types.push(output_ty.clone());
74 }
75 2 => {}
76 _ => return None,
77 }
78 Some(types)
79}
80
81fn transform_operation(std_items: &Transform, ctx: &TransformCtx, statement: &mut Statement) {
82 match &statement.kind {
83 StatementKind::Assign(
85 place,
86 Rvalue::UnaryOp(
87 UnOp::Cast(CastKind::Unsize(src_ty, tgt_ty, UnsizingMetadata::Length(_))),
88 operand,
89 ),
90 ) => {
91 if let (TyKind::Ref(_, src_ty, src_kind), TyKind::Ref(_, tgt_ty, tgt_kind)) =
92 (src_ty.kind(), tgt_ty.kind())
93 && let TyKind::Array(elem_ty, len, elem_ty_is_sized) = src_ty.kind()
94 && let TyKind::Slice(..) = tgt_ty.kind()
95 {
96 assert!(src_kind == tgt_kind);
101 let item = match src_kind {
104 RefKind::Shared => StdItem::ArrayAsSlice,
105 RefKind::Mut => StdItem::ArrayAsMutSlice,
106 };
107 let Some(&fun_id) = std_items.item_map.get(&item) else {
108 return;
109 };
110 let generics = GenericArgs::new(
111 [Region::Erased].into(),
112 [elem_ty.clone()].into(),
113 [len.clone()].into(),
114 elem_ty_is_sized.iter().cloned().collect(),
115 );
116 statement.kind = StatementKind::Call {
117 call: Call {
118 func: FnOperand::Regular(mk_fn_ptr(ctx, fun_id, generics)),
119 args: vec![operand.clone()],
120 dest: place.clone(),
121 },
122 on_unwind: Block::new_unreachable(statement.span),
123 };
124 }
125 }
126 StatementKind::Assign(place, Rvalue::Repeat(operand, ty, len, ty_is_copy)) => {
128 let Some(ty_is_copy) = ty_is_copy else {
129 return;
130 };
131 let Some(&fun_id) = std_items.item_map.get(&StdItem::ArrayRepeat) else {
134 return;
135 };
136 let TyKind::Array(_, _, ty_is_sized) = place.ty().kind() else {
137 return;
138 };
139 let clone_clause_id = TraitClauseId::new(usize::from(ty_is_sized.is_some()));
142 let Some(ty_is_clone) = ty_is_copy
143 .clone()
144 .project_parent_clause(&ctx.translated, clone_clause_id)
145 else {
146 return;
147 };
148 let generics = GenericArgs::new(
149 [].into(),
150 [ty.clone()].into(),
151 [len.clone()].into(),
152 ty_is_sized.iter().cloned().chain([ty_is_clone]).collect(),
153 );
154 statement.kind = StatementKind::Call {
155 call: Call {
156 func: FnOperand::Regular(mk_fn_ptr(ctx, fun_id, generics)),
157 args: vec![operand.clone()],
158 dest: place.clone(),
159 },
160 on_unwind: Block::new_unreachable(statement.span),
161 };
162 }
163 _ => {}
164 }
165}
166
167#[derive(Visitor)]
173struct IndexVisitor<'a, 'b> {
174 ctx: &'b mut LlbcStatementTransformCtx<'a>,
175 std_items: &'b Transform,
176 place_mutability_stack: Vec<bool>,
181}
182
183impl<'a, 'b> IndexVisitor<'a, 'b> {
184 fn transform_place(&mut self, mut_access: bool, place: &mut Place) {
186 use ProjectionElem::*;
187 let Some((subplace, pe @ (Index { .. } | Subslice { .. }))) = place.as_projection() else {
190 return;
191 };
192
193 let (ty, len, ty_is_sized) = match subplace.ty.kind() {
194 TyKind::Array(ty, len, ty_is_sized) => (ty.clone(), Some(len.clone()), ty_is_sized),
195 TyKind::Slice(ty, ty_is_sized) => (ty.clone(), None, ty_is_sized),
196 _ => unreachable!("Indexing can only be done on arrays or slices"),
197 };
198
199 let mutability = RefKind::mutable(mut_access);
200 let item = match (pe.is_subslice(), mutability) {
201 (false, RefKind::Shared) => StdItem::SliceIndex,
202 (false, RefKind::Mut) => StdItem::SliceIndexMut,
203 (true, RefKind::Shared) => StdItem::RangeIndex,
204 (true, RefKind::Mut) => StdItem::RangeIndexMut,
205 };
206 let Some(&index_fun_id) = self.std_items.item_map.get(&item) else {
207 return;
208 };
209 let Some(index_fun) = index_fun_id
210 .as_fun()
211 .and_then(|id| self.ctx.ctx.translated.fun_decls.get(*id))
212 else {
213 return;
214 };
215 let index_generics = GenericArgs::new(
216 [Region::Erased].into(),
217 [ty.clone()].into(),
218 [].into(),
219 ty_is_sized.iter().cloned().collect(),
220 );
221 let index_fn_ptr = mk_fn_ptr(self.ctx.ctx, index_fun_id, index_generics);
222 let index_ty = index_fun.signature.inputs[0]
223 .clone()
224 .substitute(&index_fn_ptr.generics);
225
226 let output_inner_ty = if matches!(pe, Index { .. }) {
227 ty.clone()
228 } else {
229 TyKind::Slice(ty.clone(), ty_is_sized.clone()).into_ty()
230 };
231 let output_ty = {
232 TyKind::Ref(
233 Region::Erased,
234 output_inner_ty.clone(),
235 RefKind::mutable(mut_access),
236 )
237 .into_ty()
238 };
239
240 let input_var =
244 self.ctx
245 .borrow_to_new_var(subplace.clone(), BorrowKind::mutable(mut_access), None);
246
247 let (last_arg, from_end) = match &pe {
249 Index {
250 offset: x,
251 from_end,
252 ..
253 }
254 | Subslice {
255 to: x, from_end, ..
256 } => (x.as_ref().clone(), *from_end),
257 _ => unreachable!(),
258 };
259 let to_idx = self
260 .ctx
261 .compute_subslice_end_idx(subplace, last_arg, from_end);
262 let index = match &pe {
263 Index { .. } => to_idx,
264 Subslice { from, .. } => {
265 let Some(range_ref) = index_ty.as_adt().cloned() else {
266 return;
267 };
268 let range_ty = TyKind::Adt(range_ref.clone()).into_ty();
269 let range_var = self.ctx.fresh_var(None, range_ty);
270 self.ctx.insert_assn_stmt(
271 range_var.clone(),
272 Rvalue::Aggregate(
273 AggregateKind::Adt(range_ref, None, None),
274 vec![from.as_ref().clone(), to_idx],
275 ),
276 );
277 Operand::Move(range_var)
278 }
279 _ => unreachable!(),
280 };
281
282 let (index_fn_ptr, args) = if let Some(len) = len {
283 let (array_item, slice_item) = match mutability {
284 RefKind::Shared => (StdItem::ArrayIndex, StdItem::SliceIndexImpl),
285 RefKind::Mut => (StdItem::ArrayIndexMut, StdItem::SliceIndexMutImpl),
286 };
287 let Some(&array_fun_id) = self.std_items.item_map.get(&array_item) else {
288 return;
289 };
290 let Some(array_fun) = array_fun_id
291 .as_fun()
292 .and_then(|id| self.ctx.ctx.translated.fun_decls.get(*id))
293 else {
294 return;
295 };
296 let Some(&slice_fun_id) = self.std_items.item_map.get(&slice_item) else {
297 return;
298 };
299 let Some(slice_fun) = slice_fun_id
300 .as_fun()
301 .and_then(|id| self.ctx.ctx.translated.fun_decls.get(*id))
302 else {
303 return;
304 };
305
306 let Some(slice_types) =
307 index_method_types(&slice_fun.generics, &ty, &index_ty, &output_inner_ty)
308 else {
309 return;
310 };
311 let mut slice_generics =
312 GenericArgs::new([Region::Erased].into(), slice_types, [].into(), [].into());
313 let Some(slice_index_trait_ref) =
314 method_impl_trait_ref(self.ctx.ctx, index_fun_id, &index_fn_ptr.generics)
315 else {
316 return;
317 };
318 let index_ty_is_sized = if ty_is_sized.is_some() {
319 let Some(clause) = slice_fun.generics.trait_clauses.get(TraitClauseId::new(1))
320 else {
321 return;
322 };
323 let trait_decl_ref = clause.trait_.clone().substitute_explicits(&slice_generics);
324 let Some(meta_sized) = slice_index_trait_ref
325 .clone()
326 .project_parent_clause(&self.ctx.ctx.translated, TraitClauseId::ZERO)
327 else {
328 return;
329 };
330 Some(TraitRef::new(
331 TraitRefKind::BuiltinOrAuto {
332 builtin_data: BuiltinImplData::Sized,
333 parent_trait_refs: [meta_sized].into_iter().collect(),
334 types: Default::default(),
335 vtable: None,
336 },
337 trait_decl_ref,
338 ))
339 } else {
340 None
341 };
342 slice_generics.trait_refs = ty_is_sized
343 .iter()
344 .cloned()
345 .chain(index_ty_is_sized.iter().cloned())
346 .chain([slice_index_trait_ref])
347 .collect();
348 let slice_fn_ptr = mk_fn_ptr(self.ctx.ctx, slice_fun_id, slice_generics);
349 let Some(slice_trait_ref) =
350 method_impl_trait_ref(self.ctx.ctx, slice_fun_id, &slice_fn_ptr.generics)
351 else {
352 return;
353 };
354
355 let Some(array_types) =
356 index_method_types(&array_fun.generics, &ty, &index_ty, &output_inner_ty)
357 else {
358 return;
359 };
360 let array_generics = GenericArgs::new(
361 [Region::Erased].into(),
362 array_types,
363 [len].into(),
364 ty_is_sized
365 .iter()
366 .cloned()
367 .chain(index_ty_is_sized)
368 .chain([slice_trait_ref])
369 .collect(),
370 );
371 (
372 mk_fn_ptr(self.ctx.ctx, array_fun_id, array_generics),
373 vec![Operand::Move(input_var), index],
374 )
375 } else {
376 (index_fn_ptr, vec![index, Operand::Move(input_var)])
377 };
378
379 let output_var = {
383 let output_var = self.ctx.fresh_var(None, output_ty);
384 let index_call = Call {
385 func: FnOperand::Regular(index_fn_ptr),
386 args,
387 dest: output_var.clone(),
388 };
389 let kind = StatementKind::Call {
390 call: index_call,
391 on_unwind: Block::new_unreachable(self.ctx.span),
392 };
393 self.ctx
394 .statements
395 .push(Statement::new(self.ctx.span, kind));
396 output_var
397 };
398
399 *place = output_var.project(ProjectionElem::Deref, output_inner_ty);
401 }
402
403 fn visit_inner_with_mutability<T>(
405 &mut self,
406 x: &mut T,
407 mutability: bool,
408 ) -> ControlFlow<Infallible>
409 where
410 T: for<'s> DriveMut<'s, BodyVisitableWrapper<Self>> + BodyVisitable,
411 {
412 self.place_mutability_stack.push(mutability);
413 self.visit_inner(x)?;
414 self.place_mutability_stack.pop();
415 Continue(())
416 }
417}
418
419impl VisitBodyMut for IndexVisitor<'_, '_> {
421 fn exit_place(&mut self, place: &mut Place) {
423 let mut_access = *self.place_mutability_stack.last().unwrap();
426 self.transform_place(mut_access, place);
427 }
428
429 fn visit_operand(&mut self, x: &mut Operand) -> ControlFlow<Infallible> {
430 match x {
431 Operand::Move(_) => self.visit_inner_with_mutability(x, true),
432 Operand::Copy(_) => self.visit_inner_with_mutability(x, false),
433 Operand::Const(..) => self.visit_inner(x),
434 }
435 }
436
437 fn visit_call(&mut self, x: &mut Call) -> ControlFlow<Infallible> {
438 self.visit_inner_with_mutability(x, true)
439 }
440
441 fn visit_fn_operand(&mut self, x: &mut FnOperand) -> ControlFlow<Infallible> {
442 match x {
443 FnOperand::Regular(_) => self.visit_inner(x),
444 FnOperand::Dynamic(_) => self.visit_inner_with_mutability(x, true),
445 }
446 }
447
448 fn visit_rvalue(&mut self, x: &mut Rvalue) -> ControlFlow<Infallible> {
449 use Rvalue::*;
450 match x {
451 RawPtr {
454 kind: RefKind::Mut, ..
455 }
456 | Ref {
457 kind: BorrowKind::Mut | BorrowKind::TwoPhaseMut | BorrowKind::UniqueImmutable,
458 ..
459 } => self.visit_inner_with_mutability(x, true),
460 RawPtr {
461 kind: RefKind::Shared,
462 ..
463 }
464 | Ref {
465 kind: BorrowKind::Shared | BorrowKind::Shallow,
466 ..
467 }
468 | Discriminant(..)
469 | Len(..) => self.visit_inner_with_mutability(x, false),
470
471 Use(..) | NullaryOp(..) | UnaryOp(..) | BinaryOp(..) | Aggregate(..) | Repeat(..) => {
472 self.visit_inner(x)
473 }
474 }
475 }
476
477 fn visit_llbc_block(&mut self, _: &mut llbc_ast::Block) -> ControlFlow<Infallible> {
478 ControlFlow::Continue(())
479 }
480}
481
482#[derive(Clone, Copy, PartialEq, Eq, Hash)]
541enum StdItem {
542 ArrayAsSlice,
543 ArrayAsMutSlice,
544 ArrayRepeat,
545 ArrayIndex,
546 ArrayIndexMut,
547 SliceIndexImpl,
548 SliceIndexMutImpl,
549 SliceIndex,
550 SliceIndexMut,
551 RangeIndex,
552 RangeIndexMut,
553}
554
555pub struct Transform {
556 item_map: HashMap<StdItem, ItemId>,
557 item_set: HashSet<ItemId>,
558}
559
560impl Transform {
561 pub fn new(ctx: &TransformCtx) -> CowBox<dyn LlbcPass> {
562 use StdItem::*;
563
564 let mut matches: [(StdItem, NamePattern, Vec<ItemId>); _] = [
565 (ArrayAsSlice, "core::array::_::as_slice"),
566 (ArrayAsMutSlice, "core::array::_::as_mut_slice"),
567 (ArrayRepeat, "core::array::repeat"),
568 (
569 SliceIndex,
570 "core::slice::index::{impl core::slice::index::SliceIndex<_> for usize}::index",
571 ),
572 (
573 SliceIndexMut,
574 "core::slice::index::{impl core::slice::index::SliceIndex<_> for usize}::index_mut",
575 ),
576 (
577 RangeIndex,
578 "core::slice::index::{impl core::slice::index::SliceIndex<_> for core::ops::range::Range<usize>}::index",
579 ),
580 (
581 RangeIndexMut,
582 "core::slice::index::{impl core::slice::index::SliceIndex<_> for core::ops::range::Range<usize>}::index_mut",
583 ),
584 ]
585 .map(|(item, path)| (item, NamePattern::parse(path).unwrap(), Vec::new()));
586
587 for (id, name) in &ctx.translated.item_names {
589 for (_, pattern, found) in &mut matches {
590 if pattern.matches(&ctx.translated, name) {
591 found.push(*id);
592 }
593 }
594 }
595
596 let mut index_impl_methods: HashMap<StdItem, Vec<ItemId>> = HashMap::new();
597 for fun in &ctx.translated.fun_decls {
598 let FunSource::TraitImpl { trait_ref, .. } = &fun.src else {
599 continue;
600 };
601 let Some(trait_decl) = ctx.translated.trait_decls.get(trait_ref.id) else {
602 continue;
603 };
604 let Some(self_ty) = trait_ref.self_ty(&ctx.translated) else {
605 continue;
606 };
607 let item = match (&trait_decl.item_meta.lang_item, self_ty.kind()) {
608 (Some(crate::ast::from_rustc::LangItem::Index), TyKind::Array(..)) => ArrayIndex,
609 (Some(crate::ast::from_rustc::LangItem::IndexMut), TyKind::Array(..)) => {
610 ArrayIndexMut
611 }
612 (Some(crate::ast::from_rustc::LangItem::Index), TyKind::Slice(..)) => {
613 SliceIndexImpl
614 }
615 (Some(crate::ast::from_rustc::LangItem::IndexMut), TyKind::Slice(..)) => {
616 SliceIndexMutImpl
617 }
618 _ => continue,
619 };
620 index_impl_methods
621 .entry(item)
622 .or_default()
623 .push(ItemId::Fun(fun.def_id));
624 }
625 let item_map = matches
626 .into_iter()
627 .map(|(item, _, found)| (item, found))
628 .chain(index_impl_methods)
629 .filter_map(|(item, found)| found.into_iter().exactly_one().ok().map(|id| (item, id)))
630 .collect::<HashMap<_, _>>();
631 let item_set = item_map.values().copied().collect();
632 CowBox::Owned(Box::new(Self { item_map, item_set }))
633 }
634}
635
636impl LlbcPass for Transform {
637 fn should_run(&self, options: &crate::options::TranslateOptions) -> bool {
638 options.ops_to_function_calls || options.index_to_function_calls
639 }
640
641 fn transform_function(&self, ctx: &mut TransformCtx, decl: &mut FunDecl) {
642 if self.item_set.contains(&ItemId::Fun(decl.def_id)) {
643 return;
644 }
645 let Some(body) = decl.body.as_structured_mut() else {
646 return;
647 };
648 if ctx.options.ops_to_function_calls {
649 body.body
650 .visit_statements(&mut |statement: &mut Statement| {
651 transform_operation(self, ctx, statement)
652 });
653 }
654 if ctx.options.index_to_function_calls {
655 decl.transform_llbc_statements(ctx, |ctx, st: &mut Statement| {
656 let mut visitor = IndexVisitor {
657 ctx,
658 std_items: self,
659 place_mutability_stack: Vec::new(),
660 };
661 use StatementKind::*;
662 match &mut st.kind {
663 Assign(..) | SetDiscriminant(..) | Drop { .. } | Call { .. } => {
664 let _ = visitor.visit_inner_with_mutability(st, true);
665 }
666 Switch { .. } | PlaceMention(..) | Borrowck(..) => {
667 let _ = visitor.visit_inner_with_mutability(st, false);
668 }
669 Nop
670 | UnwindResume
671 | Error(..)
672 | InlineAsm { .. }
673 | Assert { .. }
674 | Abort(..)
675 | StorageDead(..)
676 | StorageLive(..)
677 | Return
678 | Break(..)
679 | Continue(..)
680 | Loop(..) => {
681 let _ = st.drive_body_mut(&mut visitor);
682 }
683 }
684 })
685 }
686 }
687}