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     // The one successor remains unchanged, so no need to invalidate
  30     let (basic_blocks, local_decls, _) =
  31         body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate();
  32
  33     for block in basic_blocks {
  34         // lower `<[_]>::len` calls
  35         lower_slice_len_call(tcx, block, &*local_decls, slice_len_fn_item_def_id);
  36     }
  37 }
  38
  39 struct SliceLenPatchInformation<'tcx> {
  40     add_statement: Statement<'tcx>,
  41     new_terminator_kind: TerminatorKind<'tcx>,
  42 }
  43
  44 fn lower_slice_len_call<'tcx>(
  45     tcx: TyCtxt<'tcx>,
  46     block: &mut BasicBlockData<'tcx>,
  47     local_decls: &IndexVec<Local, LocalDecl<'tcx>>,
  48     slice_len_fn_item_def_id: DefId,
  49 ) {
  50     let mut patch_found: Option<SliceLenPatchInformation<'_>> = None;
  51
  52     let terminator = block.terminator();
  53     match &terminator.kind {
  54         TerminatorKind::Call {
  55             func,
  56             args,
  57             destination,
  58             target: Some(bb),
  59             cleanup: None,
  60             from_hir_call: true,
  61             ..
  62         } => {
  63             // some heuristics for fast rejection
  64             if args.len() != 1 {
  65                 return;
  66             }
  67             let Some(arg) = args[0].place() else { return };
  68             let func_ty = func.ty(local_decls, tcx);
  69             match func_ty.kind() {
  70                 ty::FnDef(fn_def_id, _) if fn_def_id == &slice_len_fn_item_def_id => {
  71                     // perform modifications
  72                     // from something like `_5 = core::slice::<impl [u8]>::len(move _6) -> bb1`
  73                     // into `_5 = Len(*_6)
  74                     // goto bb1
  75
  76                     // make new RValue for Len
  77                     let deref_arg = tcx.mk_place_deref(arg);
  78                     let r_value = Rvalue::Len(deref_arg);
  79                     let len_statement_kind =
  80                         StatementKind::Assign(Box::new((*destination, r_value)));
  81                     let add_statement =
  82                         Statement { kind: len_statement_kind, source_info: terminator.source_info };
  83
  84                     // modify terminator into simple Goto
  85                     let new_terminator_kind = TerminatorKind::Goto { target: *bb };
  86
  87                     let patch = SliceLenPatchInformation { add_statement, new_terminator_kind };
  88
  89                     patch_found = Some(patch);
  90                 }
  91                 _ => {}
  92             }
  93         }
  94         _ => {}
  95     }
  96
  97     if let Some(SliceLenPatchInformation { add_statement, new_terminator_kind }) = patch_found {
  98         block.statements.push(add_statement);
  99         block.terminator_mut().kind = new_terminator_kind;
 100     }
 101 }