Skip to content

Commit

Permalink
Auto merge of rust-lang#86525 - shamatar:array_len_opt, r=oli-obk
Browse files Browse the repository at this point in the history
Array `.len()` MIR optimization pass

This pass kind-of works back the `[T; N].len()` call that at the moment is first coerced as `&[T; N]` -> `&[T]` and then uses `&[T].len()`. Depends on rust-lang#86383
  • Loading branch information
bors committed Oct 7, 2021
2 parents ca8078d + a31518f commit 680ff86
Show file tree
Hide file tree
Showing 18 changed files with 979 additions and 9 deletions.
2 changes: 2 additions & 0 deletions compiler/rustc_mir_transform/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ mod lower_intrinsics;
mod lower_slice_len;
mod match_branches;
mod multiple_return_terminators;
mod normalize_array_len;
mod nrvo;
mod remove_noop_landing_pads;
mod remove_storage_markers;
Expand Down Expand Up @@ -488,6 +489,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
// machine than on MIR with async primitives.
let optimizations_with_generators: &[&dyn MirPass<'tcx>] = &[
&lower_slice_len::LowerSliceLenCalls, // has to be done before inlining, otherwise actual call will be almost always inlined. Also simple, so can just do first
&normalize_array_len::NormalizeArrayLen, // has to run after `slice::len` lowering
&unreachable_prop::UnreachablePropagation,
&uninhabited_enum_branching::UninhabitedEnumBranching,
&simplify::SimplifyCfg::new("after-uninhabited-enum-branching"),
Expand Down
287 changes: 287 additions & 0 deletions compiler/rustc_mir_transform/src/normalize_array_len.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,287 @@
//! This pass eliminates casting of arrays into slices when their length
//! is taken using `.len()` method. Handy to preserve information in MIR for const prop
use crate::MirPass;
use rustc_data_structures::fx::FxIndexMap;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::IndexVec;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, TyCtxt};

const MAX_NUM_BLOCKS: usize = 800;
const MAX_NUM_LOCALS: usize = 3000;

pub struct NormalizeArrayLen;

impl<'tcx> MirPass<'tcx> for NormalizeArrayLen {
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
if tcx.sess.mir_opt_level() < 4 {
return;
}

// early returns for edge cases of highly unrolled functions
if body.basic_blocks().len() > MAX_NUM_BLOCKS {
return;
}
if body.local_decls().len() > MAX_NUM_LOCALS {
return;
}
normalize_array_len_calls(tcx, body)
}
}

pub fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();

// do a preliminary analysis to see if we ever have locals of type `[T;N]` or `&[T;N]`
let mut interesting_locals = BitSet::new_empty(local_decls.len());
for (local, decl) in local_decls.iter_enumerated() {
match decl.ty.kind() {
ty::Array(..) => {
interesting_locals.insert(local);
}
ty::Ref(.., ty, Mutability::Not) => match ty.kind() {
ty::Array(..) => {
interesting_locals.insert(local);
}
_ => {}
},
_ => {}
}
}
if interesting_locals.is_empty() {
// we have found nothing to analyze
return;
}
let num_intesting_locals = interesting_locals.count();
let mut state = FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
let mut patches_scratchpad =
FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
let mut replacements_scratchpad =
FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
for block in basic_blocks {
// make length calls for arrays [T; N] not to decay into length calls for &[T]
// that forbids constant propagation
normalize_array_len_call(
tcx,
block,
local_decls,
&interesting_locals,
&mut state,
&mut patches_scratchpad,
&mut replacements_scratchpad,
);
state.clear();
patches_scratchpad.clear();
replacements_scratchpad.clear();
}
}

struct Patcher<'a, 'tcx> {
tcx: TyCtxt<'tcx>,
patches_scratchpad: &'a FxIndexMap<usize, usize>,
replacements_scratchpad: &'a mut FxIndexMap<usize, Local>,
local_decls: &'a mut IndexVec<Local, LocalDecl<'tcx>>,
statement_idx: usize,
}

impl<'a, 'tcx> Patcher<'a, 'tcx> {
fn patch_expand_statement(
&mut self,
statement: &mut Statement<'tcx>,
) -> Option<std::vec::IntoIter<Statement<'tcx>>> {
let idx = self.statement_idx;
if let Some(len_statemnt_idx) = self.patches_scratchpad.get(&idx).copied() {
let mut statements = Vec::with_capacity(2);

// we are at statement that performs a cast. The only sound way is
// to create another local that performs a similar copy without a cast and then
// use this copy in the Len operation

match &statement.kind {
StatementKind::Assign(box (
..,
Rvalue::Cast(
CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
operand,
_,
),
)) => {
match operand {
Operand::Copy(place) | Operand::Move(place) => {
// create new local
let ty = operand.ty(self.local_decls, self.tcx);
let local_decl =
LocalDecl::with_source_info(ty, statement.source_info.clone());
let local = self.local_decls.push(local_decl);
// make it live
let mut make_live_statement = statement.clone();
make_live_statement.kind = StatementKind::StorageLive(local);
statements.push(make_live_statement);
// copy into it

let operand = Operand::Copy(*place);
let mut make_copy_statement = statement.clone();
let assign_to = Place::from(local);
let rvalue = Rvalue::Use(operand);
make_copy_statement.kind =
StatementKind::Assign(box (assign_to, rvalue));
statements.push(make_copy_statement);

// to reorder we have to copy and make NOP
statements.push(statement.clone());
statement.make_nop();

self.replacements_scratchpad.insert(len_statemnt_idx, local);
}
_ => {
unreachable!("it's a bug in the implementation")
}
}
}
_ => {
unreachable!("it's a bug in the implementation")
}
}

self.statement_idx += 1;

Some(statements.into_iter())
} else if let Some(local) = self.replacements_scratchpad.get(&idx).copied() {
let mut statements = Vec::with_capacity(2);

match &statement.kind {
StatementKind::Assign(box (into, Rvalue::Len(place))) => {
let add_deref = if let Some(..) = place.as_local() {
false
} else if let Some(..) = place.local_or_deref_local() {
true
} else {
unreachable!("it's a bug in the implementation")
};
// replace len statement
let mut len_statement = statement.clone();
let mut place = Place::from(local);
if add_deref {
place = self.tcx.mk_place_deref(place);
}
len_statement.kind = StatementKind::Assign(box (*into, Rvalue::Len(place)));
statements.push(len_statement);

// make temporary dead
let mut make_dead_statement = statement.clone();
make_dead_statement.kind = StatementKind::StorageDead(local);
statements.push(make_dead_statement);

// make original statement NOP
statement.make_nop();
}
_ => {
unreachable!("it's a bug in the implementation")
}
}

self.statement_idx += 1;

Some(statements.into_iter())
} else {
self.statement_idx += 1;
None
}
}
}

fn normalize_array_len_call<'tcx>(
tcx: TyCtxt<'tcx>,
block: &mut BasicBlockData<'tcx>,
local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
interesting_locals: &BitSet<Local>,
state: &mut FxIndexMap<Local, usize>,
patches_scratchpad: &mut FxIndexMap<usize, usize>,
replacements_scratchpad: &mut FxIndexMap<usize, Local>,
) {
for (statement_idx, statement) in block.statements.iter_mut().enumerate() {
match &mut statement.kind {
StatementKind::Assign(box (place, rvalue)) => {
match rvalue {
Rvalue::Cast(
CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
operand,
cast_ty,
) => {
let local = if let Some(local) = place.as_local() { local } else { return };
match operand {
Operand::Copy(place) | Operand::Move(place) => {
let operand_local =
if let Some(local) = place.local_or_deref_local() {
local
} else {
return;
};
if !interesting_locals.contains(operand_local) {
return;
}
let operand_ty = local_decls[operand_local].ty;
match (operand_ty.kind(), cast_ty.kind()) {
(ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
if of_ty_src == of_ty_dst {
// this is a cast from [T; N] into [T], so we are good
state.insert(local, statement_idx);
}
}
// current way of patching doesn't allow to work with `mut`
(
ty::Ref(
ty::RegionKind::ReErased,
operand_ty,
Mutability::Not,
),
ty::Ref(ty::RegionKind::ReErased, cast_ty, Mutability::Not),
) => {
match (operand_ty.kind(), cast_ty.kind()) {
// current way of patching doesn't allow to work with `mut`
(ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
if of_ty_src == of_ty_dst {
// this is a cast from [T; N] into [T], so we are good
state.insert(local, statement_idx);
}
}
_ => {}
}
}
_ => {}
}
}
_ => {}
}
}
Rvalue::Len(place) => {
let local = if let Some(local) = place.local_or_deref_local() {
local
} else {
return;
};
if let Some(cast_statement_idx) = state.get(&local).copied() {
patches_scratchpad.insert(cast_statement_idx, statement_idx);
}
}
_ => {
// invalidate
state.remove(&place.local);
}
}
}
_ => {}
}
}

let mut patcher = Patcher {
tcx,
patches_scratchpad: &*patches_scratchpad,
replacements_scratchpad,
local_decls,
statement_idx: 0,
};

block.expand_statements(|st| patcher.patch_expand_statement(st));
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
let mut _7: usize; // in scope 0 at $DIR/slice_len.rs:5:5: 5:33
let mut _8: bool; // in scope 0 at $DIR/slice_len.rs:5:5: 5:33
let mut _9: &[u32; 3]; // in scope 0 at $DIR/slice_len.rs:5:6: 5:19
let mut _10: &[u32; 3]; // in scope 0 at $DIR/slice_len.rs:5:6: 5:19

bb0: {
StorageLive(_1); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
Expand All @@ -27,14 +28,16 @@
// + literal: Const { ty: &[u32; 3], val: Unevaluated(Unevaluated { def: WithOptConstParam { did: DefId(0:3 ~ slice_len[6547]::main), const_param_did: None }, substs_: Some([]), promoted: Some(promoted[0]) }) }
_4 = _9; // scope 0 at $DIR/slice_len.rs:5:6: 5:19
_3 = _4; // scope 0 at $DIR/slice_len.rs:5:6: 5:19
StorageLive(_10); // scope 0 at $DIR/slice_len.rs:5:6: 5:19
_10 = _3; // scope 0 at $DIR/slice_len.rs:5:6: 5:19
_2 = move _3 as &[u32] (Pointer(Unsize)); // scope 0 at $DIR/slice_len.rs:5:6: 5:19
StorageDead(_3); // scope 0 at $DIR/slice_len.rs:5:18: 5:19
StorageLive(_6); // scope 0 at $DIR/slice_len.rs:5:31: 5:32
_6 = const 1_usize; // scope 0 at $DIR/slice_len.rs:5:31: 5:32
- _7 = Len((*_2)); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
_7 = const 3_usize; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
StorageDead(_10); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
- _8 = Lt(_6, _7); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
- assert(move _8, "index out of bounds: the length is {} but the index is {}", move _7, _6) -> bb1; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
+ _7 = const 3_usize; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
+ _8 = const true; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
+ assert(const true, "index out of bounds: the length is {} but the index is {}", const 3_usize, const 1_usize) -> bb1; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
let mut _7: usize; // in scope 0 at $DIR/slice_len.rs:5:5: 5:33
let mut _8: bool; // in scope 0 at $DIR/slice_len.rs:5:5: 5:33
let mut _9: &[u32; 3]; // in scope 0 at $DIR/slice_len.rs:5:6: 5:19
let mut _10: &[u32; 3]; // in scope 0 at $DIR/slice_len.rs:5:6: 5:19

bb0: {
StorageLive(_1); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
Expand All @@ -27,14 +28,16 @@
// + literal: Const { ty: &[u32; 3], val: Unevaluated(Unevaluated { def: WithOptConstParam { did: DefId(0:3 ~ slice_len[6547]::main), const_param_did: None }, substs_: Some([]), promoted: Some(promoted[0]) }) }
_4 = _9; // scope 0 at $DIR/slice_len.rs:5:6: 5:19
_3 = _4; // scope 0 at $DIR/slice_len.rs:5:6: 5:19
StorageLive(_10); // scope 0 at $DIR/slice_len.rs:5:6: 5:19
_10 = _3; // scope 0 at $DIR/slice_len.rs:5:6: 5:19
_2 = move _3 as &[u32] (Pointer(Unsize)); // scope 0 at $DIR/slice_len.rs:5:6: 5:19
StorageDead(_3); // scope 0 at $DIR/slice_len.rs:5:18: 5:19
StorageLive(_6); // scope 0 at $DIR/slice_len.rs:5:31: 5:32
_6 = const 1_usize; // scope 0 at $DIR/slice_len.rs:5:31: 5:32
- _7 = Len((*_2)); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
_7 = const 3_usize; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
StorageDead(_10); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
- _8 = Lt(_6, _7); // scope 0 at $DIR/slice_len.rs:5:5: 5:33
- assert(move _8, "index out of bounds: the length is {} but the index is {}", move _7, _6) -> bb1; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
+ _7 = const 3_usize; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
+ _8 = const true; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
+ assert(const true, "index out of bounds: the length is {} but the index is {}", const 3_usize, const 1_usize) -> bb1; // scope 0 at $DIR/slice_len.rs:5:5: 5:33
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
let mut _20: *const T; // in scope 0 at $DIR/issue_76432.rs:9:70: 9:84
let mut _21: *const T; // in scope 0 at $DIR/issue_76432.rs:9:70: 9:84
let mut _22: !; // in scope 0 at $SRC_DIR/core/src/panic.rs:LL:COL
let mut _23: &[T; 3]; // in scope 0 at $DIR/issue_76432.rs:7:19: 7:29
scope 1 {
debug v => _2; // in scope 1 at $DIR/issue_76432.rs:7:9: 7:10
let _13: &T; // in scope 1 at $DIR/issue_76432.rs:9:10: 9:16
Expand Down Expand Up @@ -51,16 +52,17 @@
StorageDead(_6); // scope 0 at $DIR/issue_76432.rs:7:28: 7:29
_4 = &_5; // scope 0 at $DIR/issue_76432.rs:7:19: 7:29
_3 = _4; // scope 0 at $DIR/issue_76432.rs:7:19: 7:29
StorageLive(_23); // scope 0 at $DIR/issue_76432.rs:7:19: 7:29
_23 = _3; // scope 0 at $DIR/issue_76432.rs:7:19: 7:29
_2 = move _3 as &[T] (Pointer(Unsize)); // scope 0 at $DIR/issue_76432.rs:7:19: 7:29
StorageDead(_3); // scope 0 at $DIR/issue_76432.rs:7:28: 7:29
StorageDead(_4); // scope 0 at $DIR/issue_76432.rs:7:29: 7:30
StorageLive(_9); // scope 1 at $DIR/issue_76432.rs:8:5: 11:6
_10 = Len((*_2)); // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
_10 = const 3_usize; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
StorageDead(_23); // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
_11 = const 3_usize; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
- _12 = Eq(move _10, const 3_usize); // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
- switchInt(move _12) -> [false: bb1, otherwise: bb2]; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
+ nop; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
+ switchInt(move _10) -> [3_usize: bb2, otherwise: bb1]; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
_12 = const true; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
goto -> bb2; // scope 1 at $DIR/issue_76432.rs:9:9: 9:33
}

bb1: {
Expand Down
Loading

0 comments on commit 680ff86

Please sign in to comment.