Fallible from_bits, Identifier repr optimisations

Author: Bluefinger

crates/bevy_ecs/src/entity/mod.rs

@@ -267,20 +267,37 @@ impl Entity {
     /// Reconstruct an `Entity` previously destructured with [`Entity::to_bits`].
     ///
     /// Only useful when applied to results from `to_bits` in the same instance of an application.
-    #[inline(always)]
+    ///
+    /// # Panics
+    ///
+    /// This method will likely panic if given `u64` values that did not come from [`Entity::to_bits`].
+    #[inline]
     pub const fn from_bits(bits: u64) -> Self {
-        let high = IdentifierMask::get_high(bits);
-        let kind = IdentifierMask::extract_kind_from_high(high) as u8;
+        // Construct an Identifier initially to extract the kind from.
+        let id = Self::try_from_bits(bits);
 
-        // Needed to convert the kind to an integer so that the comparison can be const
-        assert!(
-            kind == IdKind::Entity as u8,
-            "Attempted to initialise invalid bits as an entity"
-        );
+        match id {
+            Ok(entity) => entity,
+            Err(_) => panic!("Attempted to initialise invalid bits as an entity"),
+        }
+    }
 
-        Self {
-            index: IdentifierMask::get_low(bits),
-            generation: high,
+    /// Reconstruct an `Entity` previously destructured with [`Entity::to_bits`].
+    ///
+    /// Only useful when applied to results from `to_bits` in the same instance of an application.
+    /// 
+    /// This method is the fallible counterpart to [`Entity::from_bits`].
+    #[inline(always)]
+    pub const fn try_from_bits(bits: u64) -> Result<Self, IdentifierError> {
+        let id = Identifier::from_bits(bits);
+        let kind = id.kind();
+
+        match kind {
+            IdKind::Entity => Ok(Self {
+                index: id.low(),
+                generation: id.high(),
+            }),
+            _ => Err(IdentifierError::InvalidEntityId(bits)),
         }
     }
 
@@ -307,23 +324,16 @@ impl Entity {
 impl TryFrom<Identifier> for Entity {
     type Error = IdentifierError;
 
+    #[inline]
     fn try_from(value: Identifier) -> Result<Self, Self::Error> {
-        let kind = IdentifierMask::extract_kind_from_high(value.high());
-
-        if kind == IdKind::Entity {
-            Ok(Entity {
-                index: value.low(),
-                generation: value.high(),
-            })
-        } else {
-            Err(IdentifierError::InvalidEntityId(value.to_bits()))
-        }
+        Self::try_from_bits(value.to_bits())
     }
 }
 
 impl From<Entity> for Identifier {
+    #[inline]
     fn from(value: Entity) -> Self {
-        Identifier::new(value.index(), value.generation(), IdKind::Entity)
+        Identifier::from_bits(value.to_bits())
     }
 }
 
@@ -343,9 +353,7 @@ impl<'de> Deserialize<'de> for Entity {
     {
         let id: u64 = serde::de::Deserialize::deserialize(deserializer)?;
 
-        // Try converting from an `Identifier` as this will allow us to determine
-        // if it is a valid entity or not
-        Entity::try_from(Identifier::from_bits(id)).map_err(serde::de::Error::custom)
+        Entity::try_from_bits(id).map_err(serde::de::Error::custom)
     }
 }
 

crates/bevy_ecs/src/identifier/kinds.rs

@@ -3,7 +3,7 @@
 /// of the ID.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 #[repr(u8)]
-pub(crate) enum IdKind {
+pub enum IdKind {
     /// An ID variant that is compatible with [`crate::entity::Entity`].
     Entity = 0,
     /// A future ID variant.

crates/bevy_ecs/src/identifier/mod.rs

@@ -4,6 +4,7 @@
 //! [`Identifier`]s cannot be created directly, only able to be converted from other
 //! compatible IDs.
 use self::{kinds::IdKind, masks::IdentifierMask};
+use std::hash::Hash;
 
 pub mod error;
 pub(crate) mod kinds;
@@ -13,16 +14,24 @@ pub(crate) mod masks;
 /// Has the same size as a `u64` integer, but the layout is split between a 32-bit low
 /// segment, a 31-bit high segment, and the significant bit reserved as type flags to denote
 /// entity kinds.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Debug, Clone, Copy)]
+// Alignment repr necessary to allow LLVM to better output
+// optimised codegen for `to_bits`, `PartialEq` and `Ord`.
+#[repr(C, align(8))]
 pub struct Identifier {
+    // Do not reorder the fields here. The ordering is explicitly used by repr(C)
+    // to make this struct equivalent to a u64.
+    #[cfg(target_endian = "little")]
     low: u32,
     high: u32,
+    #[cfg(target_endian = "big")]
+    low: u32,
 }
 
 impl Identifier {
     /// Construct a new [`Identifier`]. The `high` parameter is masked with the
     /// `kind` so to pack the high value and bit flags into the same field.
-    #[inline]
+    #[inline(always)]
     #[must_use]
     pub(crate) const fn new(low: u32, high: u32, kind: IdKind) -> Self {
         // the high bits are masked to cut off the most significant bit
@@ -38,34 +47,98 @@ impl Identifier {
     }
 
     /// Returns the value of the low segment of the [`Identifier`].
-    #[inline]
-    pub(crate) const fn low(self) -> u32 {
+    #[inline(always)]
+    pub const fn low(self) -> u32 {
         self.low
     }
 
     /// Returns the value of the high segment of the [`Identifier`]. This
     /// does not apply any masking.
-    #[inline]
-    pub(crate) const fn high(self) -> u32 {
+    #[inline(always)]
+    pub const fn high(self) -> u32 {
         self.high
     }
 
+    /// Returns the masked value of the high segment of the [`Identifier`].
+    /// Does not include the flag bits.
+    #[inline(always)]
+    pub const fn masked_high(self) -> u32 {
+        IdentifierMask::extract_value_from_high(self.high())
+    }
+
+    /// Returns the kind of [`Identifier`] from the high segment.
+    #[inline(always)]
+    pub const fn kind(self) -> IdKind {
+        IdentifierMask::extract_kind_from_high(self.high())
+    }
+
     /// Convert the [`Identifier`] into a `u64`.
-    #[inline]
-    pub(crate) const fn to_bits(self) -> u64 {
+    #[inline(always)]
+    pub const fn to_bits(self) -> u64 {
         IdentifierMask::pack_into_u64(self.low, self.high)
     }
 
     /// Convert a `u64` into an [`Identifier`].
-    #[inline]
-    pub(crate) const fn from_bits(value: u64) -> Self {
+    #[inline(always)]
+    pub const fn from_bits(value: u64) -> Self {
         Self {
             low: IdentifierMask::get_low(value),
             high: IdentifierMask::get_high(value),
         }
     }
 }
 
+// By not short-circuiting in comparisons, we get better codegen.
+// See <https://github.com/rust-lang/rust/issues/117800>
+impl PartialEq for Identifier {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        // By using `to_bits`, the codegen can be optimised out even
+        // further potentially. Relies on the correct alignment/field
+        // order of `Entity`.
+        self.to_bits() == other.to_bits()
+    }
+}
+
+impl Eq for Identifier {}
+
+// The derive macro codegen output is not optimal and can't be optimised as well
+// by the compiler. This impl resolves the issue of non-optimal codegen by relying
+// on comparing against the bit representation of `Entity` instead of comparing
+// the fields. The result is then LLVM is able to optimise the codegen for Entity
+// far beyond what the derive macro can.
+// See <https://github.com/rust-lang/rust/issues/106107>
+impl PartialOrd for Identifier {
+    #[inline]
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        // Make use of our `Ord` impl to ensure optimal codegen output
+        Some(self.cmp(other))
+    }
+}
+
+// The derive macro codegen output is not optimal and can't be optimised as well
+// by the compiler. This impl resolves the issue of non-optimal codegen by relying
+// on comparing against the bit representation of `Entity` instead of comparing
+// the fields. The result is then LLVM is able to optimise the codegen for Entity
+// far beyond what the derive macro can.
+// See <https://github.com/rust-lang/rust/issues/106107>
+impl Ord for Identifier {
+    #[inline]
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        // This will result in better codegen for ordering comparisons, plus
+        // avoids pitfalls with regards to macro codegen relying on property
+        // position when we want to compare against the bit representation.
+        self.to_bits().cmp(&other.to_bits())
+    }
+}
+
+impl Hash for Identifier {
+    #[inline]
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.to_bits().hash(state);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -100,4 +173,36 @@ mod tests {
             IdKind::Entity
         );
     }
+
+    #[rustfmt::skip]
+    #[test]
+    fn id_comparison() {
+        // This is intentionally testing `lt` and `ge` as separate functions.
+        #![allow(clippy::nonminimal_bool)]
+
+        assert!(Identifier::new(123, 456, IdKind::Entity) == Identifier::new(123, 456, IdKind::Entity));
+        assert!(Identifier::new(123, 456, IdKind::Placeholder) == Identifier::new(123, 456, IdKind::Placeholder));
+        assert!(Identifier::new(123, 789, IdKind::Entity) != Identifier::new(123, 456, IdKind::Entity));
+        assert!(Identifier::new(123, 456, IdKind::Entity) != Identifier::new(123, 789, IdKind::Entity));
+        assert!(Identifier::new(123, 456, IdKind::Entity) != Identifier::new(456, 123, IdKind::Entity));
+        assert!(Identifier::new(123, 456, IdKind::Entity) != Identifier::new(123, 456, IdKind::Placeholder));
+
+        // ordering is by flag then high then by low
+
+        assert!(Identifier::new(123, 456, IdKind::Entity) >= Identifier::new(123, 456, IdKind::Entity));
+        assert!(Identifier::new(123, 456, IdKind::Entity) <= Identifier::new(123, 456, IdKind::Entity));
+        assert!(!(Identifier::new(123, 456, IdKind::Entity) < Identifier::new(123, 456, IdKind::Entity)));
+        assert!(!(Identifier::new(123, 456, IdKind::Entity) > Identifier::new(123, 456, IdKind::Entity)));
+
+        assert!(Identifier::new(9, 1, IdKind::Entity) < Identifier::new(1, 9, IdKind::Entity));
+        assert!(Identifier::new(1, 9, IdKind::Entity) > Identifier::new(9, 1, IdKind::Entity));
+
+        assert!(Identifier::new(9, 1, IdKind::Entity) < Identifier::new(9, 1, IdKind::Placeholder));
+        assert!(Identifier::new(1, 9, IdKind::Placeholder) > Identifier::new(1, 9, IdKind::Entity));
+
+        assert!(Identifier::new(1, 1, IdKind::Entity) < Identifier::new(2, 1, IdKind::Entity));
+        assert!(Identifier::new(1, 1, IdKind::Entity) <= Identifier::new(2, 1, IdKind::Entity));
+        assert!(Identifier::new(2, 2, IdKind::Entity) > Identifier::new(1, 2, IdKind::Entity));
+        assert!(Identifier::new(2, 2, IdKind::Entity) >= Identifier::new(1, 2, IdKind::Entity));
+    }
 }