Auto merge of rust-lang#122926 - scottmcm:assume-lengths, r=<try>

Add `assume`s to slice length calls

Since `.len()` on slices is safe, let's see how this impacts things vs what we could do with rust-lang#121965 and LLVM 19

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -748,7 +748,19 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
         // use common size calculation for non zero-sized types
         let cg_value = self.codegen_place(bx, place.as_ref());
-        cg_value.len(bx.cx())
+        let length = cg_value.len(bx.cx());
+
+        if bx.sess().opts.optimize != OptLevel::No {
+            let elem_ty = cg_value.layout.field(bx, 0);
+            if let Some(elem_bytes) = std::num::NonZeroU64::new(elem_ty.size.bytes()) {
+                let isize_max = (1_u64 << (bx.sess().target.pointer_width - 1)) - 1;
+                let len_max = isize_max / elem_bytes;
+                let limit = bx.icmp(IntPredicate::IntULE, length, bx.const_usize(len_max));
+                bx.assume(limit);
+            }
+        }
+
+        length
     }
 
     /// Codegen an `Rvalue::AddressOf` or `Rvalue::Ref`

tests/codegen/issues/issue-98294-get-mut-copy-from-slice-opt.rs

@@ -8,6 +8,10 @@
 pub fn test(a: &mut [u8], offset: usize, bytes: &[u8]) {
     // CHECK-LABEL: @test(
     // CHECK-NOT: call
+    // CHECK: call void @llvm.assume
+    // CHECK-NOT: call
+    // CHECK: call void @llvm.assume
+    // CHECK-NOT: call
     // CHECK: call void @llvm.memcpy
     // CHECK-NOT: call
     // CHECK: }

tests/codegen/slice-as_chunks.rs

@@ -7,22 +7,23 @@
 // CHECK-LABEL: @chunks4
 #[no_mangle]
 pub fn chunks4(x: &[u8]) -> &[[u8; 4]] {
-    // CHECK-NEXT: start:
-    // CHECK-NEXT: lshr i64 %x.1, 2
     // CHECK-NOT: shl
     // CHECK-NOT: mul
     // CHECK-NOT: udiv
     // CHECK-NOT: urem
-    // CHECK: ret
+    // CHECK: %[[NEWLEN:.+]] = lshr i64 %x.1, 2
+    // CHECK: %[[A:.+]] = insertvalue { ptr, i64 } poison, ptr %x.0, 0
+    // CHECK: %[[B:.+]] = insertvalue { ptr, i64 } %[[A]], i64 %[[NEWLEN]], 1
+    // CHECK: ret { ptr, i64 } %[[B]]
     x.as_chunks().0
 }
 
 // CHECK-LABEL: @chunks4_with_remainder
 #[no_mangle]
 pub fn chunks4_with_remainder(x: &[u8]) -> (&[[u8; 4]], &[u8]) {
-    // CHECK-DAG: and i64 %x.1, -4
+    // CHECK-DAG: and i64 %x.1, [[#0x7FFFFFFFFFFFFFFC]]
     // CHECK-DAG: and i64 %x.1, 3
-    // CHECK-DAG: lshr
+    // CHECK-DAG: lshr i64 %x.1, 2
     // CHECK-NOT: mul
     // CHECK-NOT: udiv
     // CHECK-NOT: urem

tests/codegen/slice-iter-fold.rs

@@ -2,11 +2,13 @@
 #![crate_type = "lib"]
 
 // CHECK-LABEL: @slice_fold_to_last
+// CHECK-SAME: %slice.0, [[USIZE:i[0-9]+]] noundef %slice.1)
 #[no_mangle]
 pub fn slice_fold_to_last(slice: &[i32]) -> Option<&i32> {
-    // CHECK-NOT: loop
-    // CHECK-NOT: br
-    // CHECK-NOT: call
-    // CHECK: ret
+    // CHECK: %[[END:.+]] = getelementptr inbounds i32, ptr %slice.0, [[USIZE]] %slice.1
+    // CHECK: %[[EMPTY:.+]] = icmp eq [[USIZE]] %slice.1, 0
+    // CHECK: %[[LAST:.+]] = getelementptr i32, ptr %[[END]], i64 -1
+    // CHECK: %[[R:.+]] = select i1 %[[EMPTY]], ptr null, ptr %[[LAST]]
+    // CHECK: ret ptr %[[R]]
     slice.iter().fold(None, |_, i| Some(i))
 }

tests/codegen/slice-iter-nonnull.rs

@@ -52,9 +52,13 @@ pub fn slice_iter_next_back<'a>(it: &mut std::slice::Iter<'a, u32>) -> Option<&'
 // CHECK-LABEL: @slice_iter_new
 // CHECK-SAME: (ptr noalias noundef nonnull {{.+}} %slice.0, {{.+}} noundef %slice.1)
 #[no_mangle]
-pub fn slice_iter_new(slice: &[u32]) -> std::slice::Iter<'_, u32> {
+pub fn slice_iter_new(slice: &[u8]) -> std::slice::Iter<'_, u8> {
     // CHECK-NOT: slice
-    // CHECK: %[[END:.+]] = getelementptr inbounds i32{{.+}} %slice.0{{.+}} %slice.1
+    // CHECK: %[[NONNEG:.+]] = icmp sgt i64 %slice.1, -1
+    // CHECK-NOT: slice
+    // CHECK: call void @llvm.assume(i1 %[[NONNEG]])
+    // CHECK-NOT: slice
+    // CHECK: %[[END:.+]] = getelementptr inbounds i8{{.+}} %slice.0{{.+}} %slice.1
     // CHECK-NOT: slice
     // CHECK: insertvalue {{.+}} ptr %slice.0, 0
     // CHECK-NOT: slice
@@ -67,9 +71,13 @@ pub fn slice_iter_new(slice: &[u32]) -> std::slice::Iter<'_, u32> {
 // CHECK-LABEL: @slice_iter_mut_new
 // CHECK-SAME: (ptr noalias noundef nonnull {{.+}} %slice.0, {{.+}} noundef %slice.1)
 #[no_mangle]
-pub fn slice_iter_mut_new(slice: &mut [u32]) -> std::slice::IterMut<'_, u32> {
+pub fn slice_iter_mut_new(slice: &mut [u8]) -> std::slice::IterMut<'_, u8> {
+    // CHECK-NOT: slice
+    // CHECK: %[[NONNEG:.+]] = icmp sgt i64 %slice.1, -1
+    // CHECK-NOT: slice
+    // CHECK: call void @llvm.assume(i1 %[[NONNEG]])
     // CHECK-NOT: slice
-    // CHECK: %[[END:.+]] = getelementptr inbounds i32{{.+}} %slice.0{{.+}} %slice.1
+    // CHECK: %[[END:.+]] = getelementptr inbounds i8{{.+}} %slice.0{{.+}} %slice.1
     // CHECK-NOT: slice
     // CHECK: insertvalue {{.+}} ptr %slice.0, 0
     // CHECK-NOT: slice

tests/codegen/slice-length-limits.rs

@@ -0,0 +1,55 @@
+//@ revisions: 32BIT 64BIT
+//@ compile-flags: -O
+//@ min-llvm-version: 18.0
+//@ [32BIT] only-32bit
+//@ [64BIT] only-64bit
+
+#![crate_type = "lib"]
+
+// Confirm that the `assume` calls from the length allows LLVM to know that some
+// math on the indices can be done without overflow risk.
+
+// CHECK-LABEL: @slice_length_demo
+#[no_mangle]
+pub unsafe fn slice_length_demo(x: &[u16]) {
+    // 32BIT: %[[LIMIT:.+]] = icmp ult [[USIZE:i32]] %x.1, [[#0x40000000]]
+    // 64BIT: %[[LIMIT:.+]] = icmp ult [[USIZE:i64]] %x.1, [[#0x4000000000000000]]
+    // CHECK: tail call void @llvm.assume(i1 %[[LIMIT]])
+
+    // CHECK: %[[Y:.+]] = phi [[USIZE]]
+    // CHECK-SAME: [ 0, %start ]
+    // CHECK: %[[PLUS_ONE:.+]] = add nuw nsw [[USIZE]] %[[Y]], 1
+    // CHECK: call void @do_something([[USIZE]] noundef %[[PLUS_ONE]])
+    // CHECK: %[[TIMES_TWO:.+]] = shl nuw nsw [[USIZE]] %[[Y]], 1
+    // CHECK: call void @do_something([[USIZE]] noundef %[[TIMES_TWO]])
+    for y in 0..x.len() {
+        do_something(y + 1);
+        do_something(y * 2);
+    }
+}
+
+// CHECK-LABEL: @nested_slice_length
+#[no_mangle]
+pub unsafe fn nested_slice_length(x: &[f32], y: &[f32]) {
+    // 32BIT: %[[LIMIT:.+]] = icmp ult [[USIZE:i32]] %x.1, [[#0x20000000]]
+    // 64BIT: %[[LIMIT:.+]] = icmp ult [[USIZE:i64]] %x.1, [[#0x2000000000000000]]
+    // CHECK: tail call void @llvm.assume(i1 %[[LIMIT]])
+    // 32BIT: %[[LIMIT:.+]] = icmp ult [[USIZE]] %y.1, [[#0x20000000]]
+    // 64BIT: %[[LIMIT:.+]] = icmp ult [[USIZE]] %y.1, [[#0x2000000000000000]]
+    // CHECK: tail call void @llvm.assume(i1 %[[LIMIT]])
+
+    // CHECK: %[[J:.+]] = phi [[USIZE]]
+    // CHECK: %[[I:.+]] = phi [[USIZE]]
+    // CHECK-NOT: phi
+    // CHECK: %[[T:.+]] = add nuw nsw [[USIZE]] %[[I]], %[[J]]
+    // CHECK: call void @do_something([[USIZE]] noundef %[[T]])
+    for i in 0..x.len() {
+        for j in 0..y.len() {
+            do_something(i + j);
+        }
+    }
+}
+
+extern {
+    fn do_something(x: usize);
+}

tests/codegen/slice-ref-equality.rs

@@ -15,7 +15,7 @@ use std::num::NonZero;
 // CHECK-LABEL: @is_zero_slice_long
 #[no_mangle]
 pub fn is_zero_slice_long(data: &[u8; 456]) -> bool {
-    // CHECK: %[[BCMP:.+]] = tail call i32 @{{bcmp|memcmp}}({{.+}})
+    // CHECK: %[[BCMP:.+]] = tail call i32 @{{bcmp|memcmp}}({{.+}} 456)
     // CHECK-NEXT: %[[EQ:.+]] = icmp eq i32 %[[BCMP]], 0
     // CHECK-NEXT: ret i1 %[[EQ]]
     &data[..] == [0; 456]
@@ -48,7 +48,7 @@ pub fn is_zero_array(data: &[u8; 4]) -> bool {
 #[no_mangle]
 fn eq_slice_of_nested_u8(x: &[[u8; 3]], y: &[[u8; 3]]) -> bool {
     // CHECK: icmp eq [[USIZE]] %1, %3
-    // CHECK: %[[BYTES:.+]] = mul nsw [[USIZE]] %1, 3
+    // CHECK: %[[BYTES:.+]] = mul nuw nsw [[USIZE]] %1, 3
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y
@@ -60,7 +60,7 @@ fn eq_slice_of_nested_u8(x: &[[u8; 3]], y: &[[u8; 3]]) -> bool {
 #[no_mangle]
 fn eq_slice_of_i32(x: &[i32], y: &[i32]) -> bool {
     // CHECK: icmp eq [[USIZE]] %1, %3
-    // CHECK: %[[BYTES:.+]] = shl nsw [[USIZE]] %1, 2
+    // CHECK: %[[BYTES:.+]] = shl nuw nsw [[USIZE]] %1, 2
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y
@@ -72,7 +72,7 @@ fn eq_slice_of_i32(x: &[i32], y: &[i32]) -> bool {
 #[no_mangle]
 fn eq_slice_of_nonzero(x: &[NonZero<i32>], y: &[NonZero<i32>]) -> bool {
     // CHECK: icmp eq [[USIZE]] %1, %3
-    // CHECK: %[[BYTES:.+]] = shl nsw [[USIZE]] %1, 2
+    // CHECK: %[[BYTES:.+]] = shl nuw nsw [[USIZE]] %1, 2
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y
@@ -84,7 +84,7 @@ fn eq_slice_of_nonzero(x: &[NonZero<i32>], y: &[NonZero<i32>]) -> bool {
 #[no_mangle]
 fn eq_slice_of_option_of_nonzero(x: &[Option<NonZero<i16>>], y: &[Option<NonZero<i16>>]) -> bool {
     // CHECK: icmp eq [[USIZE]] %1, %3
-    // CHECK: %[[BYTES:.+]] = shl nsw [[USIZE]] %1, 1
+    // CHECK: %[[BYTES:.+]] = shl nuw nsw [[USIZE]] %1, 1
     // CHECK: tail call{{( noundef)?}} i32 @{{bcmp|memcmp}}(ptr
     // CHECK-SAME: , [[USIZE]]{{( noundef)?}} %[[BYTES]])
     x == y

tests/codegen/vecdeque-drain.rs

@@ -43,6 +43,8 @@ pub fn truncate(v: &mut VecDeque<i32>, n: usize) {
 // CHECK-NOT: call
 // CHECK: br
 // CHECK-NOT: call
+// CHECK: call void @llvm.assume
+// CHECK-NOT: call
 // CHECK: br
 // CHECK-NOT: call
 // CHECK: br