Optimize Entity::eq

benches/Cargo.toml

@@ -61,3 +61,8 @@ harness = false
 name = "bezier"
 path = "benches/bevy_math/bezier.rs"
 harness = false
+
+[[bench]]
+name = "utils"
+path = "benches/bevy_utils/entity_hash.rs"
+harness = false

benches/benches/bevy_utils/entity_hash.rs

@@ -0,0 +1,75 @@
+use bevy_ecs::entity::Entity;
+use bevy_utils::EntityHashSet;
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha8Rng;
+
+criterion_group!(benches, entity_set_build_and_lookup,);
+criterion_main!(benches);
+
+const SIZES: [usize; 5] = [100, 316, 1000, 3162, 10000];
+
+fn make_entity(rng: &mut impl Rng, size: usize) -> Entity {
+    // -log₂(1-x) gives an exponential distribution with median 1.0
+    // That lets us get values that are mostly small, but some are quite large
+    // * For ids, half are in [0, size), half are unboundedly larger.
+    // * For generations, half are in [0, 2), half are unboundedly larger.
+
+    let x: f64 = rng.gen();
+    let id = -(1.0 - x).log2() * (size as f64);
+    let x: f64 = rng.gen();
+    let gen = -(1.0 - x).log2() * 2.0;
+
+    // this is not reliable, but we're internal so a hack is ok
+    let bits = ((gen as u64) << 32) | (id as u64);
+    let e = Entity::from_bits(bits);
+    assert_eq!(e.index(), id as u32);
+    assert_eq!(e.generation(), gen as u32);
+    e
+}
+
+fn entity_set_build_and_lookup(c: &mut Criterion) {
+    let mut group = c.benchmark_group("entity_hash");
+    for size in SIZES {
+        // Get some random-but-consistent entities to use for all the benches below.
+        let mut rng = ChaCha8Rng::seed_from_u64(size as u64);
+        let entities = Vec::from_iter(
+            std::iter::repeat_with(|| make_entity(&mut rng, size)).take(size),
+        );
+
+        group.throughput(Throughput::Elements(size as u64));
+        group.bench_function(
+            BenchmarkId::new("entity_set_build", size),
+            |bencher| {
+                bencher.iter_with_large_drop(|| EntityHashSet::from_iter(entities.iter().copied()));
+            },
+        );
+        group.bench_function(
+            BenchmarkId::new("entity_set_lookup_hit", size),
+            |bencher| {
+                let set = EntityHashSet::from_iter(entities.iter().copied());
+                bencher.iter(|| entities.iter().copied().filter(|e| set.contains(e)).count());
+            },
+        );
+        group.bench_function(
+            BenchmarkId::new("entity_set_lookup_miss_id", size),
+            |bencher| {
+                let set = EntityHashSet::from_iter(entities.iter().copied());
+                bencher.iter(|| entities.iter()
+                    .copied()
+                    .map(|e| Entity::from_bits(e.to_bits() + 1))
+                    .filter(|e| set.contains(e)).count());
+            },
+        );
+        group.bench_function(
+            BenchmarkId::new("entity_set_lookup_miss_gen", size),
+            |bencher| {
+                let set = EntityHashSet::from_iter(entities.iter().copied());
+                bencher.iter(|| entities.iter()
+                    .copied()
+                    .map(|e| Entity::from_bits(e.to_bits() + (1 << 32)))
+                    .filter(|e| set.contains(e)).count());
+            },
+        );
+    }
+}

crates/bevy_ecs/src/entity/mod.rs

@@ -115,12 +115,21 @@ type IdCursor = isize;
 /// [`EntityCommands`]: crate::system::EntityCommands
 /// [`Query::get`]: crate::system::Query::get
 /// [`World`]: crate::world::World
-#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd)]
+#[derive(Clone, Copy, Eq, Ord, PartialOrd)]
 pub struct Entity {
     generation: u32,
     index: u32,
 }
 
+// By not short-circuiting in comparisons, we get better codegen.
+// See <https://github.com/rust-lang/rust/issues/117800>
+impl PartialEq for Entity {
+    #[inline]
+    fn eq(&self, other: &Entity) -> bool {
+        (self.generation == other.generation) & (self.index == other.index)
+    }
+}
+
 impl Hash for Entity {
     #[inline]
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
@@ -917,4 +926,30 @@ mod tests {
         assert_eq!(next_entity.index(), entity.index());
         assert!(next_entity.generation > entity.generation + GENERATIONS);
     }
+
+    #[test]
+    fn entity_comparison() {
+        // This is intentionally testing `lt` and `ge` as separate functions.
+        #![allow(clippy::nonminimal_bool)]
+
+        assert!(Entity::new(123, 456) == Entity::new(123, 456));
+        assert!(Entity::new(123, 789) != Entity::new(123, 456));
+        assert!(Entity::new(123, 456) != Entity::new(123, 789));
+        assert!(Entity::new(123, 456) != Entity::new(456, 123));
+
+        // ordering is by generation then by index
+
+        assert!(Entity::new(123, 456) >= Entity::new(123, 456));
+        assert!(Entity::new(123, 456) <= Entity::new(123, 456));
+        assert!(!(Entity::new(123, 456) < Entity::new(123, 456)));
+        assert!(!(Entity::new(123, 456) > Entity::new(123, 456)));
+
+        assert!(Entity::new(9, 1) < Entity::new(1, 9));
+        assert!(Entity::new(1, 9) > Entity::new(9, 1));
+
+        assert!(Entity::new(1, 1) < Entity::new(2, 1));
+        assert!(Entity::new(1, 1) <= Entity::new(2, 1));
+        assert!(Entity::new(2, 2) > Entity::new(1, 2));
+        assert!(Entity::new(2, 2) >= Entity::new(1, 2));
+    }
 }