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.opts.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     let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
  30
  31     for block in basic_blocks {
  32         // lower `<[_]>::len` calls
  33         lower_slice_len_call(tcx, block, &*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: Some((dest, bb)),
  56             cleanup: None,
  57             from_hir_call: true,
  58             ..
  59         } => {
  60             // some heuristics for fast rejection
  61             if args.len() != 1 {
  62                 return;
  63             }
  64             let Some(arg) = args[0].place() else { return };
  65             let func_ty = func.ty(local_decls, tcx);
  66             match func_ty.kind() {
  67                 ty::FnDef(fn_def_id, _) if fn_def_id == &slice_len_fn_item_def_id => {
  68                     // perform modifications
  69                     // from something like `_5 = core::slice::<impl [u8]>::len(move _6) -> bb1`
  70                     // into `_5 = Len(*_6)
  71                     // goto bb1
  72
  73                     // make new RValue for Len
  74                     let deref_arg = tcx.mk_place_deref(arg);
  75                     let r_value = Rvalue::Len(deref_arg);
  76                     let len_statement_kind = StatementKind::Assign(Box::new((*dest, r_value)));
  77                     let add_statement =
  78                         Statement { kind: len_statement_kind, source_info: terminator.source_info };
  79
  80                     // modify terminator into simple Goto
  81                     let new_terminator_kind = TerminatorKind::Goto { target: *bb };
  82
  83                     let patch = SliceLenPatchInformation { add_statement, new_terminator_kind };
  84
  85                     patch_found = Some(patch);
  86                 }
  87                 _ => {}
  88             }
  89         }
  90         _ => {}
  91     }
  92
  93     if let Some(SliceLenPatchInformation { add_statement, new_terminator_kind }) = patch_found {
  94         block.statements.push(add_statement);
  95         block.terminator_mut().kind = new_terminator_kind;
  96     }
  97 }