allow arbitrary formatting using a fn pointer

Author: joseph-gio

crates/bevy_utils/src/label.rs

@@ -5,6 +5,8 @@ use std::{
     hash::{Hash, Hasher},
 };
 
+use crate::Interner;
+
 pub trait DynEq: Any {
     fn as_any(&self) -> &dyn Any;
 
@@ -47,6 +49,30 @@ where
     }
 }
 
+#[doc(hidden)]
+pub struct VTable {
+    // FIXME: When const TypeId stabilizes, inline the type instead of using a fn pointer for indirection.
+    pub ty: fn() -> ::std::any::TypeId,
+    pub fmt: fn(u64, &mut ::std::fmt::Formatter) -> ::std::fmt::Result,
+}
+
+impl PartialEq for VTable {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        (self.ty)() == (other.ty)()
+    }
+}
+impl Eq for VTable {}
+
+impl Hash for VTable {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        (self.ty)().hash(state);
+    }
+}
+
+#[doc(hidden)]
+pub static STR_INTERN: Interner<&str> = Interner::new();
+
 /// Macro to define a new label trait
 ///
 /// # Example
@@ -71,48 +97,83 @@ macro_rules! define_label {
     ) => {
         $(#[$id_attr])*
         #[derive(Clone, Copy, PartialEq, Eq, Hash)]
-        pub struct $id_name(::core::any::TypeId, &'static str);
+        pub struct $id_name {
+            data: u64,
+            vtable: &'static $crate::label::VTable,
+        }
 
-        impl ::core::fmt::Debug for $id_name {
-            fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
-                write!(f, "{}", self.1)
+        impl ::std::fmt::Debug for $id_name {
+            fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
+                (self.vtable.fmt)(self.data, f)
             }
         }
 
         $(#[$label_attr])*
         pub trait $label_name: 'static {
             /// Converts this type into an opaque, strongly-typed label.
+            #[inline]
             fn as_label(&self) -> $id_name {
-                let id = self.type_id();
-                let label = self.as_str();
-                $id_name(id, label)
-            }
-            /// Returns the [`TypeId`] used to differentiate labels.
-            fn type_id(&self) -> ::core::any::TypeId {
-                ::core::any::TypeId::of::<Self>()
+                // This is just machinery that lets us store the TypeId and formatter fn in the same static reference.
+                struct VTables<L: ?::std::marker::Sized>(L);
+                impl<L: $label_name + ?::std::marker::Sized> VTables<L> {
+                    const VTABLE: $crate::label::VTable = $crate::label::VTable {
+                        ty: || ::std::any::TypeId::of::<L>(),
+                        fmt: <L as $label_name>::fmt,
+                    };
+                }
+
+                let data = self.data();
+                $id_name { data, vtable: &VTables::<Self>::VTABLE }
             }
-            /// Returns the representation of this label as a string literal.
+            /// Returns a number used to distinguish different labels of the same type.
+            fn data(&self) -> u64;
+            /// Writes debug info for a label of the current type.
+            /// * `data`: the result of calling [`data()`](#method.data) on an instance of this type.
             ///
-            /// In cases where you absolutely need a label to be determined at runtime,
-            /// you can use [`Box::leak`] to get a `'static` reference.
-            fn as_str(&self) -> &'static str;
+            /// You should not call this method directly, as it may panic for some types;
+            /// use [`as_label`](#method.as_label) instead.
+            fn fmt(data: u64, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result;
         }
 
         impl $label_name for $id_name {
+            #[inline]
             fn as_label(&self) -> Self {
                 *self
             }
-            fn type_id(&self) -> ::core::any::TypeId {
-                self.0
+            #[inline]
+            fn data(&self) -> u64 {
+                self.data
             }
-            fn as_str(&self) -> &'static str {
-                self.1
+            #[track_caller]
+            fn fmt(data: u64, f: &mut ::std::fmt::Formatter) -> std::fmt::Result {
+                let label = stringify!($label_name);
+                ::std::unimplemented!("do not call `{label}::fmt` directly -- use the result of `as_label()` for formatting instead")
+            }
+        }
+
+        impl $id_name {
+            /// Returns the [`TypeId`] of the label from which this ID was constructed.
+            ///
+            /// [`TypeId`]: ::std::any::TypeId
+            #[inline]
+            pub fn type_id(self) -> ::std::any::TypeId {
+                (self.vtable.ty)()
+            }
+            /// Returns true if this label was constructed from an instance of type `L`.
+            pub fn is<L: $label_name>(self) -> bool {
+                self.type_id() == ::std::any::TypeId::of::<L>()
             }
         }
 
         impl $label_name for &'static str {
-            fn as_str(&self) -> Self {
-                self
+            fn data(&self) -> u64 {
+                $crate::label::STR_INTERN.intern(self) as u64
+            }
+            fn fmt(idx: u64, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+                let s = $crate::label::STR_INTERN
+                    .get(idx as usize)
+                    .ok_or(::std::fmt::Error)?;
+                write!(f, "{s}")
             }
         }
     };