compiler/rustc_mir_transform/src/lower_slice_len.rs

   1 //! This pass lowers calls to core::slice::len to just Len op.
   2 //! It should run before inlining!
   3
   4 use crate::MirPass;
   5 use rustc_hir::def_id::DefId;
   6 use rustc_index::vec::IndexVec;
   7 use rustc_middle::mir::*;
   8 use rustc_middle::ty::{self, TyCtxt};
   9
  10 pub struct LowerSliceLenCalls;
  11
  12 impl<'tcx> MirPass<'tcx> for LowerSliceLenCalls {
  13     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  14         sess.mir_opt_level() > 0
  15     }
  16
  17     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  18         lower_slice_len_calls(tcx, body)
  19     }
  20 }
  21
  22 pub fn lower_slice_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  23     let language_items = tcx.lang_items();
  24     let Some(slice_len_fn_item_def_id) = language_items.slice_len_fn() else {
  25         // there is no language item to compare to :)
  26         return;
  27     };
  28
  29     // The one successor remains unchanged, so no need to invalidate
  30     let basic_blocks = body.basic_blocks.as_mut_preserves_cfg();
  31     for block in basic_blocks {
  32         // lower `<[_]>::len` calls
  33         lower_slice_len_call(tcx, block, &body.local_decls, slice_len_fn_item_def_id);
  34     }
  35 }
  36
  37 struct SliceLenPatchInformation<'tcx> {
  38     add_statement: Statement<'tcx>,
  39     new_terminator_kind: TerminatorKind<'tcx>,
  40 }
  41
  42 fn lower_slice_len_call<'tcx>(
  43     tcx: TyCtxt<'tcx>,
  44     block: &mut BasicBlockData<'tcx>,
  45     local_decls: &IndexVec<Local, LocalDecl<'tcx>>,
  46     slice_len_fn_item_def_id: DefId,
  47 ) {
  48     let mut patch_found: Option<SliceLenPatchInformation<'_>> = None;
  49
  50     let terminator = block.terminator();
  51     match &terminator.kind {
  52         TerminatorKind::Call {
  53             func,
  54             args,
  55             destination,
  56             target: Some(bb),
  57             cleanup: None,
  58             from_hir_call: true,
  59             ..
  60         } => {
  61             // some heuristics for fast rejection
  62             if args.len() != 1 {
  63                 return;
  64             }
  65             let Some(arg) = args[0].place() else { return };
  66             let func_ty = func.ty(local_decls, tcx);
  67             match func_ty.kind() {
  68                 ty::FnDef(fn_def_id, _) if fn_def_id == &slice_len_fn_item_def_id => {
  69                     // perform modifications
  70                     // from something like `_5 = core::slice::<impl [u8]>::len(move _6) -> bb1`
  71                     // into `_5 = Len(*_6)
  72                     // goto bb1
  73
  74                     // make new RValue for Len
  75                     let deref_arg = tcx.mk_place_deref(arg);
  76                     let r_value = Rvalue::Len(deref_arg);
  77                     let len_statement_kind =
  78                         StatementKind::Assign(Box::new((*destination, r_value)));
  79                     let add_statement =
  80                         Statement { kind: len_statement_kind, source_info: terminator.source_info };
  81
  82                     // modify terminator into simple Goto
  83                     let new_terminator_kind = TerminatorKind::Goto { target: *bb };
  84
  85                     let patch = SliceLenPatchInformation { add_statement, new_terminator_kind };
  86
  87                     patch_found = Some(patch);
  88                 }
  89                 _ => {}
  90             }
  91         }
  92         _ => {}
  93     }
  94
  95     if let Some(SliceLenPatchInformation { add_statement, new_terminator_kind }) = patch_found {
  96         block.statements.push(add_statement);
  97         block.terminator_mut().kind = new_terminator_kind;
  98     }
  99 }