Add impl_tag! macro to implement Tag for tagged pointer easily

compiler/rustc_data_structures/src/tagged_ptr.rs

@@ -24,6 +24,7 @@ use crate::aligned::Aligned;
 
 mod copy;
 mod drop;
+mod impl_tag;
 
 pub use copy::CopyTaggedPtr;
 pub use drop::TaggedPtr;
@@ -141,6 +142,40 @@ pub unsafe trait Tag: Copy {
     unsafe fn from_usize(tag: usize) -> Self;
 }
 
+/// Returns the number of bits available for use for tags in a pointer to `T`
+/// (this is based on `T`'s alignment).
+pub const fn bits_for<T: ?Sized + Aligned>() -> u32 {
+    crate::aligned::align_of::<T>().as_nonzero().trailing_zeros()
+}
+
+/// Returns the correct [`Tag::BITS`] constant for a set of tag values.
+pub const fn bits_for_tags(mut tags: &[usize]) -> u32 {
+    let mut bits = 0;
+
+    while let &[tag, ref rest @ ..] = tags {
+        tags = rest;
+
+        let b = bits_for_tag(tag);
+        if b > bits {
+            bits = b;
+        }
+    }
+
+    bits
+}
+
+/// Returns `(size_of::<usize>() * 8) - tag.leading_zeros()`
+const fn bits_for_tag(mut tag: usize) -> u32 {
+    let mut bits = 0;
+
+    while tag > 0 {
+        bits += 1;
+        tag >>= 1;
+    }
+
+    bits
+}
+
 unsafe impl<T: ?Sized + Aligned> Pointer for Box<T> {
     const BITS: u32 = bits_for::<Self::Target>();
 
@@ -221,12 +256,6 @@ unsafe impl<'a, T: 'a + ?Sized + Aligned> Pointer for &'a mut T {
     }
 }
 
-/// Returns the number of bits available for use for tags in a pointer to `T`
-/// (this is based on `T`'s alignment).
-pub const fn bits_for<T: ?Sized + Aligned>() -> u32 {
-    crate::aligned::align_of::<T>().as_nonzero().trailing_zeros()
-}
-
 /// A tag type used in [`CopyTaggedPtr`] and [`TaggedPtr`] tests.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 #[cfg(test)]

compiler/rustc_data_structures/src/tagged_ptr/impl_tag.rs

@@ -0,0 +1,139 @@
+/// Implements [`Tag`] for a given type.
+///
+/// You can use `impl_tag` on structs and enums.
+/// You need to specify the type and all its possible values,
+/// which can only be paths with optional fields.
+///
+/// [`Tag`]: crate::tagged_ptr::Tag
+///
+/// # Examples
+///
+/// Basic usage:
+///
+/// ```
+/// use rustc_data_structures::{impl_tag, tagged_ptr::Tag};
+///
+/// #[derive(Copy, Clone, PartialEq, Debug)]
+/// enum SomeTag {
+///     A,
+///     B,
+///     X { v: bool },
+///     Y(bool, bool),
+/// }
+///
+/// impl_tag! {
+///     // The type for which the `Tag` will be implemented
+///     impl Tag for SomeTag;
+///     // You need to specify the `{value_of_the_type} <=> {tag}` relationship
+///     SomeTag::A <=> 0,
+///     SomeTag::B <=> 1,
+///     // For variants with fields, you need to specify the fields:
+///     SomeTag::X { v: true  } <=> 2,
+///     SomeTag::X { v: false } <=> 3,
+///     // For tuple variants use named syntax:
+///     SomeTag::Y { 0: true,  1: true  } <=> 4,
+///     SomeTag::Y { 0: false, 1: true  } <=> 5,
+///     SomeTag::Y { 0: true,  1: false } <=> 6,
+///     SomeTag::Y { 0: false, 1: false } <=> 7,
+/// }
+///
+/// assert_eq!(SomeTag::A.into_usize(), 0);
+/// assert_eq!(SomeTag::X { v: false }.into_usize(), 3);
+/// assert_eq!(SomeTag::Y(false, true).into_usize(), 5);
+///
+/// assert_eq!(unsafe { SomeTag::from_usize(1) }, SomeTag::B);
+/// assert_eq!(unsafe { SomeTag::from_usize(2) }, SomeTag::X { v: true });
+/// assert_eq!(unsafe { SomeTag::from_usize(7) }, SomeTag::Y(false, false));
+/// ```
+///
+/// Structs are supported:
+///
+/// ```
+/// # use rustc_data_structures::impl_tag;
+/// #[derive(Copy, Clone)]
+/// struct Flags { a: bool, b: bool }
+///
+/// impl_tag! {
+///     impl Tag for Flags;
+///     Flags { a: true,  b: true  } <=> 3,
+///     Flags { a: false, b: true  } <=> 2,
+///     Flags { a: true,  b: false } <=> 1,
+///     Flags { a: false, b: false } <=> 0,
+/// }
+/// ```
+///
+/// Not specifying all values results in a compile error:
+///
+/// ```compile_fail,E0004
+/// # use rustc_data_structures::impl_tag;
+/// #[derive(Copy, Clone)]
+/// enum E {
+///     A,
+///     B,
+/// }
+///
+/// impl_tag! {
+///     impl Tag for E;
+///     E::A <=> 0,
+/// }
+/// ```
+#[macro_export]
+macro_rules! impl_tag {
+    (
+        impl Tag for $Self:ty;
+        $(
+            $($path:ident)::* $( { $( $fields:tt )* })? <=> $tag:literal,
+        )*
+    ) => {
+        // Safety:
+        // `into_usize` only returns one of `$tag`s,
+        // `bits_for_tags` is called on all `$tag`s,
+        // thus `BITS` constant is correct.
+        unsafe impl $crate::tagged_ptr::Tag for $Self {
+            const BITS: u32 = $crate::tagged_ptr::bits_for_tags(&[
+                $( $tag, )*
+            ]);
+
+            fn into_usize(self) -> usize {
+                // This forbids use of repeating patterns (`Enum::V`&`Enum::V`, etc)
+                // (or at least it should, see <https://github.com/rust-lang/rust/issues/110613>)
+                #[forbid(unreachable_patterns)]
+                match self {
+                    // `match` is doing heavy lifting here, by requiring exhaustiveness
+                    $(
+                        $($path)::* $( { $( $fields )* } )? => $tag,
+                    )*
+                }
+            }
+
+            unsafe fn from_usize(tag: usize) -> Self {
+                // Similarly to the above, this forbids repeating tags
+                // (or at least it should, see <https://github.com/rust-lang/rust/issues/110613>)
+                #[forbid(unreachable_patterns)]
+                match tag {
+                    $(
+                        $tag => $($path)::* $( { $( $fields )* } )?,
+                    )*
+
+                    // Safety:
+                    // `into_usize` only returns one of `$tag`s,
+                    // all `$tag`s are filtered up above,
+                    // thus if this is reached, the safety contract of this
+                    // function was already breached.
+                    _ => unsafe {
+                        debug_assert!(
+                            false,
+                            "invalid tag: {tag}\
+                             (this is a bug in the caller of `from_usize`)"
+                        );
+                        std::hint::unreachable_unchecked()
+                    },
+                }
+            }
+
+        }
+    };
+}
+
+#[cfg(test)]
+mod tests;

compiler/rustc_data_structures/src/tagged_ptr/impl_tag/tests.rs

@@ -0,0 +1,33 @@
+#[test]
+fn bits_constant() {
+    use crate::tagged_ptr::Tag;
+
+    #[derive(Copy, Clone)]
+    struct Unit;
+    impl_tag! { impl Tag for Unit; Unit <=> 0, }
+    assert_eq!(Unit::BITS, 0);
+
+    #[derive(Copy, Clone)]
+    struct Unit1;
+    impl_tag! { impl Tag for Unit1; Unit1 <=> 1, }
+    assert_eq!(Unit1::BITS, 1);
+
+    #[derive(Copy, Clone)]
+    struct Unit2;
+    impl_tag! { impl Tag for Unit2; Unit2 <=> 0b10, }
+    assert_eq!(Unit2::BITS, 2);
+
+    #[derive(Copy, Clone)]
+    struct Unit3;
+    impl_tag! { impl Tag for Unit3; Unit3 <=> 0b100, }
+    assert_eq!(Unit3::BITS, 3);
+
+    #[derive(Copy, Clone)]
+    enum Enum {
+        A,
+        B,
+        C,
+    }
+    impl_tag! { impl Tag for Enum; Enum::A <=> 0b1, Enum::B <=> 0b1000, Enum::C <=> 0b10, }
+    assert_eq!(Enum::BITS, 4);
+}

compiler/rustc_middle/src/ty/mod.rs

@@ -1626,29 +1626,12 @@ struct ParamTag {
     constness: hir::Constness,
 }
 
-unsafe impl rustc_data_structures::tagged_ptr::Tag for ParamTag {
-    const BITS: u32 = 2;
-
-    #[inline]
-    fn into_usize(self) -> usize {
-        match self {
-            Self { reveal: traits::Reveal::UserFacing, constness: hir::Constness::NotConst } => 0,
-            Self { reveal: traits::Reveal::All, constness: hir::Constness::NotConst } => 1,
-            Self { reveal: traits::Reveal::UserFacing, constness: hir::Constness::Const } => 2,
-            Self { reveal: traits::Reveal::All, constness: hir::Constness::Const } => 3,
-        }
-    }
-
-    #[inline]
-    unsafe fn from_usize(ptr: usize) -> Self {
-        match ptr {
-            0 => Self { reveal: traits::Reveal::UserFacing, constness: hir::Constness::NotConst },
-            1 => Self { reveal: traits::Reveal::All, constness: hir::Constness::NotConst },
-            2 => Self { reveal: traits::Reveal::UserFacing, constness: hir::Constness::Const },
-            3 => Self { reveal: traits::Reveal::All, constness: hir::Constness::Const },
-            _ => std::hint::unreachable_unchecked(),
-        }
-    }
+impl_tag! {
+    impl Tag for ParamTag;
+    ParamTag { reveal: traits::Reveal::UserFacing, constness: hir::Constness::NotConst } <=> 0,
+    ParamTag { reveal: traits::Reveal::All,        constness: hir::Constness::NotConst } <=> 1,
+    ParamTag { reveal: traits::Reveal::UserFacing, constness: hir::Constness::Const    } <=> 2,
+    ParamTag { reveal: traits::Reveal::All,        constness: hir::Constness::Const    } <=> 3,
 }
 
 impl<'tcx> fmt::Debug for ParamEnv<'tcx> {