Migrate to the new Entity layout

For details see bevyengine/bevy#19121
bevyengine/bevy#18704

The niche is now in the index, which makes the compression logic even
simpler.

Author: Shatur

src/compact_entity.rs

@@ -17,32 +17,34 @@
 
 use core::fmt::{self, Formatter};
 
-use bevy::prelude::*;
+use bevy::{
+    ecs::entity::{EntityGeneration, EntityRow},
+    prelude::*,
+};
 use serde::{
     Deserializer, Serialize, Serializer,
     de::{self, SeqAccess, Visitor},
 };
 
 /// Serializes an entity by writing its index and generation as separate numbers.
 ///
-/// This reduces the space used for serializers with varint encoding.
+/// This reduces the space required when using serializers with varint encoding.
 ///
-/// The index is first prepended with a bit flag to indicate if the generation
-/// is serialized or not. It is not serialized if <= 1; note that generations are [`NonZeroU32`](core::num::NonZeroU32)
-/// and a value of zero is used in [`Option<Entity>`] to signify [`None`], so generation 1 is the first
-/// generation.
+/// Since the index can never be [`u32::MAX`], we reuse that extra niche to indicate
+/// whether the generation isn't [`EntityGeneration::FIRST`]. If it doesn't,
+/// the generation is skipped during serialization.
 ///
 /// See also [`deserialize`] and [`postcard_utils::entity_to_extend_mut`](crate::postcard_utils::entity_to_extend_mut).
 pub fn serialize<S: Serializer>(entity: &Entity, serializer: S) -> Result<S::Ok, S::Error> {
-    let mut flagged_index = (entity.index() as u64) << 1;
-    let flag = entity.generation() > 1;
-    flagged_index |= flag as u64;
+    let mut index = entity.index() << 1;
+    let has_generation = entity.generation() != EntityGeneration::FIRST;
+    index |= has_generation as u32;
 
-    if flag {
-        let generation = entity.generation() - 1;
-        (flagged_index, generation).serialize(serializer)
+    if has_generation {
+        let generation = entity.generation().to_bits();
+        (index, generation).serialize(serializer)
     } else {
-        flagged_index.serialize(serializer)
+        index.serialize(serializer)
     }
 }
 
@@ -63,22 +65,23 @@ impl<'de> Visitor<'de> for EntityVisitor {
     }
 
     fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
-        let flagged_index: u64 = seq
+        let index: u32 = seq
             .next_element()?
             .ok_or_else(|| de::Error::invalid_length(0, &self))?;
 
-        let has_generation = (flagged_index & 1) != 0;
+        let has_generation = (index & 1) != 0;
 
         let generation = if has_generation {
-            let generation: u32 = seq
-                .next_element()?
-                .ok_or_else(|| de::Error::invalid_length(1, &self))?;
-            generation as u64 + 1
+            seq.next_element()?
+                .ok_or_else(|| de::Error::invalid_length(1, &self))?
         } else {
-            1
+            0
         };
 
-        let bits = (flagged_index >> 1) | (generation << 32);
-        Ok(Entity::from_bits(bits))
+        // SAFETY: `index` is non-max after shift.
+        let row = unsafe { EntityRow::from_raw_u32(index >> 1).unwrap_unchecked() };
+        let generation = EntityGeneration::from_bits(generation);
+
+        Ok(Entity::from_row_and_generation(row, generation))
     }
 }

src/postcard_utils.rs

@@ -212,11 +212,13 @@ impl<'a, T: Buf> DeFlavor<'a> for BufFlavor<'a, T> {
 
 #[cfg(test)]
 mod tests {
+    use bevy::ecs::entity::{EntityGeneration, EntityRow};
+
     use super::*;
 
     #[test]
     fn entity_without_generation() {
-        let expected_entity = Entity::from_raw(1);
+        let expected_entity = Entity::from_raw_u32(1).unwrap();
         let mut buffer = Vec::new();
         entity_to_extend_mut(&expected_entity, &mut buffer).unwrap();
         assert_eq!(buffer.len(), 1);
@@ -227,7 +229,10 @@ mod tests {
 
     #[test]
     fn entity_with_generation() {
-        let expected_entity = Entity::from_bits(1 | (2 << 32));
+        let expected_entity = Entity::from_row_and_generation(
+            EntityRow::from_raw_u32(1).unwrap(),
+            EntityGeneration::from_bits(1),
+        );
         let mut buffer = Vec::new();
         entity_to_extend_mut(&expected_entity, &mut buffer).unwrap();
         assert_eq!(buffer.len(), 2);

src/shared/server_entity_map.rs

@@ -155,8 +155,8 @@ mod tests {
 
     #[test]
     fn mapping() {
-        const SERVER_ENTITY: Entity = Entity::from_raw(0);
-        const CLIENT_ENTITY: Entity = Entity::from_raw(1);
+        const SERVER_ENTITY: Entity = Entity::from_raw_u32(0).unwrap();
+        const CLIENT_ENTITY: Entity = Entity::from_raw_u32(1).unwrap();
 
         let mut map = ServerEntityMap::default();
         assert_eq!(map.server_entry(SERVER_ENTITY).get(), None);

tests/server_event.rs

@@ -187,7 +187,8 @@ fn local_resending() {
     .add_server_event::<TestEvent>(Channel::Ordered)
     .finish();
 
-    const PLACEHOLDER_CLIENT_ID: ClientId = ClientId::Client(Entity::from_raw(1));
+    const CLIENT_ENTITY: Entity = Entity::from_raw_u32(1).unwrap();
+    const PLACEHOLDER_CLIENT_ID: ClientId = ClientId::Client(CLIENT_ENTITY);
     for (mode, events_count) in [
         (SendMode::Broadcast, 1),
         (SendMode::Direct(ClientId::Server), 1),