store key in the lower bits of the string length

* Optimizes for calling `.data()`

Author: joseph-gio

benches/benches/bevy_ecs/scheduling/schedule.rs

@@ -64,13 +64,13 @@ pub fn build_schedule(criterion: &mut Criterion) {
     // Use multiple different kinds of label to ensure that dynamic dispatch
     // doesn't somehow get optimized away.
     #[derive(Debug, Clone, Copy)]
-    struct NumLabel(usize);
+    struct NumLabel(u16);
     #[derive(Debug, Clone, Copy, SystemLabel)]
     struct DummyLabel;
 
     impl SystemLabel for NumLabel {
-        fn data(&self) -> u64 {
-            self.0 as u64
+        fn data(&self) -> u16 {
+            self.0
         }
         fn as_str(&self) -> &'static str {
             let s = self.0.to_string();
@@ -88,7 +88,7 @@ pub fn build_schedule(criterion: &mut Criterion) {
     // Also, we are performing the `as_label` operation outside of the loop since that
     // requires an allocation and a leak. This is not something that would be necessary in a
     // real scenario, just a contrivance for the benchmark.
-    let labels: Vec<_> = (0..1000).map(|i| NumLabel(i).as_label()).collect();
+    let labels: Vec<_> = (0..1000).map(|i| NumLabel(i as u16).as_label()).collect();
 
     // Benchmark graphs of different sizes.
     for graph_size in [100, 500, 1000] {

crates/bevy_ecs/src/schedule/state.rs

@@ -58,7 +58,7 @@ impl RunCriteriaLabel for DriverLabel {
     fn type_id(&self) -> core::any::TypeId {
         self.0
     }
-    fn data(&self) -> u64 {
+    fn data(&self) -> u16 {
         0
     }
     fn as_str(&self) -> &'static str {

crates/bevy_ecs/src/system/function_system.rs

@@ -458,7 +458,7 @@ impl<T: 'static> SystemLabel for SystemTypeIdLabel<T> {
         std::any::type_name::<T>()
     }
     #[inline]
-    fn data(&self) -> u64 {
+    fn data(&self) -> u16 {
         0
     }
 }

crates/bevy_macro_utils/src/lib.rs

@@ -191,7 +191,8 @@ pub fn derive_label(
                     let mut path = syn::Path::from(ident.clone());
                     path.segments.push(v.ident.clone().into());
 
-                    let i = i as u64;
+                    // hopefully no one has more than 60 thousand variants
+                    let i = i as u16;
                     data_arms.push(quote! { #path { .. } => #i });
 
                     let lit = format!("{ident}::{}", v.ident.clone());
@@ -227,7 +228,7 @@ pub fn derive_label(
 
     (quote! {
         impl #impl_generics #trait_path for #ident #ty_generics #where_clause {
-            fn data(&self) -> u64 {
+            fn data(&self) -> u16 {
                 #data
             }
             fn as_str(&self) -> &'static str {

crates/bevy_utils/src/label.rs

@@ -3,6 +3,7 @@
 use std::{
     any::Any,
     hash::{Hash, Hasher},
+    marker::PhantomData,
 };
 
 pub trait DynEq: Any {
@@ -47,6 +48,66 @@ where
     }
 }
 
+#[doc(hidden)]
+/// A string slice, stuffed with a payload of data `16` bits long.
+/// This means that the maximum length of the string is `2^(size-16)`,
+/// where `size` is either 32 or 64 depending on your computer.
+#[derive(Clone, Copy)]
+pub struct StuffedStr<'a> {
+    ptr: *const u8,
+    marker: PhantomData<&'a str>,
+
+    /// lower 16 bits stores the data,
+    /// upper 48 or 16 bits stores the length of the string.
+    /// This layout optimizes for calling `.data()`
+    meta: usize,
+}
+
+unsafe impl<'a> Send for StuffedStr<'a> {}
+unsafe impl<'a> Sync for StuffedStr<'a> {}
+
+impl<'a> StuffedStr<'a> {
+    const DATA_MASK: usize = !Self::LEN_MASK;
+    const LEN_MASK: usize = !0 << 16;
+
+    pub fn new(str: &'a str, data: u16) -> Self {
+        #[cold]
+        #[inline(never)]
+        fn panic() -> ! {
+            // needs a better message but its late
+            panic!("string exceeds {} bytes", StuffedStr::LEN_MASK >> 16);
+        }
+
+        let ptr = str.as_ptr();
+        // Make sure there's enough room to store the data.
+        if str.len().leading_zeros() < 16 {
+            panic();
+        }
+        let meta = data as usize | str.len() << 16;
+        Self {
+            ptr,
+            marker: PhantomData,
+            meta,
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.meta >> 16
+    }
+    pub fn data(&self) -> u16 {
+        let data = self.meta & Self::DATA_MASK;
+        data as u16
+    }
+    pub fn as_str(&self) -> &'a str {
+        // SAFETY: this instance was constructed from a string slice of length `len`, and lifetime `'a`.
+        // It is sound to convert it back to a slice.
+        unsafe {
+            let slice = std::slice::from_raw_parts(self.ptr, self.len());
+            std::str::from_utf8_unchecked(slice)
+        }
+    }
+}
+
 /// Macro to define a new label trait
 ///
 /// # Example
@@ -71,11 +132,7 @@ macro_rules! define_label {
     ) => {
         $(#[$id_attr])*
         #[derive(Clone, Copy)]
-        pub struct $id_name {
-            ty: ::core::any::TypeId,
-            data: u64,
-            text: &'static str
-        }
+        pub struct $id_name(::core::any::TypeId, $crate::label::StuffedStr<'static>);
 
         impl ::core::fmt::Debug for $id_name {
             fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
@@ -90,14 +147,15 @@ macro_rules! define_label {
                 let ty = self.type_id();
                 let data = self.data();
                 let text = self.as_str();
-                $id_name { ty, data, text }
+                let stuffed = $crate::label::StuffedStr::new(text, data);
+                $id_name(ty, stuffed)
             }
             /// Returns the [`TypeId`] used to differentiate labels.
             fn type_id(&self) -> ::core::any::TypeId {
                 ::core::any::TypeId::of::<Self>()
             }
             /// Returns a number used to distinguish different labels of the same type.
-            fn data(&self) -> u64;
+            fn data(&self) -> u16;
             /// Returns the representation of this label as a string literal.
             ///
             /// In cases where you absolutely need a label to be determined at runtime,
@@ -110,13 +168,13 @@ macro_rules! define_label {
                 *self
             }
             fn type_id(&self) -> ::core::any::TypeId {
-                self.ty
+                self.0
             }
-            fn data(&self) -> u64 {
-                self.data
+            fn data(&self) -> u16 {
+                self.1.data()
             }
             fn as_str(&self) -> &'static str {
-                self.text
+                self.1.as_str()
             }
         }