Auto merge of rust-lang#75590 - Mark-Simulacrum:tagged-ptr, r=ecstatic-morse

Add a packed/tagged pointer abstraction and utilize it for ParamEnv

The intent here is mostly just to add the abstraction; I suspect that there are definitely more use cases for it, and we can explore those over time now that there's a (mostly) safe abstraction that can be used in rustc.

Author: bors

src/librustc_data_structures/lib.rs

@@ -7,6 +7,7 @@
 //! This API is completely unstable and subject to change.
 
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/")]
+#![allow(incomplete_features)]
 #![feature(in_band_lifetimes)]
 #![feature(unboxed_closures)]
 #![feature(generators)]
@@ -23,6 +24,8 @@
 #![feature(associated_type_bounds)]
 #![feature(thread_id_value)]
 #![feature(extend_one)]
+#![feature(const_panic)]
+#![feature(const_generics)]
 #![allow(rustc::default_hash_types)]
 
 #[macro_use]
@@ -97,6 +100,7 @@ pub mod vec_linked_list;
 pub mod work_queue;
 pub use atomic_ref::AtomicRef;
 pub mod frozen;
+pub mod tagged_ptr;
 pub mod temp_dir;
 
 pub struct OnDrop<F: Fn()>(pub F);

src/librustc_data_structures/tagged_ptr.rs

@@ -0,0 +1,157 @@
+//! This module implements tagged pointers.
+//!
+//! In order to utilize the pointer packing, you must have two types: a pointer,
+//! and a tag.
+//!
+//! The pointer must implement the `Pointer` trait, with the primary requirement
+//! being conversion to and from a usize. Note that the pointer must be
+//! dereferenceable, so raw pointers generally cannot implement the `Pointer`
+//! trait. This implies that the pointer must also be nonzero.
+//!
+//! Many common pointer types already implement the `Pointer` trait.
+//!
+//! The tag must implement the `Tag` trait. We assert that the tag and `Pointer`
+//! are compatible at compile time.
+
+use std::mem::ManuallyDrop;
+use std::ops::Deref;
+use std::rc::Rc;
+use std::sync::Arc;
+
+mod copy;
+mod drop;
+
+pub use copy::CopyTaggedPtr;
+pub use drop::TaggedPtr;
+
+/// This describes the pointer type encaspulated by TaggedPtr.
+///
+/// # Safety
+///
+/// The usize returned from `into_usize` must be a valid, dereferenceable,
+/// pointer to `<Self as Deref>::Target`. Note that pointers to `Pointee` must
+/// be thin, even though `Pointee` may not be sized.
+///
+/// Note that the returned pointer from `into_usize` should be castable to `&mut
+/// <Self as Deref>::Target` if `Pointer: DerefMut`.
+///
+/// The BITS constant must be correct. At least `BITS` bits, least-significant,
+/// must be zero on all returned pointers from `into_usize`.
+///
+/// For example, if the alignment of `Pointee` is 2, then `BITS` should be 1.
+pub unsafe trait Pointer: Deref {
+    /// Most likely the value you want to use here is the following, unless
+    /// your Pointee type is unsized (e.g., `ty::List<T>` in rustc) in which
+    /// case you'll need to manually figure out what the right type to pass to
+    /// align_of is.
+    ///
+    /// ```rust
+    /// std::mem::align_of::<<Self as Deref>::Target>().trailing_zeros() as usize;
+    /// ```
+    const BITS: usize;
+    fn into_usize(self) -> usize;
+
+    /// # Safety
+    ///
+    /// The passed `ptr` must be returned from `into_usize`.
+    ///
+    /// 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;
+
+    /// 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
+    /// may be implement differently for the Pointer than the Pointee (e.g.,
+    /// `Rc::clone` vs cloning the inner value).
+    ///
+    /// # Safety
+    ///
+    /// The passed `ptr` must be returned from `into_usize`.
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R;
+}
+
+/// This describes tags that the `TaggedPtr` struct can hold.
+///
+/// # Safety
+///
+/// The BITS constant must be correct.
+///
+/// No more than `BITS` least significant bits may be set in the returned usize.
+pub unsafe trait Tag: Copy {
+    const BITS: usize;
+
+    fn into_usize(self) -> usize;
+
+    /// # Safety
+    ///
+    /// The passed `tag` must be returned from `into_usize`.
+    unsafe fn from_usize(tag: usize) -> Self;
+}
+
+unsafe impl<T> Pointer for Box<T> {
+    const BITS: usize = std::mem::align_of::<T>().trailing_zeros() as usize;
+    fn into_usize(self) -> usize {
+        Box::into_raw(self) as usize
+    }
+    unsafe fn from_usize(ptr: usize) -> Self {
+        Box::from_raw(ptr as *mut T)
+    }
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
+        let raw = ManuallyDrop::new(Self::from_usize(ptr));
+        f(&raw)
+    }
+}
+
+unsafe impl<T> Pointer for Rc<T> {
+    const BITS: usize = std::mem::align_of::<T>().trailing_zeros() as usize;
+    fn into_usize(self) -> usize {
+        Rc::into_raw(self) as usize
+    }
+    unsafe fn from_usize(ptr: usize) -> Self {
+        Rc::from_raw(ptr as *const T)
+    }
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
+        let raw = ManuallyDrop::new(Self::from_usize(ptr));
+        f(&raw)
+    }
+}
+
+unsafe impl<T> Pointer for Arc<T> {
+    const BITS: usize = std::mem::align_of::<T>().trailing_zeros() as usize;
+    fn into_usize(self) -> usize {
+        Arc::into_raw(self) as usize
+    }
+    unsafe fn from_usize(ptr: usize) -> Self {
+        Arc::from_raw(ptr as *const T)
+    }
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
+        let raw = ManuallyDrop::new(Self::from_usize(ptr));
+        f(&raw)
+    }
+}
+
+unsafe impl<'a, T: 'a> Pointer for &'a T {
+    const BITS: usize = std::mem::align_of::<T>().trailing_zeros() as usize;
+    fn into_usize(self) -> usize {
+        self as *const T as usize
+    }
+    unsafe fn from_usize(ptr: usize) -> Self {
+        &*(ptr as *const T)
+    }
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
+        f(&*(&ptr as *const usize as *const Self))
+    }
+}
+
+unsafe impl<'a, T: 'a> Pointer for &'a mut T {
+    const BITS: usize = std::mem::align_of::<T>().trailing_zeros() as usize;
+    fn into_usize(self) -> usize {
+        self as *mut T as usize
+    }
+    unsafe fn from_usize(ptr: usize) -> Self {
+        &mut *(ptr as *mut T)
+    }
+    unsafe fn with_ref<R, F: FnOnce(&Self) -> R>(ptr: usize, f: F) -> R {
+        f(&*(&ptr as *const usize as *const Self))
+    }
+}

src/librustc_data_structures/tagged_ptr/copy.rs

@@ -0,0 +1,183 @@
+use super::{Pointer, Tag};
+use crate::stable_hasher::{HashStable, StableHasher};
+use std::fmt;
+use std::marker::PhantomData;
+use std::num::NonZeroUsize;
+
+/// A `Copy` TaggedPtr.
+///
+/// You should use this instead of the `TaggedPtr` type in all cases where
+/// `P: 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.
+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
+    }
+}
+
+// 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>
+where
+    P: Pointer,
+    T: Tag,
+{
+    const TAG_BIT_SHIFT: usize = (8 * std::mem::size_of::<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>());
+    };
+
+    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,
+        }
+    }
+
+    pub(super) fn pointer_raw(&self) -> usize {
+        self.packed.get() << T::BITS
+    }
+    pub fn pointer(self) -> P
+    where
+        P: Copy,
+    {
+        // SAFETY: pointer_raw returns the original pointer
+        //
+        // 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()) }
+    }
+    pub fn tag(&self) -> T {
+        unsafe { T::from_usize(self.packed.get() >> Self::TAG_BIT_SHIFT) }
+    }
+    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) };
+    }
+}
+
+impl<P, T, const COMPARE_PACKED: bool> std::ops::Deref for CopyTaggedPtr<P, T, COMPARE_PACKED>
+where
+    P: Pointer,
+    T: Tag,
+{
+    type Target = P::Target;
+    fn deref(&self) -> &Self::Target {
+        self.pointer_ref()
+    }
+}
+
+impl<P, T, const COMPARE_PACKED: bool> std::ops::DerefMut for CopyTaggedPtr<P, T, COMPARE_PACKED>
+where
+    P: Pointer + std::ops::DerefMut,
+    T: Tag,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.pointer_mut()
+    }
+}
+
+impl<P, T, const COMPARE_PACKED: bool> fmt::Debug for CopyTaggedPtr<P, T, COMPARE_PACKED>
+where
+    P: Pointer,
+    P::Target: 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()
+    }
+}
+
+impl<P, T> PartialEq for CopyTaggedPtr<P, T, true>
+where
+    P: Pointer,
+    T: Tag,
+{
+    fn eq(&self, other: &Self) -> bool {
+        self.packed == other.packed
+    }
+}
+
+impl<P, T> Eq for CopyTaggedPtr<P, T, true>
+where
+    P: Pointer,
+    T: Tag,
+{
+}
+
+impl<P, T> std::hash::Hash for CopyTaggedPtr<P, T, true>
+where
+    P: Pointer,
+    T: Tag,
+{
+    fn hash<H: std::hash::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>
+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.tag().hash_stable(hcx, hasher);
+    }
+}