Bless tagged pointers (comply to strict provenance)

5 files changed, 85 insertions, 67 deletions

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index e373bd18402..0a3f5f32d16 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -29,6 +29,7 @@
 #![feature(get_mut_unchecked)]
 #![feature(lint_reasons)]
 #![feature(unwrap_infallible)]
+#![feature(strict_provenance)]
 #![allow(rustc::default_hash_types)]
 #![allow(rustc::potential_query_instability)]
 #![deny(rustc::untranslatable_diagnostic)]
diff --git a/compiler/rustc_data_structures/src/tagged_ptr.rs b/compiler/rustc_data_structures/src/tagged_ptr.rs
index e69a11dae0a..8ad2b2a41fd 100644
--- a/compiler/rustc_data_structures/src/tagged_ptr.rs
+++ b/compiler/rustc_data_structures/src/tagged_ptr.rs
@@ -15,6 +15,7 @@
 
 use std::mem::{self, ManuallyDrop};
 use std::ops::Deref;
+use std::ptr::NonNull;
 use std::rc::Rc;
 use std::sync::Arc;
 
@@ -29,21 +30,24 @@ pub use drop::TaggedPtr;
 ///
 /// # Safety
 ///
-/// The usize returned from `into_usize` must be a valid, dereferenceable,
-/// pointer to [`<Self as Deref>::Target`]. Note that pointers to
-/// [`Self::Target`] must be thin, even though [`Self::Target`] may not be
-/// `Sized`.
+/// The pointer returned from [`into_ptr`] must be a [valid], pointer to
+/// [`<Self as Deref>::Target`]. Note that pointers to [`Self::Target`] must be
+/// thin, even though [`Self::Target`] may not be `Sized`.
 ///
-/// Note that the returned pointer from `into_usize` should be castable to `&mut
-/// <Self as Deref>::Target` if `Self: DerefMut`.
+/// Note that if `Self` implements [`DerefMut`] the pointer returned from
+/// [`into_ptr`] must be valid for writes (and thus calling [`NonNull::as_mut`]
+/// on it must be safe).
 ///
-/// The BITS constant must be correct. At least `BITS` bits, least-significant,
-/// must be zero on all returned pointers from `into_usize`.
+/// The `BITS` constant must be correct. At least `BITS` bits, least-significant,
+/// must be zero on all pointers returned from [`into_ptr`].
 ///
 /// For example, if the alignment of [`Self::Target`] is 2, then `BITS` should be 1.
 ///
+/// [`into_ptr`]: Pointer::into_ptr
+/// [valid]: std::ptr#safety
 /// [`<Self as Deref>::Target`]: Deref::Target
 /// [`Self::Target`]: Deref::Target
+/// [`DerefMut`]: std::ops::DerefMut
 pub unsafe trait Pointer: Deref {
     /// Number of unused (always zero) **least significant bits** in this
     /// pointer, usually related to the pointees alignment.
@@ -63,7 +67,7 @@ pub unsafe trait Pointer: Deref {
     /// [`Self::Target`]: Deref::Target
     const BITS: usize;
 
-    fn into_usize(self) -> usize;
+    fn into_ptr(self) -> NonNull<Self::Target>;
 
     /// # Safety
     ///
@@ -71,7 +75,7 @@ pub unsafe trait Pointer: Deref {
     ///
     /// This acts as `ptr::read` semantically, it should not be called more than
     /// once on non-`Copy` `Pointer`s.
-    unsafe fn from_usize(ptr: usize) -> Self;
+    unsafe fn from_ptr(ptr: NonNull<Self::Target>) -> Self;
 
     /// This provides a reference to the `Pointer` itself, rather than the
     /// `Deref::Target`. It is used for cases where we want to call methods that
@@ -81,7 +85,7 @@ pub unsafe trait Pointer: Deref {
     /// # Safety
     ///
     /// The passed `ptr` must be returned from `into_usize`.
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R;
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: NonNull<Self::Target>, f: F) -> R;
 }
 
 /// This describes tags that the `TaggedPtr` struct can hold.
@@ -106,17 +110,18 @@ unsafe impl<T> Pointer for Box<T> {
     const BITS: usize = bits_for::<Self::Target>();
 
     #[inline]
-    fn into_usize(self) -> usize {
-        Box::into_raw(self) as usize
+    fn into_ptr(self) -> NonNull<T> {
+        // Safety: pointers from `Box::into_raw` are valid & non-null
+        unsafe { NonNull::new_unchecked(Box::into_raw(self)) }
     }
 
     #[inline]
-    unsafe fn from_usize(ptr: usize) -> Self {
-        Box::from_raw(ptr as *mut T)
+    unsafe fn from_ptr(ptr: NonNull<T>) -> Self {
+        Box::from_raw(ptr.as_ptr())
     }
 
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
-        let raw = ManuallyDrop::new(Self::from_usize(ptr));
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: NonNull<T>, f: F) -> R {
+        let raw = ManuallyDrop::new(Self::from_ptr(ptr));
         f(&raw)
     }
 }
@@ -125,17 +130,17 @@ unsafe impl<T> Pointer for Rc<T> {
     const BITS: usize = bits_for::<Self::Target>();
 
     #[inline]
-    fn into_usize(self) -> usize {
-        Rc::into_raw(self) as usize
+    fn into_ptr(self) -> NonNull<T> {
+        unsafe { NonNull::new_unchecked(Rc::into_raw(self).cast_mut()) }
     }
 
     #[inline]
-    unsafe fn from_usize(ptr: usize) -> Self {
-        Rc::from_raw(ptr as *const T)
+    unsafe fn from_ptr(ptr: NonNull<T>) -> Self {
+        Rc::from_raw(ptr.as_ptr())
     }
 
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
-        let raw = ManuallyDrop::new(Self::from_usize(ptr));
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: NonNull<T>, f: F) -> R {
+        let raw = ManuallyDrop::new(Self::from_ptr(ptr));
         f(&raw)
     }
 }
@@ -144,17 +149,17 @@ unsafe impl<T> Pointer for Arc<T> {
     const BITS: usize = bits_for::<Self::Target>();
 
     #[inline]
-    fn into_usize(self) -> usize {
-        Arc::into_raw(self) as usize
+    fn into_ptr(self) -> NonNull<T> {
+        unsafe { NonNull::new_unchecked(Arc::into_raw(self).cast_mut()) }
     }
 
     #[inline]
-    unsafe fn from_usize(ptr: usize) -> Self {
-        Arc::from_raw(ptr as *const T)
+    unsafe fn from_ptr(ptr: NonNull<T>) -> Self {
+        Arc::from_raw(ptr.as_ptr())
     }
 
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
-        let raw = ManuallyDrop::new(Self::from_usize(ptr));
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: NonNull<T>, f: F) -> R {
+        let raw = ManuallyDrop::new(Self::from_ptr(ptr));
         f(&raw)
     }
 }
@@ -163,32 +168,35 @@ unsafe impl<'a, T: 'a> Pointer for &'a T {
     const BITS: usize = bits_for::<Self::Target>();
 
     #[inline]
-    fn into_usize(self) -> usize {
-        self as *const T as usize
+    fn into_ptr(self) -> NonNull<T> {
+        NonNull::from(self)
     }
 
     #[inline]
-    unsafe fn from_usize(ptr: usize) -> Self {
-        &*(ptr as *const T)
+    unsafe fn from_ptr(ptr: NonNull<T>) -> Self {
+        ptr.as_ref()
     }
 
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
-        f(&*(&ptr as *const usize as *const Self))
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: NonNull<T>, f: F) -> R {
+        f(&ptr.as_ref())
     }
 }
 
 unsafe impl<'a, T: 'a> Pointer for &'a mut T {
     const BITS: usize = bits_for::<Self::Target>();
+
     #[inline]
-    fn into_usize(self) -> usize {
-        self as *mut T as usize
+    fn into_ptr(self) -> NonNull<T> {
+        NonNull::from(self)
     }
+
     #[inline]
-    unsafe fn from_usize(ptr: usize) -> Self {
-        &mut *(ptr as *mut T)
+    unsafe fn from_ptr(mut ptr: NonNull<T>) -> Self {
+        ptr.as_mut()
     }
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
-        f(&*(&ptr as *const usize as *const Self))
+
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(mut ptr: NonNull<T>, f: F) -> R {
+        f(&ptr.as_mut())
     }
 }
 
diff --git a/compiler/rustc_data_structures/src/tagged_ptr/copy.rs b/compiler/rustc_data_structures/src/tagged_ptr/copy.rs
index d0c2b914584..958656f9a02 100644
--- a/compiler/rustc_data_structures/src/tagged_ptr/copy.rs
+++ b/compiler/rustc_data_structures/src/tagged_ptr/copy.rs
@@ -4,6 +4,7 @@ use std::fmt;
 use std::marker::PhantomData;
 use std::num::NonZeroUsize;
 use std::ops::{Deref, DerefMut};
+use std::ptr::NonNull;
 
 /// A `Copy` TaggedPtr.
 ///
@@ -18,7 +19,7 @@ where
     P: Pointer,
     T: Tag,
 {
-    packed: NonZeroUsize,
+    packed: NonNull<P::Target>,
     data: PhantomData<(P, T)>,
 }
 
@@ -53,26 +54,36 @@ where
     const ASSERTION: () = {
         assert!(T::BITS <= P::BITS);
         // Used for the transmute_copy's below
+        // TODO(waffle): do we need this assert anymore?
         assert!(std::mem::size_of::<&P::Target>() == std::mem::size_of::<usize>());
     };
 
     pub fn new(pointer: P, tag: T) -> Self {
         // Trigger assert!
         let () = Self::ASSERTION;
+
         let packed_tag = tag.into_usize() << Self::TAG_BIT_SHIFT;
 
         Self {
-            // SAFETY: We know that the pointer is non-null, as it must be
-            // dereferenceable per `Pointer` safety contract.
-            packed: unsafe {
-                NonZeroUsize::new_unchecked((P::into_usize(pointer) >> T::BITS) | packed_tag)
-            },
+            packed: P::into_ptr(pointer).map_addr(|addr| {
+                // SAFETY:
+                // - The pointer is `NonNull` => it's address is `NonZeroUsize`
+                // - `P::BITS` least significant bits are always zero (`Pointer` contract)
+                // - `T::BITS <= P::BITS` (from `Self::ASSERTION`)
+                //
+                // Thus `addr >> T::BITS` is guaranteed to be non-zero.
+                //
+                // `{non_zero} | packed_tag` can't make the value zero.
+
+                let packed = (addr.get() >> T::BITS) | packed_tag;
+                unsafe { NonZeroUsize::new_unchecked(packed) }
+            }),
             data: PhantomData,
         }
     }
 
-    pub(super) fn pointer_raw(&self) -> usize {
-        self.packed.get() << T::BITS
+    pub(super) fn pointer_raw(&self) -> NonNull<P::Target> {
+        self.packed.map_addr(|addr| unsafe { NonZeroUsize::new_unchecked(addr.get() << T::BITS) })
     }
 
     pub fn pointer(self) -> P
@@ -83,12 +94,12 @@ where
         //
         // Note that this isn't going to double-drop or anything because we have
         // P: Copy
-        unsafe { P::from_usize(self.pointer_raw()) }
+        unsafe { P::from_ptr(self.pointer_raw()) }
     }
 
     pub fn pointer_ref(&self) -> &P::Target {
         // SAFETY: pointer_raw returns the original pointer
-        unsafe { std::mem::transmute_copy(&self.pointer_raw()) }
+        unsafe { self.pointer_raw().as_ref() }
     }
 
     pub fn pointer_mut(&mut self) -> &mut P::Target
@@ -96,22 +107,22 @@ where
         P: DerefMut,
     {
         // SAFETY: pointer_raw returns the original pointer
-        unsafe { std::mem::transmute_copy(&self.pointer_raw()) }
+        unsafe { self.pointer_raw().as_mut() }
     }
 
     #[inline]
     pub fn tag(&self) -> T {
-        unsafe { T::from_usize(self.packed.get() >> Self::TAG_BIT_SHIFT) }
+        unsafe { T::from_usize(self.packed.addr().get() >> Self::TAG_BIT_SHIFT) }
     }
 
     #[inline]
     pub fn set_tag(&mut self, tag: T) {
-        let mut packed = self.packed.get();
+        // TODO: refactor packing into a function and reuse it here
         let new_tag = T::into_usize(tag) << Self::TAG_BIT_SHIFT;
         let tag_mask = (1 << T::BITS) - 1;
-        packed &= !(tag_mask << Self::TAG_BIT_SHIFT);
-        packed |= new_tag;
-        self.packed = unsafe { NonZeroUsize::new_unchecked(packed) };
+        self.packed = self.packed.map_addr(|addr| unsafe {
+            NonZeroUsize::new_unchecked(addr.get() & !(tag_mask << Self::TAG_BIT_SHIFT) | new_tag)
+        });
     }
 }
 
diff --git a/compiler/rustc_data_structures/src/tagged_ptr/drop.rs b/compiler/rustc_data_structures/src/tagged_ptr/drop.rs
index b0315c93d93..c734dadefe9 100644
--- a/compiler/rustc_data_structures/src/tagged_ptr/drop.rs
+++ b/compiler/rustc_data_structures/src/tagged_ptr/drop.rs
@@ -76,7 +76,7 @@ where
     fn drop(&mut self) {
         // No need to drop the tag, as it's Copy
         unsafe {
-            drop(P::from_usize(self.raw.pointer_raw()));
+            drop(P::from_ptr(self.raw.pointer_raw()));
         }
     }
 }
diff --git a/compiler/rustc_middle/src/ty/list.rs b/compiler/rustc_middle/src/ty/list.rs
index 4526487cf1d..01b5e82cd22 100644
--- a/compiler/rustc_middle/src/ty/list.rs
+++ b/compiler/rustc_middle/src/ty/list.rs
@@ -8,7 +8,7 @@ use std::hash::{Hash, Hasher};
 use std::iter;
 use std::mem;
 use std::ops::Deref;
-use std::ptr;
+use std::ptr::{self, NonNull};
 use std::slice;
 
 /// `List<T>` is a bit like `&[T]`, but with some critical differences.
@@ -203,18 +203,16 @@ unsafe impl<'a, T: 'a> rustc_data_structures::tagged_ptr::Pointer for &'a List<T
     const BITS: usize = bits_for::<usize>();
 
     #[inline]
-    fn into_usize(self) -> usize {
-        self as *const List<T> as usize
+    fn into_ptr(self) -> NonNull<List<T>> {
+        NonNull::from(self)
     }
 
     #[inline]
-    unsafe fn from_usize(ptr: usize) -> &'a List<T> {
-        &*(ptr as *const List<T>)
+    unsafe fn from_ptr(ptr: NonNull<List<T>>) -> &'a List<T> {
+        ptr.as_ref()
     }
 
-    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
-        // `Self` is `&'a List<T>` which impls `Copy`, so this is fine.
-        let ptr = Self::from_usize(ptr);
-        f(&ptr)
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: NonNull<List<T>>, f: F) -> R {
+        f(&ptr.as_ref())
     }
 }