Make Entity construction more ergonomic (#21246)

# Objective

Entity serialization is necessary for networking. Entities can exist
inside components and events. After deserialization, we simply map
remote entities to local entities.

To serialize entities efficiently, we split them into index and
generation, which benefits from varint serialization. #19121 and #18704
changed the entity layout, and I like the new layout a lot. We can now
use the extra bit from the index to store whether the generation is zero
or not, avoiding the need to serialize the generation entirely.

However, constructing new entities requires relying on the internal
layout, which is not very ergonomic. For example, here is how an entity
with index = 1 and generation = 1 can be created:

```rust
let expected_entity = Entity::from_bits((1 ^ u32::MAX) as u64 | (1 << 32));
```

## Solution

- Make `Entity::from_raw_and_generation` public. While at it, I also
removed outdated comment.
- Add `EntityRow::from_raw_u32` to make the initialization nicer.

## Testing

- It's a trivial change, but I re-used `EntityRow::from_raw_u32` in unit
tests to simplify them.

## Notes

I'd probably rename `Entity::from_raw_and_generation` into
`Entity::from_row_and_generation` or
`Entity::from_index_and_generation`.

Author: Shatur

crates/bevy_ecs/src/entity/mod.rs

@@ -125,6 +125,16 @@ impl EntityRow {
         Self(index)
     }
 
+    /// Equivalent to [`new`](Self::new) except that it takes a `u32` instead of a `NonMaxU32`.
+    ///
+    /// Returns `None` if the index is `u32::MAX`.
+    pub const fn from_raw_u32(index: u32) -> Option<Self> {
+        match NonMaxU32::new(index) {
+            Some(index) => Some(Self(index)),
+            None => None,
+        }
+    }
+
     /// Gets the index of the entity.
     #[inline(always)]
     pub const fn index(self) -> u32 {
@@ -445,13 +455,9 @@ impl Hash for Entity {
 }
 
 impl Entity {
-    /// Construct an [`Entity`] from a raw `row` value and a non-zero `generation` value.
-    /// Ensure that the generation value is never greater than `0x7FFF_FFFF`.
+    /// Constructs an [`Entity`] from a raw `row` value and a non-zero `generation` value.
     #[inline(always)]
-    pub(crate) const fn from_raw_and_generation(
-        row: EntityRow,
-        generation: EntityGeneration,
-    ) -> Entity {
+    pub const fn from_raw_and_generation(row: EntityRow, generation: EntityGeneration) -> Entity {
         Self { row, generation }
     }
 
@@ -1324,12 +1330,11 @@ mod tests {
 
     #[test]
     fn entity_bits_roundtrip() {
-        let r = EntityRow::new(NonMaxU32::new(0xDEADBEEF).unwrap());
+        let r = EntityRow::from_raw_u32(0xDEADBEEF).unwrap();
         assert_eq!(EntityRow::from_bits(r.to_bits()), r);
 
-        // Generation cannot be greater than 0x7FFF_FFFF else it will be an invalid Entity id
         let e = Entity::from_raw_and_generation(
-            EntityRow::new(NonMaxU32::new(0xDEADBEEF).unwrap()),
+            EntityRow::from_raw_u32(0xDEADBEEF).unwrap(),
             EntityGeneration::from_bits(0x5AADF00D),
         );
         assert_eq!(Entity::from_bits(e.to_bits()), e);
@@ -1368,15 +1373,15 @@ mod tests {
 
     #[test]
     fn entity_const() {
-        const C1: Entity = Entity::from_raw(EntityRow::new(NonMaxU32::new(42).unwrap()));
+        const C1: Entity = Entity::from_raw(EntityRow::from_raw_u32(42).unwrap());
         assert_eq!(42, C1.index());
         assert_eq!(0, C1.generation().to_bits());
 
         const C2: Entity = Entity::from_bits(0x0000_00ff_0000_00cc);
         assert_eq!(!0x0000_00cc, C2.index());
         assert_eq!(0x0000_00ff, C2.generation().to_bits());
 
-        const C3: u32 = Entity::from_raw(EntityRow::new(NonMaxU32::new(33).unwrap())).index();
+        const C3: u32 = Entity::from_raw(EntityRow::from_raw_u32(33).unwrap()).index();
         assert_eq!(33, C3);
 
         const C4: u32 = Entity::from_bits(0x00dd_00ff_1111_1111)
@@ -1421,41 +1426,41 @@ mod tests {
     fn entity_comparison() {
         assert_eq!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ),
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             )
         );
         assert_ne!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(789)
             ),
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             )
         );
         assert_ne!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ),
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(789)
             )
         );
         assert_ne!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ),
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(456).unwrap()),
+                EntityRow::from_raw_u32(456).unwrap(),
                 EntityGeneration::from_bits(123)
             )
         );
@@ -1464,93 +1469,93 @@ mod tests {
 
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ) >= Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             )
         );
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ) <= Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             )
         );
         assert!(
             !(Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ) < Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ))
         );
         assert!(
             !(Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ) > Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(123).unwrap()),
+                EntityRow::from_raw_u32(123).unwrap(),
                 EntityGeneration::from_bits(456)
             ))
         );
 
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(9).unwrap()),
+                EntityRow::from_raw_u32(9).unwrap(),
                 EntityGeneration::from_bits(1)
             ) < Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(1).unwrap()),
+                EntityRow::from_raw_u32(1).unwrap(),
                 EntityGeneration::from_bits(9)
             )
         );
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(1).unwrap()),
+                EntityRow::from_raw_u32(1).unwrap(),
                 EntityGeneration::from_bits(9)
             ) > Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(9).unwrap()),
+                EntityRow::from_raw_u32(9).unwrap(),
                 EntityGeneration::from_bits(1)
             )
         );
 
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(1).unwrap()),
+                EntityRow::from_raw_u32(1).unwrap(),
                 EntityGeneration::from_bits(1)
             ) > Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(2).unwrap()),
+                EntityRow::from_raw_u32(2).unwrap(),
                 EntityGeneration::from_bits(1)
             )
         );
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(1).unwrap()),
+                EntityRow::from_raw_u32(1).unwrap(),
                 EntityGeneration::from_bits(1)
             ) >= Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(2).unwrap()),
+                EntityRow::from_raw_u32(2).unwrap(),
                 EntityGeneration::from_bits(1)
             )
         );
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(2).unwrap()),
+                EntityRow::from_raw_u32(2).unwrap(),
                 EntityGeneration::from_bits(2)
             ) < Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(1).unwrap()),
+                EntityRow::from_raw_u32(1).unwrap(),
                 EntityGeneration::from_bits(2)
             )
         );
         assert!(
             Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(2).unwrap()),
+                EntityRow::from_raw_u32(2).unwrap(),
                 EntityGeneration::from_bits(2)
             ) <= Entity::from_raw_and_generation(
-                EntityRow::new(NonMaxU32::new(1).unwrap()),
+                EntityRow::from_raw_u32(1).unwrap(),
                 EntityGeneration::from_bits(2)
             )
         );
@@ -1564,15 +1569,12 @@ mod tests {
         let hash = EntityHash;
 
         let first_id = 0xC0FFEE << 8;
-        let first_hash = hash.hash_one(Entity::from_raw(EntityRow::new(
-            NonMaxU32::new(first_id).unwrap(),
-        )));
+        let first_hash =
+            hash.hash_one(Entity::from_raw(EntityRow::from_raw_u32(first_id).unwrap()));
 
         for i in 1..=255 {
             let id = first_id + i;
-            let hash = hash.hash_one(Entity::from_raw(EntityRow::new(
-                NonMaxU32::new(id).unwrap(),
-            )));
+            let hash = hash.hash_one(Entity::from_raw(EntityRow::from_raw_u32(id).unwrap()));
             assert_eq!(first_hash.wrapping_sub(hash) as u32, i);
         }
     }
@@ -1584,15 +1586,12 @@ mod tests {
         let hash = EntityHash;
 
         let first_id = 0xC0FFEE;
-        let first_hash = hash.hash_one(Entity::from_raw(EntityRow::new(
-            NonMaxU32::new(first_id).unwrap(),
-        ))) >> 57;
+        let first_hash =
+            hash.hash_one(Entity::from_raw(EntityRow::from_raw_u32(first_id).unwrap())) >> 57;
 
         for bit in 0..u32::BITS {
             let id = first_id ^ (1 << bit);
-            let hash = hash.hash_one(Entity::from_raw(EntityRow::new(
-                NonMaxU32::new(id).unwrap(),
-            ))) >> 57;
+            let hash = hash.hash_one(Entity::from_raw(EntityRow::from_raw_u32(id).unwrap())) >> 57;
             assert_ne!(hash, first_hash);
         }
     }
@@ -1617,7 +1616,7 @@ mod tests {
 
     #[test]
     fn entity_debug() {
-        let entity = Entity::from_raw(EntityRow::new(NonMaxU32::new(42).unwrap()));
+        let entity = Entity::from_raw(EntityRow::from_raw_u32(42).unwrap());
         let string = format!("{entity:?}");
         assert_eq!(string, "42v0");
 
@@ -1628,7 +1627,7 @@ mod tests {
 
     #[test]
     fn entity_display() {
-        let entity = Entity::from_raw(EntityRow::new(NonMaxU32::new(42).unwrap()));
+        let entity = Entity::from_raw(EntityRow::from_raw_u32(42).unwrap());
         let string = format!("{entity}");
         assert_eq!(string, "42v0");
 

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -666,19 +666,18 @@ mod tests {
         storage::SparseSet,
     };
     use alloc::{vec, vec::Vec};
-    use nonmax::NonMaxU32;
 
     #[derive(Debug, Eq, PartialEq)]
     struct Foo(usize);
 
     #[test]
     fn sparse_set() {
         let mut set = SparseSet::<Entity, Foo>::default();
-        let e0 = Entity::from_raw(EntityRow::new(NonMaxU32::new(0).unwrap()));
-        let e1 = Entity::from_raw(EntityRow::new(NonMaxU32::new(1).unwrap()));
-        let e2 = Entity::from_raw(EntityRow::new(NonMaxU32::new(2).unwrap()));
-        let e3 = Entity::from_raw(EntityRow::new(NonMaxU32::new(3).unwrap()));
-        let e4 = Entity::from_raw(EntityRow::new(NonMaxU32::new(4).unwrap()));
+        let e0 = Entity::from_raw(EntityRow::from_raw_u32(0).unwrap());
+        let e1 = Entity::from_raw(EntityRow::from_raw_u32(1).unwrap());
+        let e2 = Entity::from_raw(EntityRow::from_raw_u32(2).unwrap());
+        let e3 = Entity::from_raw(EntityRow::from_raw_u32(3).unwrap());
+        let e4 = Entity::from_raw(EntityRow::from_raw_u32(4).unwrap());
 
         set.insert(e1, Foo(1));
         set.insert(e2, Foo(2));

crates/bevy_ecs/src/storage/table/mod.rs

@@ -859,7 +859,6 @@ mod tests {
         storage::{TableBuilder, TableId, TableRow, Tables},
     };
     use alloc::vec::Vec;
-    use nonmax::NonMaxU32;
 
     #[derive(Component)]
     struct W<T>(T);
@@ -889,7 +888,7 @@ mod tests {
             .add_column(components.get_info(component_id).unwrap())
             .build();
         let entities = (0..200)
-            .map(|index| Entity::from_raw(EntityRow::new(NonMaxU32::new(index).unwrap())))
+            .map(|index| Entity::from_raw(EntityRow::from_raw_u32(index).unwrap()))
             .collect::<Vec<_>>();
         for entity in &entities {
             // SAFETY: we allocate and immediately set data afterwards