Skip to content

Commit 5f5b207

Browse files
scottmcmscottmcm
committed
Redo SliceIndex implementations
1 parent 4314bd7 commit 5f5b207

8 files changed

+375
-51
lines changed

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

+90-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,51 @@ 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 compile-time.
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+
unsafe {
135+
let ptr = get_noubcheck(ptr, offset);
136+
crate::intrinsics::aggregate_raw_ptr(ptr, len)
137+
}
138+
}
139+
140+
#[inline(always)]
141+
const unsafe fn get_offset_len_mut_noubcheck<T>(
142+
ptr: *mut [T],
143+
offset: usize,
144+
len: usize,
145+
) -> *mut [T] {
146+
// SAFETY: The caller already checked these preconditions
147+
unsafe {
148+
let ptr = get_mut_noubcheck(ptr, offset);
149+
crate::intrinsics::aggregate_raw_ptr(ptr, len)
150+
}
151+
}
152+
109153
mod private_slice_index {
110154
use super::ops;
111155
#[stable(feature = "slice_get_slice", since = "1.28.0")]
@@ -203,13 +247,17 @@ unsafe impl<T> SliceIndex<[T]> for usize {
203247
#[inline]
204248
fn get(self, slice: &[T]) -> Option<&T> {
205249
// SAFETY: `self` is checked to be in bounds.
206-
if self < slice.len() { unsafe { Some(&*self.get_unchecked(slice)) } } else { None }
250+
if self < slice.len() { unsafe { Some(&*get_noubcheck(slice, self)) } } else { None }
207251
}
208252

209253
#[inline]
210254
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 }
255+
if self < slice.len() {
256+
// SAFETY: `self` is checked to be in bounds.
257+
unsafe { Some(&mut *get_mut_noubcheck(slice, self)) }
258+
} else {
259+
None
260+
}
213261
}
214262

215263
#[inline]
@@ -227,7 +275,7 @@ unsafe impl<T> SliceIndex<[T]> for usize {
227275
// Use intrinsics::assume instead of hint::assert_unchecked so that we don't check the
228276
// precondition of this function twice.
229277
crate::intrinsics::assume(self < slice.len());
230-
slice.as_ptr().add(self)
278+
get_noubcheck(slice, self)
231279
}
232280
}
233281

@@ -239,7 +287,7 @@ unsafe impl<T> SliceIndex<[T]> for usize {
239287
(this: usize = self, len: usize = slice.len()) => this < len
240288
);
241289
// SAFETY: see comments for `get_unchecked` above.
242-
unsafe { slice.as_mut_ptr().add(self) }
290+
unsafe { get_mut_noubcheck(slice, self) }
243291
}
244292

245293
#[inline]
@@ -265,7 +313,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
265313
fn get(self, slice: &[T]) -> Option<&[T]> {
266314
if self.end() <= slice.len() {
267315
// SAFETY: `self` is checked to be valid and in bounds above.
268-
unsafe { Some(&*self.get_unchecked(slice)) }
316+
unsafe { Some(&*get_offset_len_noubcheck(slice, self.start(), self.len())) }
269317
} else {
270318
None
271319
}
@@ -275,7 +323,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
275323
fn get_mut(self, slice: &mut [T]) -> Option<&mut [T]> {
276324
if self.end() <= slice.len() {
277325
// SAFETY: `self` is checked to be valid and in bounds above.
278-
unsafe { Some(&mut *self.get_unchecked_mut(slice)) }
326+
unsafe { Some(&mut *get_offset_len_mut_noubcheck(slice, self.start(), self.len())) }
279327
} else {
280328
None
281329
}
@@ -292,7 +340,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
292340
// cannot be longer than `isize::MAX`. They also guarantee that
293341
// `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
294342
// so the call to `add` is safe.
295-
unsafe { ptr::slice_from_raw_parts(slice.as_ptr().add(self.start()), self.len()) }
343+
unsafe { get_offset_len_noubcheck(slice, self.start(), self.len()) }
296344
}
297345

298346
#[inline]
@@ -304,14 +352,14 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
304352
);
305353

306354
// SAFETY: see comments for `get_unchecked` above.
307-
unsafe { ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start()), self.len()) }
355+
unsafe { get_offset_len_mut_noubcheck(slice, self.start(), self.len()) }
308356
}
309357

310358
#[inline]
311359
fn index(self, slice: &[T]) -> &[T] {
312360
if self.end() <= slice.len() {
313361
// SAFETY: `self` is checked to be valid and in bounds above.
314-
unsafe { &*self.get_unchecked(slice) }
362+
unsafe { &*get_offset_len_noubcheck(slice, self.start(), self.len()) }
315363
} else {
316364
slice_end_index_len_fail(self.end(), slice.len())
317365
}
@@ -321,7 +369,7 @@ unsafe impl<T> SliceIndex<[T]> for ops::IndexRange {
321369
fn index_mut(self, slice: &mut [T]) -> &mut [T] {
322370
if self.end() <= slice.len() {
323371
// SAFETY: `self` is checked to be valid and in bounds above.
324-
unsafe { &mut *self.get_unchecked_mut(slice) }
372+
unsafe { &mut *get_offset_len_mut_noubcheck(slice, self.start(), self.len()) }
325373
} else {
326374
slice_end_index_len_fail(self.end(), slice.len())
327375
}
@@ -338,21 +386,26 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
338386

339387
#[inline]
340388
fn get(self, slice: &[T]) -> Option<&[T]> {
341-
if self.start > self.end || self.end > slice.len() {
342-
None
343-
} else {
389+
// Using checked_sub is a safe way to get `SubUnchecked` in MIR
390+
if let Some(new_len) = usize::checked_sub(self.end, self.start)
391+
&& self.end <= slice.len()
392+
{
344393
// SAFETY: `self` is checked to be valid and in bounds above.
345-
unsafe { Some(&*self.get_unchecked(slice)) }
394+
unsafe { Some(&*get_offset_len_noubcheck(slice, self.start, new_len)) }
395+
} else {
396+
None
346397
}
347398
}
348399

349400
#[inline]
350401
fn get_mut(self, slice: &mut [T]) -> Option<&mut [T]> {
351-
if self.start > self.end || self.end > slice.len() {
352-
None
353-
} else {
402+
if let Some(new_len) = usize::checked_sub(self.end, self.start)
403+
&& self.end <= slice.len()
404+
{
354405
// SAFETY: `self` is checked to be valid and in bounds above.
355-
unsafe { Some(&mut *self.get_unchecked_mut(slice)) }
406+
unsafe { Some(&mut *get_offset_len_mut_noubcheck(slice, self.start, new_len)) }
407+
} else {
408+
None
356409
}
357410
}
358411

@@ -373,8 +426,10 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
373426
// `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
374427
// so the call to `add` is safe and the length calculation cannot overflow.
375428
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)
429+
// Using the intrinsic avoids a superfluous UB check,
430+
// since the one on this method already checked `end >= start`.
431+
let new_len = crate::intrinsics::unchecked_sub(self.end, self.start);
432+
get_offset_len_noubcheck(slice, self.start, new_len)
378433
}
379434
}
380435

@@ -391,31 +446,34 @@ unsafe impl<T> SliceIndex<[T]> for ops::Range<usize> {
391446
);
392447
// SAFETY: see comments for `get_unchecked` above.
393448
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)
449+
let new_len = crate::intrinsics::unchecked_sub(self.end, self.start);
450+
get_offset_len_mut_noubcheck(slice, self.start, new_len)
396451
}
397452
}
398453

399454
#[inline(always)]
400455
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() {
456+
// Using checked_sub is a safe way to get `SubUnchecked` in MIR
457+
let Some(new_len) = usize::checked_sub(self.end, self.start) else {
458+
slice_index_order_fail(self.start, self.end)
459+
};
460+
if self.end > slice.len() {
404461
slice_end_index_len_fail(self.end, slice.len());
405462
}
406463
// SAFETY: `self` is checked to be valid and in bounds above.
407-
unsafe { &*self.get_unchecked(slice) }
464+
unsafe { &*get_offset_len_noubcheck(slice, self.start, new_len) }
408465
}
409466

410467
#[inline]
411468
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() {
469+
let Some(new_len) = usize::checked_sub(self.end, self.start) else {
470+
slice_index_order_fail(self.start, self.end)
471+
};
472+
if self.end > slice.len() {
415473
slice_end_index_len_fail(self.end, slice.len());
416474
}
417475
// SAFETY: `self` is checked to be valid and in bounds above.
418-
unsafe { &mut *self.get_unchecked_mut(slice) }
476+
unsafe { &mut *get_offset_len_mut_noubcheck(slice, self.start, new_len) }
419477
}
420478
}
421479

‎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)