Auto merge of rust-lang#136771 - scottmcm:poke-slice-iter-next, r=<try>

Simplify `slice::Iter::next` enough that it inlines

Inspired by this zulip conversation: <https://rust-lang.zulipchat.com/#narrow/channel/189540-t-compiler.2Fwg-mir-opt/topic/Feedback.20on.20a.20MIR.20optimization.20idea/near/498579990>

Draft for now because it needs rust-lang#136735 to get the codegen tests to pass.

Author: bors

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(

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

library/core/src/slice/iter.rs

@@ -11,7 +11,7 @@ use crate::iter::{
 use crate::marker::PhantomData;
 use crate::mem::{self, SizedTypeProperties};
 use crate::num::NonZero;
-use crate::ptr::{NonNull, without_provenance, without_provenance_mut};
+use crate::ptr::{NonNull, null, without_provenance, without_provenance_mut};
 use crate::{cmp, fmt};
 
 #[stable(feature = "boxed_slice_into_iter", since = "1.80.0")]

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
         }
     }};
@@ -154,16 +154,26 @@ macro_rules! iterator {
 
             #[inline]
             fn next(&mut self) -> Option<$elem> {
-                // could be implemented with slices, but this avoids bounds checks
+                // intentionally not using the helpers because this is
+                // one of the most mono'd things in the library.
 
-                // SAFETY: The call to `next_unchecked` is
-                // safe since we check if the iterator is empty first.
+                let ptr = self.ptr;
+                let end_or_len = self.end_or_len;
+                // SAFETY: Type invariants.
                 unsafe {
-                    if is_empty!(self) {
-                        None
+                    if T::IS_ZST {
+                        let byte_end = end_or_len as *const u8;
+                        if byte_end == null() {
+                            return None;
+                        }
+                        self.end_or_len = byte_end.wrapping_sub(1) as _;
                     } else {
-                        Some(self.next_unchecked())
+                        if ptr == crate::intrinsics::transmute::<*const T, NonNull<T>>(end_or_len) {
+                            return None;
+                        }
+                        self.ptr = ptr.add(1);
                     }
+                    crate::intrinsics::transmute::<NonNull<T>, Option<$elem>>(ptr)
                 }
             }
 

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)
+}

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

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

@@ -1,4 +1,5 @@
 //@ compile-flags: -O
+//@ 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
 }

tests/codegen/slice-iter-nonnull.rs

@@ -14,11 +14,11 @@
 // CHECK-LABEL: @slice_iter_next(
 #[no_mangle]
 pub fn slice_iter_next<'a>(it: &mut std::slice::Iter<'a, u32>) -> Option<&'a u32> {
-    // CHECK: %[[ENDP:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %it, {{i32 4|i64 8}}
-    // CHECK: %[[END:.+]] = load ptr, ptr %[[ENDP]]
+    // CHECK: %[[START:.+]] = load ptr, ptr %it,
     // CHECK-SAME: !nonnull
     // CHECK-SAME: !noundef
-    // CHECK: %[[START:.+]] = load ptr, ptr %it,
+    // CHECK: %[[ENDP:.+]] = getelementptr inbounds{{( nuw)?}} i8, ptr %it, {{i32 4|i64 8}}
+    // CHECK: %[[END:.+]] = load ptr, ptr %[[ENDP]]
     // CHECK-SAME: !nonnull
     // CHECK-SAME: !noundef
     // CHECK: icmp eq ptr %[[START]], %[[END]]