Auto merge of #123575 - scottmcm:assume-raw-parts, r=<try>

Add some targeted `assume(bytes <= isize::MAX)` around slices

Putting it on every single slice length was pretty expensive, but let's see if doing it just in `slice::from_raw_parts(_mut)` and `slice::Iter(Mut)::len` works better.

Author: bors

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -1,6 +1,7 @@
 use super::operand::{OperandRef, OperandValue};
 use super::place::PlaceRef;
 use super::FunctionCx;
+use crate::common::IntPredicate;
 use crate::errors;
 use crate::errors::InvalidMonomorphization;
 use crate::meth;
@@ -499,6 +500,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     // The `_unsigned` version knows the relative ordering of the pointers,
                     // so can use `sub nuw` and `udiv exact` instead of dealing in signed.
                     let d = bx.unchecked_usub(a, b);
+                    if let OptLevel::Default | OptLevel::Aggressive = bx.sess().opts.optimize {
+                        // Despite dealing in `usize`, it's still not allowed for the
+                        // pointers to be more than `isize::MAX` apart.
+                        let inrange = bx.icmp(IntPredicate::IntSGE, d, bx.const_usize(0));
+                        bx.assume(inrange);
+                    }
                     bx.exactudiv(d, pointee_size)
                 }
             }

library/core/src/mem/mod.rs

@@ -1235,6 +1235,17 @@ pub trait SizedTypeProperties: Sized {
     #[doc(hidden)]
     #[unstable(feature = "sized_type_properties", issue = "none")]
     const IS_ZST: bool = size_of::<Self>() == 0;
+
+    /// The largest safe length for a `[Self]`.
+    ///
+    /// Anything larger than this would make `size_of_val` overflow `isize::MAX`,
+    /// which is never allowed for a single object.
+    #[doc(hidden)]
+    #[unstable(feature = "sized_type_properties", issue = "none")]
+    const MAX_SLICE_LEN: usize = match size_of::<Self>() {
+        0 => usize::MAX,
+        n => (isize::MAX as usize) / n,
+    };
 }
 #[doc(hidden)]
 #[unstable(feature = "sized_type_properties", issue = "none")]

library/core/src/slice/raw.rs

@@ -1,7 +1,8 @@
 //! Free functions to create `&[T]` and `&mut [T]`.
 
 use crate::array;
-use crate::mem::{align_of, size_of};
+use crate::intrinsics;
+use crate::mem::{align_of, SizedTypeProperties};
 use crate::ops::Range;
 use crate::ptr;
 use crate::ub_checks;
@@ -98,13 +99,14 @@ pub const unsafe fn from_raw_parts<'a, T>(data: *const T, len: usize) -> &'a [T]
             "slice::from_raw_parts requires the pointer to be aligned and non-null, and the total size of the slice not to exceed `isize::MAX`",
             (
                 data: *mut () = data as *mut (),
-                size: usize = size_of::<T>(),
                 align: usize = align_of::<T>(),
                 len: usize = len,
+                max_len: usize = T::MAX_SLICE_LEN,
             ) =>
             ub_checks::is_aligned_and_not_null(data, align)
-                && ub_checks::is_valid_allocation_size(size, len)
+                && len <= max_len
         );
+        intrinsics::assume(len <= T::MAX_SLICE_LEN);
         &*ptr::slice_from_raw_parts(data, len)
     }
 }
@@ -152,13 +154,14 @@ pub const unsafe fn from_raw_parts_mut<'a, T>(data: *mut T, len: usize) -> &'a m
             "slice::from_raw_parts_mut requires the pointer to be aligned and non-null, and the total size of the slice not to exceed `isize::MAX`",
             (
                 data: *mut () = data as *mut (),
-                size: usize = size_of::<T>(),
                 align: usize = align_of::<T>(),
                 len: usize = len,
+                max_len: usize = T::MAX_SLICE_LEN,
             ) =>
             ub_checks::is_aligned_and_not_null(data, align)
-                && ub_checks::is_valid_allocation_size(size, len)
+                && len <= max_len
         );
+        intrinsics::assume(len <= T::MAX_SLICE_LEN);
         &mut *ptr::slice_from_raw_parts_mut(data, len)
     }
 }

library/core/src/ub_checks.rs

@@ -115,12 +115,6 @@ pub(crate) const fn is_aligned_and_not_null(ptr: *const (), align: usize) -> boo
     !ptr.is_null() && ptr.is_aligned_to(align)
 }
 
-#[inline]
-pub(crate) const fn is_valid_allocation_size(size: usize, len: usize) -> bool {
-    let max_len = if size == 0 { usize::MAX } else { isize::MAX as usize / size };
-    len <= max_len
-}
-
 /// Checks whether the regions of memory starting at `src` and `dst` of size
 /// `count * size` do *not* overlap.
 ///

tests/codegen/intrinsics/offset_from.rs

@@ -1,4 +1,6 @@
-//@ compile-flags: -C opt-level=1
+//@ revisions: OPT1 OPT2
+//@ [OPT1] compile-flags: -Copt-level=1
+//@ [OPT2] compile-flags: -Copt-level=2
 //@ only-64bit (because we're using [ui]size)
 
 #![crate_type = "lib"]
@@ -17,8 +19,9 @@ pub unsafe fn offset_from_odd_size(a: *const RGB, b: *const RGB) -> isize {
     // CHECK: start
     // CHECK-NEXT: ptrtoint
     // CHECK-NEXT: ptrtoint
-    // CHECK-NEXT: sub i64
-    // CHECK-NEXT: sdiv exact i64 %{{[0-9]+}}, 3
+    // CHECK-NEXT: %[[D:.+]] = sub i64
+    // CHECK-NOT: assume
+    // CHECK-NEXT: sdiv exact i64 %[[D]], 3
     // CHECK-NEXT: ret i64
     ptr_offset_from(a, b)
 }
@@ -29,8 +32,11 @@ pub unsafe fn offset_from_unsigned_odd_size(a: *const RGB, b: *const RGB) -> usi
     // CHECK: start
     // CHECK-NEXT: ptrtoint
     // CHECK-NEXT: ptrtoint
-    // CHECK-NEXT: sub nuw i64
-    // CHECK-NEXT: udiv exact i64 %{{[0-9]+}}, 3
+    // CHECK-NEXT: %[[D:.+]] = sub nuw i64
+    // OPT1-NOT: assume
+    // OPT2-NEXT: %[[POS:.+]] = icmp sgt i64 %[[D]], -1
+    // OPT2-NEXT: tail call void @llvm.assume(i1 %[[POS]])
+    // CHECK-NEXT: udiv exact i64 %[[D]], 3
     // CHECK-NEXT: ret i64
     ptr_offset_from_unsigned(a, b)
 }

tests/codegen/slice-iter-len-eq-zero.rs

@@ -4,10 +4,21 @@
 type Demo = [u8; 3];
 
 // CHECK-LABEL: @slice_iter_len_eq_zero
+// CHECK-SAME: ptr noundef nonnull %0
+// CHECK-SAME: ptr noundef %1
 #[no_mangle]
 pub fn slice_iter_len_eq_zero(y: std::slice::Iter<'_, Demo>) -> bool {
-    // CHECK-NOT: sub
-    // CHECK: %[[RET:.+]] = icmp eq ptr {{%1|%0}}, {{%1|%0}}
+    // There's no `nonnull` in the signature, so we end up with an assume.
+    // CHECK: %[[NNULL:.+]] = icmp ne ptr %1, null
+    // CHECK: tail call void @llvm.assume(i1 %[[NNULL]])
+
+    // CHECK-DAG: %[[E:.+]] = ptrtoint ptr %1 to i
+    // CHECK-DAG: %[[B:.+]] = ptrtoint ptr %0 to i
+    // CHECK: %[[D:.+]] = sub nuw {{.+}} %[[E]], %[[B]]
+    // CHECK: %[[NNEG:.+]] = icmp sgt {{.+}} %[[D]], -1
+    // CHECK: tail call void @llvm.assume(i1 %[[NNEG]])
+
+    // CHECK: %[[RET:.+]] = icmp eq ptr {{%1, %0|%0, %1}}
     // CHECK: ret i1 %[[RET]]
     y.len() == 0
 }

tests/codegen/vec-in-place.rs

@@ -38,6 +38,13 @@ pub struct Baz {
 pub fn vec_iterator_cast_primitive(vec: Vec<i8>) -> Vec<u8> {
     // CHECK-NOT: loop
     // CHECK-NOT: call
+
+    // CHECK: %[[LEN_IN_ISIZE:.+]] = icmp sgt i{{.+}}, -1
+    // CHECK-NEXT: call void @llvm.assume(i1 %[[LEN_IN_ISIZE]])
+
+    // CHECK-NOT: loop
+    // CHECK-NOT: call
+
     vec.into_iter().map(|e| e as u8).collect()
 }
 
@@ -46,6 +53,13 @@ pub fn vec_iterator_cast_primitive(vec: Vec<i8>) -> Vec<u8> {
 pub fn vec_iterator_cast_wrapper(vec: Vec<u8>) -> Vec<Wrapper<u8>> {
     // CHECK-NOT: loop
     // CHECK-NOT: call
+
+    // CHECK: %[[LEN_IN_ISIZE:.+]] = icmp sgt i{{.+}}, -1
+    // CHECK-NEXT: call void @llvm.assume(i1 %[[LEN_IN_ISIZE]])
+
+    // CHECK-NOT: loop
+    // CHECK-NOT: call
+
     vec.into_iter().map(|e| Wrapper(e)).collect()
 }
 
@@ -54,6 +68,13 @@ pub fn vec_iterator_cast_wrapper(vec: Vec<u8>) -> Vec<Wrapper<u8>> {
 pub fn vec_iterator_cast_unwrap(vec: Vec<Wrapper<u8>>) -> Vec<u8> {
     // CHECK-NOT: loop
     // CHECK-NOT: call
+
+    // CHECK: %[[LEN_IN_ISIZE:.+]] = icmp sgt i{{.+}}, -1
+    // CHECK-NEXT: call void @llvm.assume(i1 %[[LEN_IN_ISIZE]])
+
+    // CHECK-NOT: loop
+    // CHECK-NOT: call
+
     vec.into_iter().map(|e| e.0).collect()
 }
 
@@ -90,3 +111,9 @@ pub fn vec_iterator_cast_deaggregate_fold(vec: Vec<Baz>) -> Vec<[u64; 4]> {
     // correct.
     vec.into_iter().map(|e| unsafe { std::mem::transmute(e) }).collect()
 }
+
+// Make sure that the `NOT`s above don't trigger on the `nocallback` attribute further down.
+
+// CHECK-LABEL: @last_definition
+#[no_mangle]
+pub fn last_definition() {}