Rollup merge of #95556 - declanvk:nonnull-provenance, r=dtolnay

Implement provenance preserving methods on NonNull

### Description
 Add the `addr`, `with_addr`, `map_addr` methods to the `NonNull` type, and map the address type to `NonZeroUsize`.

 ### Motivation
 The `NonNull` type is useful for implementing pointer types which have  the 0-niche. It is currently possible to implement these provenance  preserving functions by calling `NonNull::as_ptr` and `new_unchecked`. The adding these methods makes it more ergonomic.

 ### Testing
 Added a unit test of a non-null tagged pointer type. This is based on some real code I have elsewhere, that currently routes the pointer through a `NonZeroUsize` and back out to produce a usable pointer. I wanted to produce an ideal version of the same tagged pointer struct that preserved pointer provenance.

### Related

Extension of APIs proposed in #95228 . I can also split this out into a separate tracking issue if that is better (though I may need some pointers on how to do that).

library/core/src/ptr/non_null.rs

@@ -4,6 +4,7 @@
 use crate::hash;
 use crate::marker::Unsize;
 use crate::mem::{self, MaybeUninit};
+use crate::num::NonZeroUsize;
 use crate::ops::{CoerceUnsized, DispatchFromDyn};
 use crate::ptr::Unique;
 use crate::slice::{self, SliceIndex};
@@ -253,6 +254,53 @@ pub const fn to_raw_parts(self) -> (NonNull<()>, <T as super::Pointee>::Metadata
         (self.cast(), super::metadata(self.as_ptr()))
     }
 
+    /// Gets the "address" portion of the pointer.
+    ///
+    /// This API and its claimed semantics are part of the Strict Provenance experiment,
+    /// see the [module documentation][crate::ptr] for details.
+    #[must_use]
+    #[inline]
+    #[unstable(feature = "strict_provenance", issue = "95228")]
+    pub fn addr(self) -> NonZeroUsize
+    where
+        T: Sized,
+    {
+        // SAFETY: The pointer is guaranteed by the type to be non-null,
+        // meaning that the address will be non-zero.
+        unsafe { NonZeroUsize::new_unchecked(self.pointer.addr()) }
+    }
+
+    /// Creates a new pointer with the given address.
+    ///
+    /// This API and its claimed semantics are part of the Strict Provenance experiment,
+    /// see the [module documentation][crate::ptr] for details.
+    #[must_use]
+    #[inline]
+    #[unstable(feature = "strict_provenance", issue = "95228")]
+    pub fn with_addr(self, addr: NonZeroUsize) -> Self
+    where
+        T: Sized,
+    {
+        // SAFETY: The result of `ptr::from::with_addr` is non-null because `addr` is guaranteed to be non-zero.
+        unsafe { NonNull::new_unchecked(self.pointer.with_addr(addr.get()) as *mut _) }
+    }
+
+    /// Creates a new pointer by mapping `self`'s address to a new one.
+    ///
+    /// This is a convenience for [`with_addr`][Self::with_addr], see that method for details.
+    ///
+    /// This API and its claimed semantics are part of the Strict Provenance experiment,
+    /// see the [module documentation][crate::ptr] for details.
+    #[must_use]
+    #[inline]
+    #[unstable(feature = "strict_provenance", issue = "95228")]
+    pub fn map_addr(self, f: impl FnOnce(NonZeroUsize) -> NonZeroUsize) -> Self
+    where
+        T: Sized,
+    {
+        self.with_addr(f(self.addr()))
+    }
+
     /// Acquires the underlying `*mut` pointer.
     ///
     /// # Examples

library/core/tests/lib.rs

@@ -86,6 +86,7 @@
 #![feature(int_roundings)]
 #![feature(slice_group_by)]
 #![feature(split_array)]
+#![feature(strict_provenance)]
 #![feature(trusted_random_access)]
 #![feature(unsize)]
 #![feature(unzip_option)]

library/core/tests/ptr.rs

@@ -1,4 +1,5 @@
 use core::cell::RefCell;
+use core::num::NonZeroUsize;
 use core::ptr;
 use core::ptr::*;
 use std::fmt::{Debug, Display};
@@ -691,3 +692,80 @@ fn drop(&mut self) {
         }
     }
 }
+
+#[test]
+fn nonnull_tagged_pointer_with_provenance() {
+    let raw_pointer = Box::into_raw(Box::new(10));
+
+    let mut p = TaggedPointer::new(raw_pointer).unwrap();
+    assert_eq!(p.tag(), 0);
+
+    p.set_tag(1);
+    assert_eq!(p.tag(), 1);
+    assert_eq!(unsafe { *p.pointer().as_ptr() }, 10);
+
+    p.set_tag(3);
+    assert_eq!(p.tag(), 3);
+    assert_eq!(unsafe { *p.pointer().as_ptr() }, 10);
+
+    unsafe { Box::from_raw(p.pointer().as_ptr()) };
+
+    /// A non-null pointer type which carries several bits of metadata and maintains provenance.
+    #[repr(transparent)]
+    pub struct TaggedPointer<T>(NonNull<T>);
+
+    impl<T> Clone for TaggedPointer<T> {
+        fn clone(&self) -> Self {
+            Self(self.0)
+        }
+    }
+
+    impl<T> Copy for TaggedPointer<T> {}
+
+    impl<T> TaggedPointer<T> {
+        /// The ABI-required minimum alignment of the `P` type.
+        pub const ALIGNMENT: usize = core::mem::align_of::<T>();
+        /// A mask for data-carrying bits of the address.
+        pub const DATA_MASK: usize = !Self::ADDRESS_MASK;
+        /// Number of available bits of storage in the address.
+        pub const NUM_BITS: u32 = Self::ALIGNMENT.trailing_zeros();
+        /// A mask for the non-data-carrying bits of the address.
+        pub const ADDRESS_MASK: usize = usize::MAX << Self::NUM_BITS;
+
+        /// Create a new tagged pointer from a possibly null pointer.
+        pub fn new(pointer: *mut T) -> Option<TaggedPointer<T>> {
+            Some(TaggedPointer(NonNull::new(pointer)?))
+        }
+
+        /// Consume this tagged pointer and produce a raw mutable pointer to the
+        /// memory location.
+        pub fn pointer(self) -> NonNull<T> {
+            // SAFETY: The `addr` guaranteed to have bits set in the Self::ADDRESS_MASK, so the result will be non-null.
+            self.0.map_addr(|addr| unsafe {
+                NonZeroUsize::new_unchecked(addr.get() & Self::ADDRESS_MASK)
+            })
+        }
+
+        /// Consume this tagged pointer and produce the data it carries.
+        pub fn tag(&self) -> usize {
+            self.0.addr().get() & Self::DATA_MASK
+        }
+
+        /// Update the data this tagged pointer carries to a new value.
+        pub fn set_tag(&mut self, data: usize) {
+            assert_eq!(
+                data & Self::ADDRESS_MASK,
+                0,
+                "cannot set more data beyond the lowest NUM_BITS"
+            );
+            let data = data & Self::DATA_MASK;
+
+            // SAFETY: This value will always be non-zero because the upper bits (from
+            // ADDRESS_MASK) will always be non-zero. This a property of the type and its
+            // construction.
+            self.0 = self.0.map_addr(|addr| unsafe {
+                NonZeroUsize::new_unchecked((addr.get() & Self::ADDRESS_MASK) | data)
+            })
+        }
+    }
+}