1 files changed, 232 insertions, 89 deletions

diff --git a/compiler/rustc_data_structures/src/tagged_ptr/copy.rs b/compiler/rustc_data_structures/src/tagged_ptr/copy.rs
index e1d3e0bd3..e893a2c78 100644
--- a/compiler/rustc_data_structures/src/tagged_ptr/copy.rs
+++ b/compiler/rustc_data_structures/src/tagged_ptr/copy.rs
@@ -1,78 +1,94 @@
 use super::{Pointer, Tag};
 use crate::stable_hasher::{HashStable, StableHasher};
 use std::fmt;
+use std::hash::{Hash, Hasher};
 use std::marker::PhantomData;
+use std::mem::ManuallyDrop;
 use std::num::NonZeroUsize;
+use std::ops::{Deref, DerefMut};
+use std::ptr::NonNull;
 
-/// A `Copy` TaggedPtr.
+/// A [`Copy`] tagged pointer.
 ///
-/// You should use this instead of the `TaggedPtr` type in all cases where
-/// `P: Copy`.
+/// This is essentially `{ pointer: P, tag: T }` packed in a single pointer.
+///
+/// You should use this instead of the [`TaggedPtr`] type in all cases where
+/// `P` implements [`Copy`].
 ///
 /// If `COMPARE_PACKED` is true, then the pointers will be compared and hashed without
-/// unpacking. Otherwise we don't implement PartialEq/Eq/Hash; if you want that,
-/// wrap the TaggedPtr.
+/// unpacking. Otherwise we don't implement [`PartialEq`], [`Eq`] and [`Hash`];
+/// if you want that, wrap the [`CopyTaggedPtr`].
+///
+/// [`TaggedPtr`]: crate::tagged_ptr::TaggedPtr
 pub struct CopyTaggedPtr<P, T, const COMPARE_PACKED: bool>
 where
     P: Pointer,
     T: Tag,
 {
-    packed: NonZeroUsize,
-    data: PhantomData<(P, T)>,
-}
-
-impl<P, T, const COMPARE_PACKED: bool> Copy for CopyTaggedPtr<P, T, COMPARE_PACKED>
-where
-    P: Pointer,
-    T: Tag,
-    P: Copy,
-{
-}
-
-impl<P, T, const COMPARE_PACKED: bool> Clone for CopyTaggedPtr<P, T, COMPARE_PACKED>
-where
-    P: Pointer,
-    T: Tag,
-    P: Copy,
-{
-    fn clone(&self) -> Self {
-        *self
-    }
+    /// This is semantically a pair of `pointer: P` and `tag: T` fields,
+    /// however we pack them in a single pointer, to save space.
+    ///
+    /// We pack the tag into the **most**-significant bits of the pointer to
+    /// ease retrieval of the value. A left shift is a multiplication and
+    /// those are embeddable in instruction encoding, for example:
+    ///
+    /// ```asm
+    /// // (<https://godbolt.org/z/jqcYPWEr3>)
+    /// example::shift_read3:
+    ///     mov     eax, dword ptr [8*rdi]
+    ///     ret
+    ///
+    /// example::mask_read3:
+    ///     and     rdi, -8
+    ///     mov     eax, dword ptr [rdi]
+    ///     ret
+    /// ```
+    ///
+    /// This is ASM outputted by rustc for reads of values behind tagged
+    /// pointers for different approaches of tagging:
+    /// - `shift_read3` uses `<< 3` (the tag is in the most-significant bits)
+    /// - `mask_read3` uses `& !0b111` (the tag is in the least-significant bits)
+    ///
+    /// The shift approach thus produces less instructions and is likely faster
+    /// (see <https://godbolt.org/z/Y913sMdWb>).
+    ///
+    /// Encoding diagram:
+    /// ```text
+    /// [ packed.addr                     ]
+    /// [ tag ] [ pointer.addr >> T::BITS ] <-- usize::BITS - T::BITS bits
+    ///    ^
+    ///    |
+    /// T::BITS bits
+    /// ```
+    ///
+    /// The tag can be retrieved by `packed.addr() >> T::BITS` and the pointer
+    /// can be retrieved by `packed.map_addr(|addr| addr << T::BITS)`.
+    packed: NonNull<P::Target>,
+    tag_ghost: PhantomData<T>,
 }
 
-// We pack the tag into the *upper* bits of the pointer to ease retrieval of the
-// value; a left shift is a multiplication and those are embeddable in
-// instruction encoding.
-impl<P, T, const COMPARE_PACKED: bool> CopyTaggedPtr<P, T, COMPARE_PACKED>
+// Note that even though `CopyTaggedPtr` is only really expected to work with
+// `P: Copy`, can't add `P: Copy` bound, because `CopyTaggedPtr` is used in the
+// `TaggedPtr`'s implementation.
+impl<P, T, const CP: bool> CopyTaggedPtr<P, T, CP>
 where
     P: Pointer,
     T: Tag,
 {
-    const TAG_BIT_SHIFT: usize = usize::BITS as usize - T::BITS;
-    const ASSERTION: () = {
-        assert!(T::BITS <= P::BITS);
-        // Used for the transmute_copy's below
-        assert!(std::mem::size_of::<&P::Target>() == std::mem::size_of::<usize>());
-    };
-
+    /// Tags `pointer` with `tag`.
+    ///
+    /// Note that this leaks `pointer`: it won't be dropped when
+    /// `CopyTaggedPtr` is dropped. If you have a pointer with a significant
+    /// drop, use [`TaggedPtr`] instead.
+    ///
+    /// [`TaggedPtr`]: crate::tagged_ptr::TaggedPtr
+    #[inline]
     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)
-            },
-            data: PhantomData,
-        }
+        Self { packed: Self::pack(P::into_ptr(pointer), tag), tag_ghost: PhantomData }
     }
 
-    pub(super) fn pointer_raw(&self) -> usize {
-        self.packed.get() << T::BITS
-    }
+    /// Retrieves the pointer.
+    #[inline]
     pub fn pointer(self) -> P
     where
         P: Copy,
@@ -81,66 +97,143 @@ where
         //
         // Note that this isn't going to double-drop or anything because we have
         // P: Copy
-        unsafe { P::from_usize(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()) }
-    }
-    pub fn pointer_mut(&mut self) -> &mut P::Target
-    where
-        P: std::ops::DerefMut,
-    {
-        // SAFETY: pointer_raw returns the original pointer
-        unsafe { std::mem::transmute_copy(&self.pointer_raw()) }
+        unsafe { P::from_ptr(self.pointer_raw()) }
     }
+
+    /// Retrieves the tag.
     #[inline]
     pub fn tag(&self) -> T {
-        unsafe { T::from_usize(self.packed.get() >> Self::TAG_BIT_SHIFT) }
+        // Unpack the tag, according to the `self.packed` encoding scheme
+        let tag = self.packed.addr().get() >> Self::TAG_BIT_SHIFT;
+
+        // Safety:
+        // The shift retrieves the original value from `T::into_usize`,
+        // satisfying `T::from_usize`'s preconditions.
+        unsafe { T::from_usize(tag) }
     }
+
+    /// Sets the tag to a new value.
     #[inline]
     pub fn set_tag(&mut self, tag: T) {
-        let mut packed = self.packed.get();
-        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::pack(self.pointer_raw(), tag);
+    }
+
+    const TAG_BIT_SHIFT: u32 = usize::BITS - T::BITS;
+    const ASSERTION: () = { assert!(T::BITS <= P::BITS) };
+
+    /// Pack pointer `ptr` that comes from [`P::into_ptr`] with a `tag`,
+    /// according to `self.packed` encoding scheme.
+    ///
+    /// [`P::into_ptr`]: Pointer::into_ptr
+    #[inline]
+    fn pack(ptr: NonNull<P::Target>, tag: T) -> NonNull<P::Target> {
+        // Trigger assert!
+        let () = Self::ASSERTION;
+
+        let packed_tag = tag.into_usize() << Self::TAG_BIT_SHIFT;
+
+        ptr.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) }
+        })
     }
+
+    /// Retrieves the original raw pointer from `self.packed`.
+    #[inline]
+    pub(super) fn pointer_raw(&self) -> NonNull<P::Target> {
+        self.packed.map_addr(|addr| unsafe { NonZeroUsize::new_unchecked(addr.get() << T::BITS) })
+    }
+
+    /// This provides a reference to the `P` pointer itself, rather than the
+    /// `Deref::Target`. It is used for cases where we want to call methods
+    /// that may be implement differently for the Pointer than the Pointee
+    /// (e.g., `Rc::clone` vs cloning the inner value).
+    pub(super) fn with_pointer_ref<R>(&self, f: impl FnOnce(&P) -> R) -> R {
+        // Safety:
+        // - `self.raw.pointer_raw()` is originally returned from `P::into_ptr`
+        //   and as such is valid for `P::from_ptr`.
+        //   - This also allows us to not care whatever `f` panics or not.
+        // - Even though we create a copy of the pointer, we store it inside
+        //   `ManuallyDrop` and only access it by-ref, so we don't double-drop.
+        //
+        // Semantically this is just `f(&self.pointer)` (where `self.pointer`
+        // is non-packed original pointer).
+        //
+        // Note that even though `CopyTaggedPtr` is only really expected to
+        // work with `P: Copy`, we have to assume `P: ?Copy`, because
+        // `CopyTaggedPtr` is used in the `TaggedPtr`'s implementation.
+        let ptr = unsafe { ManuallyDrop::new(P::from_ptr(self.pointer_raw())) };
+        f(&ptr)
+    }
+}
+
+impl<P, T, const CP: bool> Copy for CopyTaggedPtr<P, T, CP>
+where
+    P: Pointer + Copy,
+    T: Tag,
+{
 }
 
-impl<P, T, const COMPARE_PACKED: bool> std::ops::Deref for CopyTaggedPtr<P, T, COMPARE_PACKED>
+impl<P, T, const CP: bool> Clone for CopyTaggedPtr<P, T, CP>
+where
+    P: Pointer + Copy,
+    T: Tag,
+{
+    #[inline]
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<P, T, const CP: bool> Deref for CopyTaggedPtr<P, T, CP>
 where
     P: Pointer,
     T: Tag,
 {
     type Target = P::Target;
+
+    #[inline]
     fn deref(&self) -> &Self::Target {
-        self.pointer_ref()
+        // Safety:
+        // `pointer_raw` returns the original pointer from `P::into_ptr` which,
+        // by the `Pointer`'s contract, must be valid.
+        unsafe { self.pointer_raw().as_ref() }
     }
 }
 
-impl<P, T, const COMPARE_PACKED: bool> std::ops::DerefMut for CopyTaggedPtr<P, T, COMPARE_PACKED>
+impl<P, T, const CP: bool> DerefMut for CopyTaggedPtr<P, T, CP>
 where
-    P: Pointer + std::ops::DerefMut,
+    P: Pointer + DerefMut,
     T: Tag,
 {
+    #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
-        self.pointer_mut()
+        // Safety:
+        // `pointer_raw` returns the original pointer from `P::into_ptr` which,
+        // by the `Pointer`'s contract, must be valid for writes if
+        // `P: DerefMut`.
+        unsafe { self.pointer_raw().as_mut() }
     }
 }
 
-impl<P, T, const COMPARE_PACKED: bool> fmt::Debug for CopyTaggedPtr<P, T, COMPARE_PACKED>
+impl<P, T, const CP: bool> fmt::Debug for CopyTaggedPtr<P, T, CP>
 where
-    P: Pointer,
-    P::Target: fmt::Debug,
+    P: Pointer + fmt::Debug,
     T: Tag + fmt::Debug,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("CopyTaggedPtr")
-            .field("pointer", &self.pointer_ref())
-            .field("tag", &self.tag())
-            .finish()
+        self.with_pointer_ref(|ptr| {
+            f.debug_struct("CopyTaggedPtr").field("pointer", ptr).field("tag", &self.tag()).finish()
+        })
     }
 }
 
@@ -149,6 +242,7 @@ where
     P: Pointer,
     T: Tag,
 {
+    #[inline]
     fn eq(&self, other: &Self) -> bool {
         self.packed == other.packed
     }
@@ -161,25 +255,74 @@ where
 {
 }
 
-impl<P, T> std::hash::Hash for CopyTaggedPtr<P, T, true>
+impl<P, T> Hash for CopyTaggedPtr<P, T, true>
 where
     P: Pointer,
     T: Tag,
 {
-    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+    #[inline]
+    fn hash<H: Hasher>(&self, state: &mut H) {
         self.packed.hash(state);
     }
 }
 
-impl<P, T, HCX, const COMPARE_PACKED: bool> HashStable<HCX> for CopyTaggedPtr<P, T, COMPARE_PACKED>
+impl<P, T, HCX, const CP: bool> HashStable<HCX> for CopyTaggedPtr<P, T, CP>
 where
     P: Pointer + HashStable<HCX>,
     T: Tag + HashStable<HCX>,
 {
     fn hash_stable(&self, hcx: &mut HCX, hasher: &mut StableHasher) {
-        unsafe {
-            Pointer::with_ref(self.pointer_raw(), |p: &P| p.hash_stable(hcx, hasher));
-        }
+        self.with_pointer_ref(|ptr| ptr.hash_stable(hcx, hasher));
         self.tag().hash_stable(hcx, hasher);
     }
 }
+
+// Safety:
+// `CopyTaggedPtr<P, T, ..>` is semantically just `{ ptr: P, tag: T }`, as such
+// it's ok to implement `Sync` as long as `P: Sync, T: Sync`
+unsafe impl<P, T, const CP: bool> Sync for CopyTaggedPtr<P, T, CP>
+where
+    P: Sync + Pointer,
+    T: Sync + Tag,
+{
+}
+
+// Safety:
+// `CopyTaggedPtr<P, T, ..>` is semantically just `{ ptr: P, tag: T }`, as such
+// it's ok to implement `Send` as long as `P: Send, T: Send`
+unsafe impl<P, T, const CP: bool> Send for CopyTaggedPtr<P, T, CP>
+where
+    P: Send + Pointer,
+    T: Send + Tag,
+{
+}
+
+/// Test that `new` does not compile if there is not enough alignment for the
+/// tag in the pointer.
+///
+/// ```compile_fail,E0080
+/// use rustc_data_structures::tagged_ptr::{CopyTaggedPtr, Tag};
+///
+/// #[derive(Copy, Clone, Debug, PartialEq, Eq)]
+/// enum Tag2 { B00 = 0b00, B01 = 0b01, B10 = 0b10, B11 = 0b11 };
+///
+/// unsafe impl Tag for Tag2 {
+///     const BITS: u32 = 2;
+///
+///     fn into_usize(self) -> usize { todo!() }
+///     unsafe fn from_usize(tag: usize) -> Self { todo!() }
+/// }
+///
+/// let value = 12u16;
+/// let reference = &value;
+/// let tag = Tag2::B01;
+///
+/// let _ptr = CopyTaggedPtr::<_, _, true>::new(reference, tag);
+/// ```
+// For some reason miri does not get the compile error
+// probably it `check`s instead of `build`ing?
+#[cfg(not(miri))]
+const _: () = ();
+
+#[cfg(test)]
+mod tests;