From 3554036280525cec34103a8f66049b0881b14d27 Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Sat, 4 Mar 2023 19:10:36 -0800
Subject: [PATCH] Use `nuw` when calculating slice lengths from `Range`s

An `assume` would definitely not be worth it, but since the flag is almost free we might as well tell LLVM this, especially on `_unchecked` calls where there's no obvious way for it to deduce it.

(Today neither safe nor unsafe indexing gets it: <https://rust.godbolt.org/z/G1jYT548s>)
---
 library/core/src/slice/index.rs               |  9 +++--
 tests/codegen/slice-indexing.rs               | 35 +++++++++++++++++++
 .../const-eval/ub-slice-get-unchecked.rs      |  9 +++++
 .../const-eval/ub-slice-get-unchecked.stderr  | 18 ++++++++++
 4 files changed, 68 insertions(+), 3 deletions(-)
 create mode 100644 tests/codegen/slice-indexing.rs
 create mode 100644 tests/ui/consts/const-eval/ub-slice-get-unchecked.rs
 create mode 100644 tests/ui/consts/const-eval/ub-slice-get-unchecked.stderr

diff --git a/library/core/src/slice/index.rs b/library/core/src/slice/index.rs
index f0e5ea53d7d..3539353240a 100644
--- a/library/core/src/slice/index.rs
+++ b/library/core/src/slice/index.rs
@@ -2,6 +2,7 @@
 
 use crate::intrinsics::assert_unsafe_precondition;
 use crate::intrinsics::const_eval_select;
+use crate::intrinsics::unchecked_sub;
 use crate::ops;
 use crate::ptr;
 
@@ -375,14 +376,15 @@ unsafe fn get_unchecked(self, slice: *const [T]) -> *const [T] {
         // SAFETY: the caller guarantees that `slice` is not dangling, so it
         // cannot be longer than `isize::MAX`. They also guarantee that
         // `self` is in bounds of `slice` so `self` cannot overflow an `isize`,
-        // so the call to `add` is safe.
+        // so the call to `add` is safe and the length calculation cannot overflow.
         unsafe {
             assert_unsafe_precondition!(
                 "slice::get_unchecked requires that the range is within the slice",
                 [T](this: ops::Range<usize>, slice: *const [T]) =>
                 this.end >= this.start && this.end <= slice.len()
             );
-            ptr::slice_from_raw_parts(slice.as_ptr().add(self.start), self.end - self.start)
+            let new_len = unchecked_sub(self.end, self.start);
+            ptr::slice_from_raw_parts(slice.as_ptr().add(self.start), new_len)
         }
     }
 
@@ -396,7 +398,8 @@ unsafe fn get_unchecked_mut(self, slice: *mut [T]) -> *mut [T] {
                 [T](this: ops::Range<usize>, slice: *mut [T]) =>
                 this.end >= this.start && this.end <= slice.len()
             );
-            ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start), self.end - self.start)
+            let new_len = unchecked_sub(self.end, self.start);
+            ptr::slice_from_raw_parts_mut(slice.as_mut_ptr().add(self.start), new_len)
         }
     }
 
diff --git a/tests/codegen/slice-indexing.rs b/tests/codegen/slice-indexing.rs
new file mode 100644
index 00000000000..c40d59fb0cf
--- /dev/null
+++ b/tests/codegen/slice-indexing.rs
@@ -0,0 +1,35 @@
+// compile-flags: -O
+// only-64bit (because the LLVM type of i64 for usize shows up)
+// ignore-debug: the debug assertions get in the way
+
+#![crate_type = "lib"]
+
+use std::ops::Range;
+
+// CHECK-LABEL: @index_by_range(
+#[no_mangle]
+pub fn index_by_range(x: &[u16], r: Range<usize>) -> &[u16] {
+    // CHECK: sub nuw i64
+    &x[r]
+}
+
+// CHECK-LABEL: @get_unchecked_by_range(
+#[no_mangle]
+pub unsafe fn get_unchecked_by_range(x: &[u16], r: Range<usize>) -> &[u16] {
+    // CHECK: sub nuw i64
+    x.get_unchecked(r)
+}
+
+// CHECK-LABEL: @index_mut_by_range(
+#[no_mangle]
+pub fn index_mut_by_range(x: &mut [i32], r: Range<usize>) -> &mut [i32] {
+    // CHECK: sub nuw i64
+    &mut x[r]
+}
+
+// CHECK-LABEL: @get_unchecked_mut_by_range(
+#[no_mangle]
+pub unsafe fn get_unchecked_mut_by_range(x: &mut [i32], r: Range<usize>) -> &mut [i32] {
+    // CHECK: sub nuw i64
+    x.get_unchecked_mut(r)
+}
diff --git a/tests/ui/consts/const-eval/ub-slice-get-unchecked.rs b/tests/ui/consts/const-eval/ub-slice-get-unchecked.rs
new file mode 100644
index 00000000000..d9a74b4f3e2
--- /dev/null
+++ b/tests/ui/consts/const-eval/ub-slice-get-unchecked.rs
@@ -0,0 +1,9 @@
+#![feature(const_slice_index)]
+
+const A: [(); 5] = [(), (), (), (), ()];
+
+// Since the indexing is on a ZST, the addresses are all fine,
+// but we should still catch the bad range.
+const B: &[()] = unsafe { A.get_unchecked(3..1) };
+
+fn main() {}
diff --git a/tests/ui/consts/const-eval/ub-slice-get-unchecked.stderr b/tests/ui/consts/const-eval/ub-slice-get-unchecked.stderr
new file mode 100644
index 00000000000..775e475dfeb
--- /dev/null
+++ b/tests/ui/consts/const-eval/ub-slice-get-unchecked.stderr
@@ -0,0 +1,18 @@
+error[E0080]: evaluation of constant value failed
+  --> $SRC_DIR/core/src/slice/index.rs:LL:COL
+   |
+   = note: overflow executing `unchecked_sub`
+   |
+note: inside `<std::ops::Range<usize> as SliceIndex<[()]>>::get_unchecked`
+  --> $SRC_DIR/core/src/slice/index.rs:LL:COL
+note: inside `core::slice::<impl [()]>::get_unchecked::<std::ops::Range<usize>>`
+  --> $SRC_DIR/core/src/slice/mod.rs:LL:COL
+note: inside `B`
+  --> $DIR/ub-slice-get-unchecked.rs:7:27
+   |
+LL | const B: &[()] = unsafe { A.get_unchecked(3..1) };
+   |                           ^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0080`.