Skip to content

Commit a9aed86

Browse files
cjgillotcjgillot
committed
Reimplement NormalizeArrayLen.
1 parent c7bf469 commit a9aed86

17 files changed

+191
-435
lines changed

‎compiler/rustc_mir_transform/src/lib.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -549,13 +549,13 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
549549
&[
550550
&reveal_all::RevealAll, // has to be done before inlining, since inlined code is in RevealAll mode.
551551
&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
552-
&normalize_array_len::NormalizeArrayLen, // has to run after `slice::len` lowering
553552
&unreachable_prop::UnreachablePropagation,
554553
&uninhabited_enum_branching::UninhabitedEnumBranching,
555554
&o1(simplify::SimplifyCfg::new("after-uninhabited-enum-branching")),
556555
&inline::Inline,
557556
&remove_storage_markers::RemoveStorageMarkers,
558557
&remove_zsts::RemoveZsts,
558+
&normalize_array_len::NormalizeArrayLen, // has to run after `slice::len` lowering
559559
&const_goto::ConstGoto,
560560
&remove_unneeded_drops::RemoveUnneededDrops,
561561
&sroa::ScalarReplacementOfAggregates,

‎compiler/rustc_mir_transform/src/normalize_array_len.rs

+69-252
Original file line numberDiff line numberDiff line change
@@ -1,288 +1,105 @@
11
//! This pass eliminates casting of arrays into slices when their length
22
//! is taken using `.len()` method. Handy to preserve information in MIR for const prop
33
4+
use crate::ssa::SsaLocals;
45
use crate::MirPass;
5-
use rustc_data_structures::fx::FxIndexMap;
6-
use rustc_data_structures::intern::Interned;
7-
use rustc_index::bit_set::BitSet;
86
use rustc_index::vec::IndexVec;
7+
use rustc_middle::mir::visit::*;
98
use rustc_middle::mir::*;
10-
use rustc_middle::ty::{self, ReErased, Region, TyCtxt};
11-
12-
const MAX_NUM_BLOCKS: usize = 800;
13-
const MAX_NUM_LOCALS: usize = 3000;
9+
use rustc_middle::ty::{self, TyCtxt};
10+
use rustc_mir_dataflow::impls::borrowed_locals;
1411

1512
pub struct NormalizeArrayLen;
1613

1714
impl<'tcx> MirPass<'tcx> for NormalizeArrayLen {
1815
fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
1916
// See #105929
20-
sess.mir_opt_level() >= 4 && sess.opts.unstable_opts.unsound_mir_opts
17+
sess.mir_opt_level() >= 3
2118
}
2219

20+
#[instrument(level = "trace", skip(self, tcx, body))]
2321
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
24-
// early returns for edge cases of highly unrolled functions
25-
if body.basic_blocks.len() > MAX_NUM_BLOCKS {
26-
return;
27-
}
28-
if body.local_decls.len() > MAX_NUM_LOCALS {
29-
return;
30-
}
22+
debug!(def_id = ?body.source.def_id());
3123
normalize_array_len_calls(tcx, body)
3224
}
3325
}
3426

35-
pub fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
36-
// We don't ever touch terminators, so no need to invalidate the CFG cache
37-
let basic_blocks = body.basic_blocks.as_mut_preserves_cfg();
38-
let local_decls = &mut body.local_decls;
27+
fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
28+
let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
29+
let borrowed_locals = borrowed_locals(body);
30+
let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
3931

40-
// do a preliminary analysis to see if we ever have locals of type `[T;N]` or `&[T;N]`
41-
let mut interesting_locals = BitSet::new_empty(local_decls.len());
42-
for (local, decl) in local_decls.iter_enumerated() {
43-
match decl.ty.kind() {
44-
ty::Array(..) => {
45-
interesting_locals.insert(local);
46-
}
47-
ty::Ref(.., ty, Mutability::Not) => match ty.kind() {
48-
ty::Array(..) => {
49-
interesting_locals.insert(local);
50-
}
51-
_ => {}
52-
},
53-
_ => {}
54-
}
55-
}
56-
if interesting_locals.is_empty() {
57-
// we have found nothing to analyze
58-
return;
59-
}
60-
let num_intesting_locals = interesting_locals.count();
61-
let mut state = FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
62-
let mut patches_scratchpad =
63-
FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
64-
let mut replacements_scratchpad =
65-
FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default());
66-
for block in basic_blocks {
67-
// make length calls for arrays [T; N] not to decay into length calls for &[T]
68-
// that forbids constant propagation
69-
normalize_array_len_call(
70-
tcx,
71-
block,
72-
local_decls,
73-
&interesting_locals,
74-
&mut state,
75-
&mut patches_scratchpad,
76-
&mut replacements_scratchpad,
77-
);
78-
state.clear();
79-
patches_scratchpad.clear();
80-
replacements_scratchpad.clear();
81-
}
82-
}
32+
let slice_lengths = compute_slice_length(tcx, &ssa, body);
33+
debug!(?slice_lengths);
8334

84-
struct Patcher<'a, 'tcx> {
85-
tcx: TyCtxt<'tcx>,
86-
patches_scratchpad: &'a FxIndexMap<usize, usize>,
87-
replacements_scratchpad: &'a mut FxIndexMap<usize, Local>,
88-
local_decls: &'a mut IndexVec<Local, LocalDecl<'tcx>>,
89-
statement_idx: usize,
35+
Replacer { tcx, slice_lengths }.visit_body_preserves_cfg(body);
9036
}
9137

92-
impl<'tcx> Patcher<'_, 'tcx> {
93-
fn patch_expand_statement(
94-
&mut self,
95-
statement: &mut Statement<'tcx>,
96-
) -> Option<std::vec::IntoIter<Statement<'tcx>>> {
97-
let idx = self.statement_idx;
98-
if let Some(len_statemnt_idx) = self.patches_scratchpad.get(&idx).copied() {
99-
let mut statements = Vec::with_capacity(2);
100-
101-
// we are at statement that performs a cast. The only sound way is
102-
// to create another local that performs a similar copy without a cast and then
103-
// use this copy in the Len operation
104-
105-
match &statement.kind {
106-
StatementKind::Assign(box (
107-
..,
108-
Rvalue::Cast(
109-
CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
110-
operand,
111-
_,
112-
),
113-
)) => {
114-
match operand {
115-
Operand::Copy(place) | Operand::Move(place) => {
116-
// create new local
117-
let ty = operand.ty(self.local_decls, self.tcx);
118-
let local_decl = LocalDecl::with_source_info(ty, statement.source_info);
119-
let local = self.local_decls.push(local_decl);
120-
// make it live
121-
let mut make_live_statement = statement.clone();
122-
make_live_statement.kind = StatementKind::StorageLive(local);
123-
statements.push(make_live_statement);
124-
// copy into it
125-
126-
let operand = Operand::Copy(*place);
127-
let mut make_copy_statement = statement.clone();
128-
let assign_to = Place::from(local);
129-
let rvalue = Rvalue::Use(operand);
130-
make_copy_statement.kind =
131-
StatementKind::Assign(Box::new((assign_to, rvalue)));
132-
statements.push(make_copy_statement);
133-
134-
// to reorder we have to copy and make NOP
135-
statements.push(statement.clone());
136-
statement.make_nop();
137-
138-
self.replacements_scratchpad.insert(len_statemnt_idx, local);
139-
}
140-
_ => {
141-
unreachable!("it's a bug in the implementation")
142-
}
143-
}
144-
}
145-
_ => {
146-
unreachable!("it's a bug in the implementation")
38+
fn compute_slice_length<'tcx>(
39+
tcx: TyCtxt<'tcx>,
40+
ssa: &SsaLocals,
41+
body: &Body<'tcx>,
42+
) -> IndexVec<Local, Option<ty::Const<'tcx>>> {
43+
let mut slice_lengths = IndexVec::from_elem(None, &body.local_decls);
44+
45+
for (local, rvalue) in ssa.assignments(body) {
46+
match rvalue {
47+
Rvalue::Cast(
48+
CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
49+
operand,
50+
cast_ty,
51+
) => {
52+
let operand_ty = operand.ty(body, tcx);
53+
debug!(?operand_ty);
54+
if let Some(operand_ty) = operand_ty.builtin_deref(true)
55+
&& let ty::Array(_, len) = operand_ty.ty.kind()
56+
&& let Some(cast_ty) = cast_ty.builtin_deref(true)
57+
&& let ty::Slice(..) = cast_ty.ty.kind()
58+
{
59+
slice_lengths[local] = Some(*len);
14760
}
14861
}
149-
150-
self.statement_idx += 1;
151-
152-
Some(statements.into_iter())
153-
} else if let Some(local) = self.replacements_scratchpad.get(&idx).copied() {
154-
let mut statements = Vec::with_capacity(2);
155-
156-
match &statement.kind {
157-
StatementKind::Assign(box (into, Rvalue::Len(place))) => {
158-
let add_deref = if let Some(..) = place.as_local() {
159-
false
160-
} else if let Some(..) = place.local_or_deref_local() {
161-
true
162-
} else {
163-
unreachable!("it's a bug in the implementation")
164-
};
165-
// replace len statement
166-
let mut len_statement = statement.clone();
167-
let mut place = Place::from(local);
168-
if add_deref {
169-
place = self.tcx.mk_place_deref(place);
170-
}
171-
len_statement.kind =
172-
StatementKind::Assign(Box::new((*into, Rvalue::Len(place))));
173-
statements.push(len_statement);
174-
175-
// make temporary dead
176-
let mut make_dead_statement = statement.clone();
177-
make_dead_statement.kind = StatementKind::StorageDead(local);
178-
statements.push(make_dead_statement);
179-
180-
// make original statement NOP
181-
statement.make_nop();
62+
// The length information is stored in the fat pointer, so we treat `operand` as a value.
63+
Rvalue::Use(operand) => {
64+
if let Some(rhs) = operand.place() && let Some(rhs) = rhs.as_local() {
65+
slice_lengths[local] = slice_lengths[rhs];
18266
}
183-
_ => {
184-
unreachable!("it's a bug in the implementation")
67+
}
68+
// The length information is stored in the fat pointer.
69+
// Reborrowing copies length information from one pointer to the other.
70+
Rvalue::Ref(_, _, rhs) | Rvalue::AddressOf(_, rhs) => {
71+
if let [PlaceElem::Deref] = rhs.projection[..] {
72+
slice_lengths[local] = slice_lengths[rhs.local];
18573
}
18674
}
187-
188-
self.statement_idx += 1;
189-
190-
Some(statements.into_iter())
191-
} else {
192-
self.statement_idx += 1;
193-
None
75+
_ => {}
19476
}
19577
}
78+
79+
slice_lengths
19680
}
19781

198-
fn normalize_array_len_call<'tcx>(
82+
struct Replacer<'tcx> {
19983
tcx: TyCtxt<'tcx>,
200-
block: &mut BasicBlockData<'tcx>,
201-
local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
202-
interesting_locals: &BitSet<Local>,
203-
state: &mut FxIndexMap<Local, usize>,
204-
patches_scratchpad: &mut FxIndexMap<usize, usize>,
205-
replacements_scratchpad: &mut FxIndexMap<usize, Local>,
206-
) {
207-
for (statement_idx, statement) in block.statements.iter_mut().enumerate() {
208-
match &mut statement.kind {
209-
StatementKind::Assign(box (place, rvalue)) => {
210-
match rvalue {
211-
Rvalue::Cast(
212-
CastKind::Pointer(ty::adjustment::PointerCast::Unsize),
213-
operand,
214-
cast_ty,
215-
) => {
216-
let Some(local) = place.as_local() else { return };
217-
match operand {
218-
Operand::Copy(place) | Operand::Move(place) => {
219-
let Some(operand_local) = place.local_or_deref_local() else { return; };
220-
if !interesting_locals.contains(operand_local) {
221-
return;
222-
}
223-
let operand_ty = local_decls[operand_local].ty;
224-
match (operand_ty.kind(), cast_ty.kind()) {
225-
(ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
226-
if of_ty_src == of_ty_dst {
227-
// this is a cast from [T; N] into [T], so we are good
228-
state.insert(local, statement_idx);
229-
}
230-
}
231-
// current way of patching doesn't allow to work with `mut`
232-
(
233-
ty::Ref(
234-
Region(Interned(ReErased, _)),
235-
operand_ty,
236-
Mutability::Not,
237-
),
238-
ty::Ref(
239-
Region(Interned(ReErased, _)),
240-
cast_ty,
241-
Mutability::Not,
242-
),
243-
) => {
244-
match (operand_ty.kind(), cast_ty.kind()) {
245-
// current way of patching doesn't allow to work with `mut`
246-
(ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => {
247-
if of_ty_src == of_ty_dst {
248-
// this is a cast from [T; N] into [T], so we are good
249-
state.insert(local, statement_idx);
250-
}
251-
}
252-
_ => {}
253-
}
254-
}
255-
_ => {}
256-
}
257-
}
258-
_ => {}
259-
}
260-
}
261-
Rvalue::Len(place) => {
262-
let Some(local) = place.local_or_deref_local() else {
263-
return;
264-
};
265-
if let Some(cast_statement_idx) = state.get(&local).copied() {
266-
patches_scratchpad.insert(cast_statement_idx, statement_idx);
267-
}
268-
}
269-
_ => {
270-
// invalidate
271-
state.remove(&place.local);
272-
}
273-
}
274-
}
275-
_ => {}
276-
}
277-
}
84+
slice_lengths: IndexVec<Local, Option<ty::Const<'tcx>>>,
85+
}
27886

279-
let mut patcher = Patcher {
280-
tcx,
281-
patches_scratchpad: &*patches_scratchpad,
282-
replacements_scratchpad,
283-
local_decls,
284-
statement_idx: 0,
285-
};
87+
impl<'tcx> MutVisitor<'tcx> for Replacer<'tcx> {
88+
fn tcx(&self) -> TyCtxt<'tcx> {
89+
self.tcx
90+
}
28691

287-
block.expand_statements(|st| patcher.patch_expand_statement(st));
92+
fn visit_rvalue(&mut self, rvalue: &mut Rvalue<'tcx>, loc: Location) {
93+
if let Rvalue::Len(place) = rvalue
94+
&& let [PlaceElem::Deref] = &place.projection[..]
95+
&& let Some(len) = self.slice_lengths[place.local]
96+
{
97+
*rvalue = Rvalue::Use(Operand::Constant(Box::new(Constant {
98+
span: rustc_span::DUMMY_SP,
99+
user_ty: None,
100+
literal: ConstantKind::from_const(len, self.tcx),
101+
})));
102+
}
103+
self.super_rvalue(rvalue, loc);
104+
}
288105
}

‎tests/mir-opt/const_prop/bad_op_unsafe_oob_for_slices.main.ConstProp.32bit.diff

+3-3
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,11 @@
3333
StorageLive(_4); // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:13: +3:15
3434
StorageLive(_5); // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:23: +3:24
3535
_5 = const 3_usize; // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:23: +3:24
36-
_6 = Len((*_1)); // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
36+
_6 = const 3_usize; // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
3737
- _7 = Lt(_5, _6); // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
3838
- assert(move _7, "index out of bounds: the length is {} but the index is {}", move _6, _5) -> bb1; // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
39-
+ _7 = Lt(const 3_usize, _6); // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
40-
+ assert(move _7, "index out of bounds: the length is {} but the index is {}", move _6, const 3_usize) -> bb1; // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
39+
+ _7 = const false; // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
40+
+ assert(const false, "index out of bounds: the length is {} but the index is {}", const 3_usize, const 3_usize) -> bb1; // scope 2 at $DIR/bad_op_unsafe_oob_for_slices.rs:+3:18: +3:25
4141
}
4242

4343
bb1: {

0 commit comments

Comments
 (0)