transmute should also assume non-null pointers

scottmcm · scottmcm · commit 0cc14b688d27 · 2025-02-12T23:01:27.000-08:00
Previously it only did integer-ABI things, but this way it does data pointers too.  That gives more information in general to the backend, and allows slightly simplifying one of the helpers in slice iterators.
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -435,18 +435,22 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         scalar: abi::Scalar,
         backend_ty: Bx::Type,
     ) {
-        if matches!(self.cx.sess().opts.optimize, OptLevel::No)
-            // For now, the critical niches are all over `Int`eger values.
-            // Should floating-point values or pointers ever get more complex
-            // niches, then this code will probably want to handle them too.
-            || !matches!(scalar.primitive(), abi::Primitive::Int(..))
-            || scalar.is_always_valid(self.cx)
-        {
+        if matches!(self.cx.sess().opts.optimize, OptLevel::No) || scalar.is_always_valid(self.cx) {
             return;
         }
 
-        let range = scalar.valid_range(self.cx);
-        bx.assume_integer_range(imm, backend_ty, range);
+        match scalar.primitive() {
+            abi::Primitive::Int(..) => {
+                let range = scalar.valid_range(self.cx);
+                bx.assume_integer_range(imm, backend_ty, range);
+            }
+            abi::Primitive::Pointer(abi::AddressSpace::DATA)
+                if !scalar.valid_range(self.cx).contains(0) =>
+            {
+                bx.assume_nonnull(imm);
+            }
+            abi::Primitive::Pointer(..) | abi::Primitive::Float(..) => {}
+        }
     }
 
     pub(crate) fn codegen_rvalue_unsized(
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -243,6 +243,19 @@ pub trait BuilderMethods<'a, 'tcx>:
         self.assume(cmp);
     }
 
+    /// Emits an `assume` that the `val` of pointer type is non-null.
+    ///
+    /// You may want to check the optimization level before bothering calling this.
+    fn assume_nonnull(&mut self, val: Self::Value) {
+        // Arguably in LLVM it'd be better to emit an assume operand bundle instead
+        // <https://llvm.org/docs/LangRef.html#assume-operand-bundles>
+        // but this works fine for all backends.
+
+        let null = self.const_null(self.type_ptr());
+        let is_null = self.icmp(IntPredicate::IntNE, val, null);
+        self.assume(is_null);
+    }
+
     fn range_metadata(&mut self, load: Self::Value, range: WrappingRange);
     fn nonnull_metadata(&mut self, load: Self::Value);
 
diff --git a/library/core/src/slice/iter/macros.rs b/library/core/src/slice/iter/macros.rs
@@ -30,7 +30,7 @@ macro_rules! if_zst {
             $zst_body
         } else {
             // SAFETY: for non-ZSTs, the type invariant ensures it cannot be null
-            let $end = unsafe { *(&raw const $this.end_or_len).cast::<NonNull<T>>() };
+            let $end = unsafe { mem::transmute::<*const T, NonNull<T>>($this.end_or_len) };
             $other_body
         }
     }};
diff --git a/tests/codegen/intrinsics/transmute-niched.rs b/tests/codegen/intrinsics/transmute-niched.rs
@@ -6,6 +6,7 @@
 
 use std::mem::transmute;
 use std::num::NonZero;
+use std::ptr::NonNull;
 
 #[repr(u8)]
 pub enum SmallEnum {
@@ -192,3 +193,31 @@ pub unsafe fn check_bool_to_ordering(x: bool) -> std::cmp::Ordering {
 
     transmute(x)
 }
+
+// CHECK-LABEL: @check_nonnull_to_ptr(
+#[no_mangle]
+pub unsafe fn check_nonnull_to_ptr(x: NonNull<u8>) -> *const u8 {
+    // CHECK-NOT: icmp
+    // CHECK-NOT: assume
+    // OPT: %0 = icmp ne ptr %x, null
+    // OPT: call void @llvm.assume(i1 %0)
+    // CHECK-NOT: icmp
+    // CHECK-NOT: assume
+    // CHECK: ret ptr %x
+
+    transmute(x)
+}
+
+// CHECK-LABEL: @check_ptr_to_nonnull(
+#[no_mangle]
+pub unsafe fn check_ptr_to_nonnull(x: *const u8) -> NonNull<u8> {
+    // CHECK-NOT: icmp
+    // CHECK-NOT: assume
+    // OPT: %0 = icmp ne ptr %x, null
+    // OPT: call void @llvm.assume(i1 %0)
+    // CHECK-NOT: icmp
+    // CHECK-NOT: assume
+    // CHECK: ret ptr %x
+
+    transmute(x)
+}
diff --git a/tests/codegen/intrinsics/transmute.rs b/tests/codegen/intrinsics/transmute.rs
@@ -370,9 +370,11 @@ pub unsafe fn check_issue_110005(x: (usize, bool)) -> Option<Box<[u8]>> {
 #[no_mangle]
 pub unsafe fn check_pair_to_dst_ref<'a>(x: (usize, usize)) -> &'a [u8] {
     // CHECK: %_0.0 = getelementptr i8, ptr null, i64 %x.0
-    // CHECK: %0 = insertvalue { ptr, i64 } poison, ptr %_0.0, 0
-    // CHECK: %1 = insertvalue { ptr, i64 } %0, i64 %x.1, 1
-    // CHECK: ret { ptr, i64 } %1
+    // CHECK: %0 = icmp ne ptr %_0.0, null
+    // CHECK: call void @llvm.assume(i1 %0)
+    // CHECK: %1 = insertvalue { ptr, i64 } poison, ptr %_0.0, 0
+    // CHECK: %2 = insertvalue { ptr, i64 } %1, i64 %x.1, 1
+    // CHECK: ret { ptr, i64 } %2
     transmute(x)
 }
 
diff --git a/tests/codegen/slice-iter-len-eq-zero.rs b/tests/codegen/slice-iter-len-eq-zero.rs
@@ -1,4 +1,5 @@
 //@ compile-flags: -Copt-level=3
+//@ needs-deterministic-layouts (opposite scalar pair orders breaks it)
 #![crate_type = "lib"]
 
 type Demo = [u8; 3];
@@ -7,7 +8,40 @@ type Demo = [u8; 3];
 #[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}}
+    // CHECK: %[[RET:.+]] = icmp eq ptr {{%y.0, %y.1|%y.1, %y.0}}
+    // CHECK: ret i1 %[[RET]]
+    y.len() == 0
+}
+
+// CHECK-LABEL: @slice_iter_len_eq_zero_ref
+#[no_mangle]
+pub fn slice_iter_len_eq_zero_ref(y: &mut std::slice::Iter<'_, Demo>) -> bool {
+    // CHECK-NOT: sub
+    // CHECK: %[[A:.+]] = load ptr
+    // CHECK-SAME: !nonnull
+    // CHECK: %[[B:.+]] = load ptr
+    // CHECK-SAME: !nonnull
+    // CHECK: %[[RET:.+]] = icmp eq ptr %[[A]], %[[B]]
+    // CHECK: ret i1 %[[RET]]
+    y.len() == 0
+}
+
+struct MyZST;
+
+// CHECK-LABEL: @slice_zst_iter_len_eq_zero
+#[no_mangle]
+pub fn slice_zst_iter_len_eq_zero(y: std::slice::Iter<'_, MyZST>) -> bool {
+    // CHECK: %[[RET:.+]] = icmp eq ptr %y.1, null
+    // CHECK: ret i1 %[[RET]]
+    y.len() == 0
+}
+
+// CHECK-LABEL: @slice_zst_iter_len_eq_zero_ref
+#[no_mangle]
+pub fn slice_zst_iter_len_eq_zero_ref(y: &mut std::slice::Iter<'_, MyZST>) -> bool {
+    // CHECK: %[[LEN:.+]] = load ptr
+    // CHECK-NOT: !nonnull
+    // CHECK: %[[RET:.+]] = icmp eq ptr %[[LEN]], null
     // CHECK: ret i1 %[[RET]]
     y.len() == 0
 }
diff --git a/tests/mir-opt/pre-codegen/slice_iter.slice_iter_generic_is_empty.PreCodegen.after.panic-abort.mir b/tests/mir-opt/pre-codegen/slice_iter.slice_iter_generic_is_empty.PreCodegen.after.panic-abort.mir
@@ -4,73 +4,67 @@ fn slice_iter_generic_is_empty(_1: &std::slice::Iter<'_, T>) -> bool {
     debug it => _1;
     let mut _0: bool;
     scope 1 (inlined <std::slice::Iter<'_, T> as ExactSizeIterator>::is_empty) {
-        let mut _2: *const *const T;
-        let mut _3: *const std::ptr::NonNull<T>;
-        let mut _8: *const T;
+        let mut _2: *const T;
+        let mut _7: *const T;
         scope 2 {
-            let _4: std::ptr::NonNull<T>;
-            let _9: usize;
+            let _3: std::ptr::NonNull<T>;
+            let _8: usize;
             scope 3 {
             }
             scope 4 {
-                scope 8 (inlined <NonNull<T> as PartialEq>::eq) {
-                    let mut _5: std::ptr::NonNull<T>;
+                scope 7 (inlined <NonNull<T> as PartialEq>::eq) {
+                    let mut _4: std::ptr::NonNull<T>;
+                    let mut _5: *mut T;
                     let mut _6: *mut T;
-                    let mut _7: *mut T;
-                    scope 9 (inlined NonNull::<T>::as_ptr) {
+                    scope 8 (inlined NonNull::<T>::as_ptr) {
                     }
-                    scope 10 (inlined NonNull::<T>::as_ptr) {
+                    scope 9 (inlined NonNull::<T>::as_ptr) {
                     }
                 }
             }
             scope 5 (inlined std::ptr::const_ptr::<impl *const T>::addr) {
                 scope 6 (inlined std::ptr::const_ptr::<impl *const T>::cast::<()>) {
                 }
             }
-            scope 7 (inlined std::ptr::const_ptr::<impl *const *const T>::cast::<NonNull<T>>) {
-            }
         }
     }
 
     bb0: {
-        StorageLive(_9);
         StorageLive(_8);
-        StorageLive(_4);
+        StorageLive(_7);
+        StorageLive(_3);
         switchInt(const <T as std::mem::SizedTypeProperties>::IS_ZST) -> [0: bb1, otherwise: bb2];
     }
 
     bb1: {
-        StorageLive(_3);
         StorageLive(_2);
-        _2 = &raw const ((*_1).1: *const T);
-        _3 = copy _2 as *const std::ptr::NonNull<T> (PtrToPtr);
+        _2 = copy ((*_1).1: *const T);
+        _3 = move _2 as std::ptr::NonNull<T> (Transmute);
         StorageDead(_2);
-        _4 = copy (*_3);
-        StorageDead(_3);
-        StorageLive(_6);
         StorageLive(_5);
-        _5 = copy ((*_1).0: std::ptr::NonNull<T>);
-        _6 = copy _5 as *mut T (Transmute);
-        StorageDead(_5);
-        StorageLive(_7);
-        _7 = copy _4 as *mut T (Transmute);
-        _0 = Eq(move _6, move _7);
-        StorageDead(_7);
+        StorageLive(_4);
+        _4 = copy ((*_1).0: std::ptr::NonNull<T>);
+        _5 = copy _4 as *mut T (Transmute);
+        StorageDead(_4);
+        StorageLive(_6);
+        _6 = copy _3 as *mut T (Transmute);
+        _0 = Eq(move _5, move _6);
         StorageDead(_6);
+        StorageDead(_5);
         goto -> bb3;
     }
 
     bb2: {
-        _8 = copy ((*_1).1: *const T);
-        _9 = copy _8 as usize (Transmute);
-        _0 = Eq(copy _9, const 0_usize);
+        _7 = copy ((*_1).1: *const T);
+        _8 = copy _7 as usize (Transmute);
+        _0 = Eq(copy _8, const 0_usize);
         goto -> bb3;
     }
 
     bb3: {
-        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_7);
         StorageDead(_8);
-        StorageDead(_9);
         return;
     }
 }
diff --git a/tests/mir-opt/pre-codegen/slice_iter.slice_iter_generic_is_empty.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/pre-codegen/slice_iter.slice_iter_generic_is_empty.PreCodegen.after.panic-unwind.mir
@@ -4,73 +4,67 @@ fn slice_iter_generic_is_empty(_1: &std::slice::Iter<'_, T>) -> bool {
     debug it => _1;
     let mut _0: bool;
     scope 1 (inlined <std::slice::Iter<'_, T> as ExactSizeIterator>::is_empty) {
-        let mut _2: *const *const T;
-        let mut _3: *const std::ptr::NonNull<T>;
-        let mut _8: *const T;
+        let mut _2: *const T;
+        let mut _7: *const T;
         scope 2 {
-            let _4: std::ptr::NonNull<T>;
-            let _9: usize;
+            let _3: std::ptr::NonNull<T>;
+            let _8: usize;
             scope 3 {
             }
             scope 4 {
-                scope 8 (inlined <NonNull<T> as PartialEq>::eq) {
-                    let mut _5: std::ptr::NonNull<T>;
+                scope 7 (inlined <NonNull<T> as PartialEq>::eq) {
+                    let mut _4: std::ptr::NonNull<T>;
+                    let mut _5: *mut T;
                     let mut _6: *mut T;
-                    let mut _7: *mut T;
-                    scope 9 (inlined NonNull::<T>::as_ptr) {
+                    scope 8 (inlined NonNull::<T>::as_ptr) {
                     }
-                    scope 10 (inlined NonNull::<T>::as_ptr) {
+                    scope 9 (inlined NonNull::<T>::as_ptr) {
                     }
                 }
             }
             scope 5 (inlined std::ptr::const_ptr::<impl *const T>::addr) {
                 scope 6 (inlined std::ptr::const_ptr::<impl *const T>::cast::<()>) {
                 }
             }
-            scope 7 (inlined std::ptr::const_ptr::<impl *const *const T>::cast::<NonNull<T>>) {
-            }
         }
     }
 
     bb0: {
-        StorageLive(_9);
         StorageLive(_8);
-        StorageLive(_4);
+        StorageLive(_7);
+        StorageLive(_3);
         switchInt(const <T as std::mem::SizedTypeProperties>::IS_ZST) -> [0: bb1, otherwise: bb2];
     }
 
     bb1: {
-        StorageLive(_3);
         StorageLive(_2);
-        _2 = &raw const ((*_1).1: *const T);
-        _3 = copy _2 as *const std::ptr::NonNull<T> (PtrToPtr);
+        _2 = copy ((*_1).1: *const T);
+        _3 = move _2 as std::ptr::NonNull<T> (Transmute);
         StorageDead(_2);
-        _4 = copy (*_3);
-        StorageDead(_3);
-        StorageLive(_6);
         StorageLive(_5);
-        _5 = copy ((*_1).0: std::ptr::NonNull<T>);
-        _6 = copy _5 as *mut T (Transmute);
-        StorageDead(_5);
-        StorageLive(_7);
-        _7 = copy _4 as *mut T (Transmute);
-        _0 = Eq(move _6, move _7);
-        StorageDead(_7);
+        StorageLive(_4);
+        _4 = copy ((*_1).0: std::ptr::NonNull<T>);
+        _5 = copy _4 as *mut T (Transmute);
+        StorageDead(_4);
+        StorageLive(_6);
+        _6 = copy _3 as *mut T (Transmute);
+        _0 = Eq(move _5, move _6);
         StorageDead(_6);
+        StorageDead(_5);
         goto -> bb3;
     }
 
     bb2: {
-        _8 = copy ((*_1).1: *const T);
-        _9 = copy _8 as usize (Transmute);
-        _0 = Eq(copy _9, const 0_usize);
+        _7 = copy ((*_1).1: *const T);
+        _8 = copy _7 as usize (Transmute);
+        _0 = Eq(copy _8, const 0_usize);
         goto -> bb3;
     }
 
     bb3: {
-        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_7);
         StorageDead(_8);
-        StorageDead(_9);
         return;
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ macro_rules! if_zst {`
`30`	`30`	`$zst_body`
`31`	`31`	`} else {`
`32`	`32`	`// SAFETY: for non-ZSTs, the type invariant ensures it cannot be null`
`33`		`- let $end = unsafe { *(&raw const $this.end_or_len).cast::<NonNull<T>>() };`
	`33`	`+ let $end = unsafe { mem::transmute::<*const T, NonNull<T>>($this.end_or_len) };`
`34`	`34`	`$other_body`
`35`	`35`	`}`
`36`	`36`	`}};`