Skip to content

Commit 33c4817

Browse files
scottmcmscottmcm
committed
Redo SliceIndex implementations
1 parent 17a9d34 commit 33c4817

8 files changed

+371
-51
lines changed

‎library/core/src/slice/index.rs

+86-32
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
33
use crate::intrinsics::const_eval_select;
44
use crate::ops;
5-
use crate::ptr;
65
use crate::ub_checks::assert_unsafe_precondition;
76

87
#[stable(feature = "rust1", since = "1.0.0")]
@@ -106,6 +105,47 @@ const fn slice_end_index_overflow_fail() -> ! {
106105
panic!("attempted to index slice up to maximum usize");
107106
}
108107

108+
// The UbChecks are great for catching bugs in the unsafe methods, but including
109+
// them in safe indexing is unnecessary and hurts inlining and debug runtime perf.
110+
// Both the safe and unsafe public methods share these helpers,
111+
// which use intrinsics directly to get *no* extra checks.
112+
113+
#[inline(always)]
114+
const unsafe fn get_noubcheck<T>(ptr: *const [T], index: usize) -> *const T {
115+
let ptr = ptr as *const T;
116+
// SAFETY: The caller already checked these preconditions
117+
unsafe { crate::intrinsics::offset(ptr, index) }
118+
}
119+
120+
#[inline(always)]
121+
const unsafe fn get_mut_noubcheck<T>(ptr: *mut [T], index: usize) -> *mut T {
122+
let ptr = ptr as *mut T;
123+
// SAFETY: The caller already checked these preconditions
124+
unsafe { crate::intrinsics::offset(ptr, index) }
125+
}
126+
127+
#[inline(always)]
128+
const unsafe fn get_offset_len_noubcheck<T>(
129+
ptr: *const [T],
130+
offset: usize,
131+
len: usize,
132+
) -> *const [T] {
133+
// SAFETY: The caller already checked these preconditions
134+
let ptr = unsafe { get_noubcheck(ptr, offset) };
135+
crate::intrinsics::aggregate_raw_ptr(ptr, len)
136+
}
137+
138+
#[inline(always)]
139+
const unsafe fn get_offset_len_mut_noubcheck<T>(
140+
ptr: *mut [T],
141+
offset: usize,
142+
len: usize,
143+
) -> *mut [T] {
144+
// SAFETY: The caller already checked these preconditions
145+
let ptr = unsafe { get_mut_noubcheck(ptr, offset) };
146+
crate::intrinsics::aggregate_raw_ptr(ptr, len)
147+
}
148+
109149
mod private_slice_index {
110150
use super::ops;
111151
#[stable(feature = "slice_get_slice", since = "1.28.0")]
@@ -203,13 +243,17 @@ unsafe impl<T> SliceIndex<[T]> for usize {
203243
#[inline]
204244
fn get(self, slice: &[T]) -> Option<&T> {
205245
// SAFETY: `self` is checked to be in bounds.
206-
if self < slice.len() { unsafe { Some(&*self.get_unchecked(slice)) } } else { None }
246+
if self < slice.len() { unsafe { Some(&*get_noubcheck(slice, self)) } } else { None }
207247
}
208248

209249
#[inline]
210250
fn get_mut(self, slice: &mut [T]) -> Option<&mut T> {
211-
// SAFETY: `self` is checked to be in bounds.
212-
if self < slice.len() { unsafe { Some(&mut *self.get_unchecked_mut(slice)) } } else { None }
251+
if self < slice.len() {
252+
// SAFETY: `self` is checked to be in bounds.
253+
unsafe { Some(&mut *get_mut_noubcheck(slice, self)) }
254+
} else {
255+
None
256+
}
213257
}
214258

215259
#[inline]
@@ -227,7 +271,7 @@ unsafe impl<T> SliceIndex<[T]> for usize {
227271
// Use intrinsics::assume instead of hint::assert_unchecked so that we don't check the
228272
// precondition of this function twice.
229273
crate::intrinsics::assume(self < slice.len());
230-
slice.as_ptr().add(self)
274+
get_noubcheck(slice, self)
231275
}
232276
}
233277

@@ -239,7 +283,7 @@ unsafe impl<T> SliceIndex<[T]> for usize {
239283
(this: usize = self, len: usize = slice.len()) => this < len
240284
);
241285
// SAFETY: see comments for `get_unchecked` above.
242-
unsafe { slice.as_mut_ptr().add(self) }
286+
unsafe { get_mut_noubcheck(slice, self) }
243287
}
244288

245289
#[inline]
@@ -265,7 +309,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
265309
fn get(self, slice: &[T]) -> Option<&[T]> {
266310
if self.end() <= slice.len() {
267311
// SAFETY: `self` is checked to be valid and in bounds above.
268-
unsafe { Some(&*self.get_unchecked(slice)) }
312+
unsafe { Some(&*get_offset_len_noubcheck(slice, self.start(), self.len())) }
269313
} else {
270314
None
271315
}
@@ -275,7 +319,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
275319
fn get_mut(self, slice: &mut [T]) -> Option<&mut [T]> {
276320
if self.end() <= slice.len() {
277321
// SAFETY: `self` is checked to be valid and in bounds above.
278-
unsafe { Some(&mut *self.get_unchecked_mut(slice)) }
322+
unsafe { Some(&mut *get_offset_len_mut_noubcheck(slice, self.start(), self.len())) }
279323
} else {
280324
None
281325
}
@@ -292,7 +336,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
292336
// cannot be longer than `isize::MAX`. They also guarantee that
293337
// `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
294338
// so the call to `add` is safe.
295-
unsafe { ptr::slice_from_raw_parts(slice.as_ptr().add(self.start()), self.len()) }
339+
unsafe { get_offset_len_noubcheck(slice, self.start(), self.len()) }
296340
}
297341

298342
#[inline]
@@ -304,14 +348,14 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
304348
);
305349

306350
// SAFETY: see comments for `get_unchecked` above.
307-
unsafe { ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start()), self.len()) }
351+
unsafe { get_offset_len_mut_noubcheck(slice, self.start(), self.len()) }
308352
}
309353

310354
#[inline]
311355
fn index(self, slice: &[T]) -> &[T] {
312356
if self.end() <= slice.len() {
313357
// SAFETY: `self` is checked to be valid and in bounds above.
314-
unsafe { &*self.get_unchecked(slice) }
358+
unsafe { &*get_offset_len_noubcheck(slice, self.start(), self.len()) }
315359
} else {
316360
slice_end_index_len_fail(self.end(), slice.len())
317361
}
@@ -321,7 +365,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
321365
fn index_mut(self, slice: &mut [T]) -> &mut [T] {
322366
if self.end() <= slice.len() {
323367
// SAFETY: `self` is checked to be valid and in bounds above.
324-
unsafe { &mut *self.get_unchecked_mut(slice) }
368+
unsafe { &mut *get_offset_len_mut_noubcheck(slice, self.start(), self.len()) }
325369
} else {
326370
slice_end_index_len_fail(self.end(), slice.len())
327371
}
@@ -338,21 +382,26 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
338382

339383
#[inline]
340384
fn get(self, slice: &[T]) -> Option<&[T]> {
341-
if self.start > self.end || self.end > slice.len() {
342-
None
343-
} else {
385+
// Using checked_sub is a safe way to get `SubUnchecked` in MIR
386+
if let Some(new_len) = usize::checked_sub(self.end, self.start)
387+
&& self.end <= slice.len()
388+
{
344389
// SAFETY: `self` is checked to be valid and in bounds above.
345-
unsafe { Some(&*self.get_unchecked(slice)) }
390+
unsafe { Some(&*get_offset_len_noubcheck(slice, self.start, new_len)) }
391+
} else {
392+
None
346393
}
347394
}
348395

349396
#[inline]
350397
fn get_mut(self, slice: &mut [T]) -> Option<&mut [T]> {
351-
if self.start > self.end || self.end > slice.len() {
352-
None
353-
} else {
398+
if let Some(new_len) = usize::checked_sub(self.end, self.start)
399+
&& self.end <= slice.len()
400+
{
354401
// SAFETY: `self` is checked to be valid and in bounds above.
355-
unsafe { Some(&mut *self.get_unchecked_mut(slice)) }
402+
unsafe { Some(&mut *get_offset_len_mut_noubcheck(slice, self.start, new_len)) }
403+
} else {
404+
None
356405
}
357406
}
358407

@@ -373,8 +422,10 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
373422
// `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
374423
// so the call to `add` is safe and the length calculation cannot overflow.
375424
unsafe {
376-
let new_len = self.end.unchecked_sub(self.start);
377-
ptr::slice_from_raw_parts(slice.as_ptr().add(self.start), new_len)
425+
// Using the intrinsic avoids a superfluous UB check,
426+
// since the one on this method already checked `end >= start`.
427+
let new_len = crate::intrinsics::unchecked_sub(self.end, self.start);
428+
get_offset_len_noubcheck(slice, self.start, new_len)
378429
}
379430
}
380431

@@ -391,31 +442,34 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
391442
);
392443
// SAFETY: see comments for `get_unchecked` above.
393444
unsafe {
394-
let new_len = self.end.unchecked_sub(self.start);
395-
ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start), new_len)
445+
let new_len = crate::intrinsics::unchecked_sub(self.end, self.start);
446+
get_offset_len_mut_noubcheck(slice, self.start, new_len)
396447
}
397448
}
398449

399450
#[inline(always)]
400451
fn index(self, slice: &[T]) -> &[T] {
401-
if self.start > self.end {
402-
slice_index_order_fail(self.start, self.end);
403-
} else if self.end > slice.len() {
452+
// Using checked_sub is a safe way to get `SubUnchecked` in MIR
453+
let Some(new_len) = usize::checked_sub(self.end, self.start) else {
454+
slice_index_order_fail(self.start, self.end)
455+
};
456+
if self.end > slice.len() {
404457
slice_end_index_len_fail(self.end, slice.len());
405458
}
406459
// SAFETY: `self` is checked to be valid and in bounds above.
407-
unsafe { &*self.get_unchecked(slice) }
460+
unsafe { &*get_offset_len_noubcheck(slice, self.start, new_len) }
408461
}
409462

410463
#[inline]
411464
fn index_mut(self, slice: &mut [T]) -> &mut [T] {
412-
if self.start > self.end {
413-
slice_index_order_fail(self.start, self.end);
414-
} else if self.end > slice.len() {
465+
let Some(new_len) = usize::checked_sub(self.end, self.start) else {
466+
slice_index_order_fail(self.start, self.end)
467+
};
468+
if self.end > slice.len() {
415469
slice_end_index_len_fail(self.end, slice.len());
416470
}
417471
// SAFETY: `self` is checked to be valid and in bounds above.
418-
unsafe { &mut *self.get_unchecked_mut(slice) }
472+
unsafe { &mut *get_offset_len_mut_noubcheck(slice, self.start, new_len) }
419473
}
420474
}
421475

‎tests/mir-opt/pre-codegen/slice_index.rs

+25-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
// skip-filecheck
21
//@ compile-flags: -O -C debuginfo=0 -Zmir-opt-level=2 -Z ub-checks=yes
32
// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
43

@@ -9,21 +8,39 @@ use std::ops::Range;
98

109
// EMIT_MIR slice_index.slice_index_usize.PreCodegen.after.mir
1110
pub fn slice_index_usize(slice: &[u32], index: usize) -> u32 {
11+
// CHECK-LABEL: slice_index_usize
12+
// CHECK: [[LEN:_[0-9]+]] = Len((*_1))
13+
// CHECK: Lt(_2, [[LEN]])
14+
// CHECK-NOT: precondition_check
15+
// CHECK: _0 = (*_1)[_2];
1216
slice[index]
1317
}
1418

1519
// EMIT_MIR slice_index.slice_get_mut_usize.PreCodegen.after.mir
1620
pub fn slice_get_mut_usize(slice: &mut [u32], index: usize) -> Option<&mut u32> {
21+
// CHECK-LABEL: slice_get_mut_usize
22+
// CHECK: [[LEN:_[0-9]+]] = Len((*_1))
23+
// CHECK: Lt(_2, move [[LEN]])
24+
// CHECK-NOT: precondition_check
1725
slice.get_mut(index)
1826
}
1927

2028
// EMIT_MIR slice_index.slice_index_range.PreCodegen.after.mir
2129
pub fn slice_index_range(slice: &[u32], index: Range<usize>) -> &[u32] {
30+
// CHECK-LABEL: slice_index_range
2231
&slice[index]
2332
}
2433

2534
// EMIT_MIR slice_index.slice_get_unchecked_mut_range.PreCodegen.after.mir
2635
pub unsafe fn slice_get_unchecked_mut_range(slice: &mut [u32], index: Range<usize>) -> &mut [u32] {
36+
// CHECK-LABEL: slice_get_unchecked_mut_range
37+
// CHECK: [[START:_[0-9]+]] = move (_2.0: usize);
38+
// CHECK: [[END:_[0-9]+]] = move (_2.1: usize);
39+
// CHECK: precondition_check
40+
// CHECK: [[LEN:_[0-9]+]] = SubUnchecked([[END]], [[START]]);
41+
// CHECK: [[PTR:_[0-9]+]] = Offset({{_[0-9]+}}, [[START]]);
42+
// CHECK: [[SLICE:_[0-9]+]] = *mut [u32] from ([[PTR]], [[LEN]])
43+
// CHECK: _0 = &mut (*[[SLICE]]);
2744
slice.get_unchecked_mut(index)
2845
}
2946

@@ -32,5 +49,12 @@ pub unsafe fn slice_ptr_get_unchecked_range(
3249
slice: *const [u32],
3350
index: Range<usize>,
3451
) -> *const [u32] {
52+
// CHECK-LABEL: slice_ptr_get_unchecked_range
53+
// CHECK: [[START:_[0-9]+]] = move (_2.0: usize);
54+
// CHECK: [[END:_[0-9]+]] = move (_2.1: usize);
55+
// CHECK: precondition_check
56+
// CHECK: [[LEN:_[0-9]+]] = SubUnchecked([[END]], [[START]]);
57+
// CHECK: [[PTR:_[0-9]+]] = Offset({{_[0-9]+}}, [[START]]);
58+
// CHECK: _0 = *const [u32] from ([[PTR]], [[LEN]])
3559
slice.get_unchecked(index)
3660
}

‎tests/mir-opt/pre-codegen/slice_index.slice_get_mut_usize.PreCodegen.after.panic-abort.mir

+42-1
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,54 @@ fn slice_get_mut_usize(_1: &mut [u32], _2: usize) -> Option<&mut u32> {
55
debug index => _2;
66
let mut _0: std::option::Option<&mut u32>;
77
scope 1 (inlined core::slice::<impl [u32]>::get_mut::<usize>) {
8+
scope 2 (inlined <usize as SliceIndex<[u32]>>::get_mut) {
9+
let mut _3: usize;
10+
let mut _4: bool;
11+
let mut _5: *mut [u32];
12+
let mut _7: *mut u32;
13+
let mut _8: &mut u32;
14+
scope 3 (inlined core::slice::index::get_mut_noubcheck::<u32>) {
15+
let _6: *mut u32;
16+
scope 4 {
17+
}
18+
}
19+
}
820
}
921

1022
bb0: {
11-
_0 = <usize as SliceIndex<[u32]>>::get_mut(move _2, move _1) -> [return: bb1, unwind unreachable];
23+
StorageLive(_7);
24+
StorageLive(_4);
25+
StorageLive(_3);
26+
_3 = Len((*_1));
27+
_4 = Lt(_2, move _3);
28+
switchInt(move _4) -> [0: bb1, otherwise: bb2];
1229
}
1330

1431
bb1: {
32+
StorageDead(_3);
33+
_0 = const Option::<&mut u32>::None;
34+
goto -> bb3;
35+
}
36+
37+
bb2: {
38+
StorageDead(_3);
39+
StorageLive(_8);
40+
StorageLive(_5);
41+
_5 = &raw mut (*_1);
42+
StorageLive(_6);
43+
_6 = _5 as *mut u32 (PtrToPtr);
44+
_7 = Offset(_6, _2);
45+
StorageDead(_6);
46+
StorageDead(_5);
47+
_8 = &mut (*_7);
48+
_0 = Option::<&mut u32>::Some(move _8);
49+
StorageDead(_8);
50+
goto -> bb3;
51+
}
52+
53+
bb3: {
54+
StorageDead(_4);
55+
StorageDead(_7);
1556
return;
1657
}
1758
}

0 commit comments

Comments
 (0)