Auto merge of #95837 - scottmcm:ptr-offset-from-unsigned, r=oli-obk

Add `sub_ptr` on pointers (the `usize` version of `offset_from`)

We have `add`/`sub` which are the `usize` versions of `offset`, this adds the `usize` equivalent of `offset_from`.  Like how `.add(d)` replaced a whole bunch of `.offset(d as isize)`, you can see from the changes here that it's fairly common that code actually knows the order between the pointers and *wants* a `usize`, not an `isize`.

As a bonus, this can do `sub nuw`+`udiv exact`, rather than `sub`+`sdiv exact`, which can be optimized slightly better because it doesn't have to worry about negatives.  That's why the slice iterators weren't using `offset_from`, though I haven't updated that code in this PR because slices are so perf-critical that I'll do it as its own change.

This is an intrinsic, like `offset_from`, so that it can eventually be allowed in CTFE.  It also allows checking the extra safety condition -- see the test confirming that CTFE catches it if you pass the pointers in the wrong order.

Author: bors

alloc/src/lib.rs

@@ -127,6 +127,7 @@
 #![feature(pattern)]
 #![feature(ptr_internals)]
 #![feature(ptr_metadata)]
+#![feature(ptr_sub_ptr)]
 #![feature(receiver_trait)]
 #![feature(set_ptr_value)]
 #![feature(slice_group_by)]

alloc/src/slice.rs

@@ -1056,7 +1056,7 @@ where
         fn drop(&mut self) {
             // `T` is not a zero-sized type, and these are pointers into a slice's elements.
             unsafe {
-                let len = self.end.offset_from(self.start) as usize;
+                let len = self.end.sub_ptr(self.start);
                 ptr::copy_nonoverlapping(self.start, self.dest, len);
             }
         }

alloc/src/vec/drain.rs

@@ -163,7 +163,7 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
             // it from the original vec but also avoid creating a &mut to the front since that could
             // invalidate raw pointers to it which some unsafe code might rely on.
             let vec_ptr = vec.as_mut().as_mut_ptr();
-            let drop_offset = drop_ptr.offset_from(vec_ptr) as usize;
+            let drop_offset = drop_ptr.sub_ptr(vec_ptr);
             let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
             ptr::drop_in_place(to_drop);
         }

alloc/src/vec/in_place_collect.rs

@@ -250,7 +250,7 @@ where
         let sink =
             self.try_fold::<_, _, Result<_, !>>(sink, write_in_place_with_drop(end)).unwrap();
         // iteration succeeded, don't drop head
-        unsafe { ManuallyDrop::new(sink).dst.offset_from(dst_buf) as usize }
+        unsafe { ManuallyDrop::new(sink).dst.sub_ptr(dst_buf) }
     }
 }
 

alloc/src/vec/in_place_drop.rs

@@ -10,7 +10,7 @@ pub(super) struct InPlaceDrop<T> {
 
 impl<T> InPlaceDrop<T> {
     fn len(&self) -> usize {
-        unsafe { self.dst.offset_from(self.inner) as usize }
+        unsafe { self.dst.sub_ptr(self.inner) }
     }
 }
 

alloc/src/vec/into_iter.rs

@@ -169,7 +169,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         let exact = if mem::size_of::<T>() == 0 {
             self.end.addr().wrapping_sub(self.ptr.addr())
         } else {
-            unsafe { self.end.offset_from(self.ptr) as usize }
+            unsafe { self.end.sub_ptr(self.ptr) }
         };
         (exact, Some(exact))
     }

core/src/intrinsics.rs

@@ -1903,6 +1903,11 @@ extern "rust-intrinsic" {
     #[rustc_const_unstable(feature = "const_ptr_offset_from", issue = "92980")]
     pub fn ptr_offset_from<T>(ptr: *const T, base: *const T) -> isize;
 
+    /// See documentation of `<*const T>::sub_ptr` for details.
+    #[rustc_const_unstable(feature = "const_ptr_offset_from", issue = "92980")]
+    #[cfg(not(bootstrap))]
+    pub fn ptr_offset_from_unsigned<T>(ptr: *const T, base: *const T) -> usize;
+
     /// See documentation of `<*const T>::guaranteed_eq` for details.
     ///
     /// Note that, unlike most intrinsics, this is safe to call;
@@ -2385,3 +2390,11 @@ where
 {
     called_in_const.call_once(arg)
 }
+
+/// Bootstrap polyfill
+#[cfg(bootstrap)]
+pub const unsafe fn ptr_offset_from_unsigned<T>(ptr: *const T, base: *const T) -> usize {
+    // SAFETY: we have stricter preconditions than `ptr_offset_from`, so can
+    // call it, and its output has to be positive, so we can just cast.
+    unsafe { ptr_offset_from(ptr, base) as _ }
+}

core/src/lib.rs

@@ -126,6 +126,7 @@
 #![feature(const_option)]
 #![feature(const_option_ext)]
 #![feature(const_pin)]
+#![feature(const_ptr_sub_ptr)]
 #![feature(const_replace)]
 #![feature(const_ptr_as_ref)]
 #![feature(const_ptr_is_null)]

core/src/ptr/const_ptr.rs

@@ -611,6 +611,83 @@ impl<T: ?Sized> *const T {
         unsafe { intrinsics::ptr_offset_from(self, origin) }
     }
 
+    /// Calculates the distance between two pointers, *where it's known that
+    /// `self` is equal to or greater than `origin`*. The returned value is in
+    /// units of T: the distance in bytes is divided by `mem::size_of::<T>()`.
+    ///
+    /// This computes the same value that [`offset_from`](#method.offset_from)
+    /// would compute, but with the added precondition that that the offset is
+    /// guaranteed to be non-negative.  This method is equivalent to
+    /// `usize::from(self.offset_from(origin)).unwrap_unchecked()`,
+    /// but it provides slightly more information to the optimizer, which can
+    /// sometimes allow it to optimize slightly better with some backends.
+    ///
+    /// This method can be though of as recovering the `count` that was passed
+    /// to [`add`](#method.add) (or, with the parameters in the other order,
+    /// to [`sub`](#method.sub)).  The following are all equivalent, assuming
+    /// that their safety preconditions are met:
+    /// ```rust
+    /// # #![feature(ptr_sub_ptr)]
+    /// # unsafe fn blah(ptr: *const i32, origin: *const i32, count: usize) -> bool {
+    /// ptr.sub_ptr(origin) == count
+    /// # &&
+    /// origin.add(count) == ptr
+    /// # &&
+    /// ptr.sub(count) == origin
+    /// # }
+    /// ```
+    ///
+    /// # Safety
+    ///
+    /// - The distance between the pointers must be non-negative (`self >= origin`)
+    ///
+    /// - *All* the safety conditions of [`offset_from`](#method.offset_from)
+    ///   apply to this method as well; see it for the full details.
+    ///
+    /// Importantly, despite the return type of this method being able to represent
+    /// a larger offset, it's still *not permitted* to pass pointers which differ
+    /// by more than `isize::MAX` *bytes*.  As such, the result of this method will
+    /// always be less than or equal to `isize::MAX as usize`.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `T` is a Zero-Sized Type ("ZST").
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(ptr_sub_ptr)]
+    ///
+    /// let a = [0; 5];
+    /// let ptr1: *const i32 = &a[1];
+    /// let ptr2: *const i32 = &a[3];
+    /// unsafe {
+    ///     assert_eq!(ptr2.sub_ptr(ptr1), 2);
+    ///     assert_eq!(ptr1.add(2), ptr2);
+    ///     assert_eq!(ptr2.sub(2), ptr1);
+    ///     assert_eq!(ptr2.sub_ptr(ptr2), 0);
+    /// }
+    ///
+    /// // This would be incorrect, as the pointers are not correctly ordered:
+    /// // ptr1.offset_from(ptr2)
+    /// ```
+    #[unstable(feature = "ptr_sub_ptr", issue = "95892")]
+    #[rustc_const_unstable(feature = "const_ptr_sub_ptr", issue = "95892")]
+    #[inline]
+    pub const unsafe fn sub_ptr(self, origin: *const T) -> usize
+    where
+        T: Sized,
+    {
+        // SAFETY: The comparison has no side-effects, and the intrinsic
+        // does this check internally in the CTFE implementation.
+        unsafe { assert_unsafe_precondition!(self >= origin) };
+
+        let pointee_size = mem::size_of::<T>();
+        assert!(0 < pointee_size && pointee_size <= isize::MAX as usize);
+        // SAFETY: the caller must uphold the safety contract for `ptr_offset_from_unsigned`.
+        unsafe { intrinsics::ptr_offset_from_unsigned(self, origin) }
+    }
+
     /// Returns whether two pointers are guaranteed to be equal.
     ///
     /// At runtime this function behaves like `self == other`.

core/src/ptr/mut_ptr.rs

@@ -787,6 +787,78 @@ impl<T: ?Sized> *mut T {
         unsafe { (self as *const T).offset_from(origin) }
     }
 
+    /// Calculates the distance between two pointers, *where it's known that
+    /// `self` is equal to or greater than `origin`*. The returned value is in
+    /// units of T: the distance in bytes is divided by `mem::size_of::<T>()`.
+    ///
+    /// This computes the same value that [`offset_from`](#method.offset_from)
+    /// would compute, but with the added precondition that that the offset is
+    /// guaranteed to be non-negative.  This method is equivalent to
+    /// `usize::from(self.offset_from(origin)).unwrap_unchecked()`,
+    /// but it provides slightly more information to the optimizer, which can
+    /// sometimes allow it to optimize slightly better with some backends.
+    ///
+    /// This method can be though of as recovering the `count` that was passed
+    /// to [`add`](#method.add) (or, with the parameters in the other order,
+    /// to [`sub`](#method.sub)).  The following are all equivalent, assuming
+    /// that their safety preconditions are met:
+    /// ```rust
+    /// # #![feature(ptr_sub_ptr)]
+    /// # unsafe fn blah(ptr: *mut i32, origin: *mut i32, count: usize) -> bool {
+    /// ptr.sub_ptr(origin) == count
+    /// # &&
+    /// origin.add(count) == ptr
+    /// # &&
+    /// ptr.sub(count) == origin
+    /// # }
+    /// ```
+    ///
+    /// # Safety
+    ///
+    /// - The distance between the pointers must be non-negative (`self >= origin`)
+    ///
+    /// - *All* the safety conditions of [`offset_from`](#method.offset_from)
+    ///   apply to this method as well; see it for the full details.
+    ///
+    /// Importantly, despite the return type of this method being able to represent
+    /// a larger offset, it's still *not permitted* to pass pointers which differ
+    /// by more than `isize::MAX` *bytes*.  As such, the result of this method will
+    /// always be less than or equal to `isize::MAX as usize`.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `T` is a Zero-Sized Type ("ZST").
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(ptr_sub_ptr)]
+    ///
+    /// let mut a = [0; 5];
+    /// let p: *mut i32 = a.as_mut_ptr();
+    /// unsafe {
+    ///     let ptr1: *mut i32 = p.add(1);
+    ///     let ptr2: *mut i32 = p.add(3);
+    ///
+    ///     assert_eq!(ptr2.sub_ptr(ptr1), 2);
+    ///     assert_eq!(ptr1.add(2), ptr2);
+    ///     assert_eq!(ptr2.sub(2), ptr1);
+    ///     assert_eq!(ptr2.sub_ptr(ptr2), 0);
+    /// }
+    ///
+    /// // This would be incorrect, as the pointers are not correctly ordered:
+    /// // ptr1.offset_from(ptr2)
+    #[unstable(feature = "ptr_sub_ptr", issue = "95892")]
+    #[rustc_const_unstable(feature = "const_ptr_sub_ptr", issue = "95892")]
+    #[inline]
+    pub const unsafe fn sub_ptr(self, origin: *const T) -> usize
+    where
+        T: Sized,
+    {
+        // SAFETY: the caller must uphold the safety contract for `sub_ptr`.
+        unsafe { (self as *const T).sub_ptr(origin) }
+    }
+
     /// Calculates the offset from a pointer (convenience for `.offset(count as isize)`).
     ///
     /// `count` is in units of T; e.g., a `count` of 3 represents a pointer