Refactor _EntityPool to use generations and free slots

src/batch_ops.jl

@@ -966,17 +966,17 @@ end
             if isempty(table)
                 continue
             end
-            for entity in table.entities
-                $(world_has_rel ?
-                  :(
+            $(world_has_rel ?
+              :(
+                for entity in table.entities
                     if world_state._targets[entity._id]
                         push!(cleanup, entity)
                     end
-                ) :
-                  (:(nothing))
-                )
-                _recycle(world_state._entity_pool, entity)
-            end
+                end
+            ) :
+              (:(nothing))
+            )
+            _recycle_entities!(world_state._entity_pool, table.entities._data)
             empty!(table)
             for comp in world_state._archetypes[table.archetype].components
                 _clear_component_data!(world_state, stores, comp, table.id)

src/pool.jl

@@ -1,71 +1,80 @@
 
 mutable struct _EntityPool
-    const entities::Vector{Entity}
-    next::Int
+    const gens::Vector{UInt32}
+    const free::Vector{UInt32}
 end
 
 function _EntityPool(cap::UInt32)
-    v = [_new_entity(UInt32(0), typemax(UInt32))]
-    sizehint!(v, cap)
+    gens = UInt32[typemax(UInt32)]
+    sizehint!(gens, cap)
 
-    return _EntityPool(v, 0)
+    return _EntityPool(gens, UInt32[])
 end
 
-function _get_entity(p::_EntityPool)::Entity
-    if p.next == 0
+@inline function _get_entity(p::_EntityPool)::Entity
+    if isempty(p.free)
         return _get_new_entity(p)
     end
-    curr = p.next
-    temp = p.entities[curr]
-
-    p.next = temp._id
-    entity = _Entity(curr % UInt32, temp._gen)
-    p.entities[curr] = entity
-
-    return entity
+    id = pop!(p.free)
+    @inbounds gen = p.gens[id]
+    return _Entity(id, gen)
 end
 
-function _get_new_entity(p::_EntityPool)::Entity
-    e = _new_entity(length(p.entities) + 1, 0)
-    push!(p.entities, e)
+@inline function _get_new_entity(p::_EntityPool)::Entity
+    e = _new_entity(length(p.gens) + 1, 0)
+    push!(p.gens, UInt32(0))
     return e
 end
 
-function _get_pending_entity(p::_EntityPool)::Entity
+@inline function _get_pending_entity(p::_EntityPool)::Entity
     entity = _get_entity(p)
     return _new_entity(entity._id, entity._gen + UInt32(1))
 end
 
-function _activate_entity!(p::_EntityPool, e::Entity)
-    @inbounds p.entities[e._id] = e
+@inline function _activate_entity!(p::_EntityPool, e::Entity)
+    @inbounds p.gens[e._id] = e._gen
     return nothing
 end
 
 function _get_new_entities!(p::_EntityPool, n::Integer)
-    old_len = length(p.entities)
+    old_len = length(p.gens)
     new_len = old_len + n
-    resize!(p.entities, new_len)
+    resize!(p.gens, new_len)
     for i in (old_len+1):new_len
-        @inbounds p.entities[i] = _new_entity(i % UInt32, UInt32(0))
+        @inbounds p.gens[i] = UInt32(0)
     end
     return
 end
 
-function _recycle(p::_EntityPool, e::Entity)
+@inline function _recycle(p::_EntityPool, e::Entity)
     if e._id < 2
         throw(ArgumentError("can't recycle the reserved zero entity"))
     end
-    temp = p.next
-    p.next = e._id
-    p.entities[e._id] = _new_entity(temp % UInt32, e._gen + UInt32(1))
+    push!(p.free, e._id)
+    @inbounds p.gens[e._id] = e._gen + UInt32(1)
     return nothing
 end
 
-function _is_alive(p::_EntityPool, e::Entity)::Bool
-    @inbounds return e._gen == p.entities[e._id]._gen
+@inline function _recycle_entities!(p::_EntityPool, entities::Vector{Entity})
+    n = length(entities)
+    n == 0 && return nothing
+    free = p.free
+    old_len = length(free)
+    resize!(free, old_len + n)
+    @inbounds for i in 1:n
+        e = entities[i]
+        free[old_len+i] = e._id
+        p.gens[e._id] = e._gen + UInt32(1)
+    end
+    return nothing
+end
+
+@inline function _is_alive(p::_EntityPool, e::Entity)::Bool
+    @inbounds return e._gen == p.gens[e._id]
 end
 
 function _reset!(p::_EntityPool)
-    resize!(p.entities, 1)
-    p.next = 0
+    resize!(p.gens, 1)
+    empty!(p.free)
+    return
 end

src/world.jl

@@ -1836,30 +1836,34 @@ end
 
         i = old_length + 1
         # Pop from free list
-        @inbounds while i <= new_length && pool.next != 0
-            entity = _get_entity(pool)
-            entities[i] = entity
-            id = Int(entity._id)
-            state._entities[id] = _EntityIndex(table_index, UInt32(i))
-            $(world_has_rel ? :(state._targets[id] = false) : (:(nothing)))
-            i += 1
+        nfree = length(pool.free)
+        take = min(nfree, new_length - i + 1)
+        if take > 0
+            @inbounds for k in 0:(take-1)
+                id = Int(pool.free[nfree-k])
+                entities[i] = _Entity(UInt32(id), pool.gens[id])
+                state._entities[id] = _EntityIndex(table_index, UInt32(i))
+                $(world_has_rel ? :(state._targets[id] = false) : (:(nothing)))
+                i += 1
+            end
+            resize!(pool.free, nfree - take)
         end
 
         # Bulk-allocate the rest
         if i <= new_length
             rem = new_length - i + 1
-            old_pool_len = length(pool.entities)
+            old_pool_len = length(pool.gens)
             @check old_pool_len == length(state._entities)
             _get_new_entities!(pool, rem)
 
-            new_pool_len = length(pool.entities)
+            new_pool_len = length(pool.gens)
             resize!(state._entities, new_pool_len)
             $(world_has_rel ? :(resize!(state._targets, new_pool_len)) : nothing)
             $(world_has_rel ? :(view(state._targets, (old_pool_len+1):new_pool_len) .= false) : nothing)
 
             @inbounds @simd for j in 1:rem
                 id = old_pool_len + j
-                entity = pool.entities[id]
+                entity = _new_entity(id, 0)
                 entities[i] = entity
                 state._entities[id] = _EntityIndex(table_index, UInt32(i))
                 i += 1

test/include_internals.jl

@@ -11,7 +11,7 @@ using Ark: _Mask, _Not
 using Ark: _get_bit, _contains_all, _contains_any, _and, _or, _clear_bits, _active_bit_indices
 using Ark: _is_zero, _is_not_zero
 using Ark: _MutableMask, _get_bit, _set_bit!, _clear_bit!, _equals, _clear_mask!
-using Ark: _EntityPool, _get_entity, _recycle, _is_alive
+using Ark: _EntityPool, _get_entity, _get_pending_entity, _activate_entity!, _recycle, _is_alive
 using Ark: _Lock, _lock, _unlock, _is_locked
 using Ark: _VecMap, _get_map, _set_map!
 using Ark: _Linear_Map, _LOAD_FACTOR

test/test_pool.jl

@@ -4,17 +4,17 @@
     pool = _EntityPool(initialCap)
 
     @test isa(pool, _EntityPool)
-    @test length(pool.entities) == 1
-    @test all(e -> e._gen == typemax(UInt32), pool.entities)
-    @test pool.next == 0
+    @test length(pool.gens) == 1
+    @test pool.gens[1] == typemax(UInt32)
+    @test isempty(pool.free)
 end
 
 @testset "_EntityPool logic" begin
     # Setup
-    pool = _EntityPool(UInt32(10))  # creates 2 reserved entities
+    pool = _EntityPool(UInt32(10))
 
-    @test length(pool.entities) == 1
-    @test pool.next == 0
+    @test length(pool.gens) == 1
+    @test isempty(pool.free)
 
     @test _is_alive(pool, zero_entity) == false
 
@@ -23,24 +23,25 @@ end
     @test isa(e1, Entity)
     @test e1._id == 2
     @test e1._gen == 0
-    @test length(pool.entities) == 2
+    @test length(pool.gens) == 2
+    @test isempty(pool.free)
 
     # Test _get_entity again
     e2 = _get_entity(pool)
     @test e2._id == 3
     @test e2._gen == 0
-    @test length(pool.entities) == 3
+    @test length(pool.gens) == 3
 
     # Test _recycle with non-reserved entity
     _recycle(pool, e1)
-    @test pool.next == e1._id
-    @test pool.entities[e1._id]._gen == e1._gen + 1
+    @test pool.free == [e1._id]
+    @test pool.gens[e1._id] == e1._gen + 1
 
     # Test _get_entity now uses recycled entity
     e3 = _get_entity(pool)
     @test e3._id == e1._id
     @test e3._gen == e1._gen + 1
-    @test pool.next == 0
+    @test isempty(pool.free)
 
     # Test _alive
     @test _is_alive(pool, e2) == true
@@ -50,3 +51,20 @@ end
     # Test _recycle throws on reserved entity
     @test_throws "ArgumentError: can't recycle the reserved zero entity" _recycle(pool, zero_entity)
 end
+
+@testset "_EntityPool pending entities" begin
+    pool = _EntityPool(UInt32(10))
+
+    e1 = _get_entity(pool)
+    pending = _get_pending_entity(pool)
+
+    @test pending._gen == e1._gen + 1
+    @test _is_alive(pool, e1) == true
+    @test _is_alive(pool, pending) == false
+
+    _activate_entity!(pool, pending)
+    @test _is_alive(pool, pending) == true
+
+    _recycle(pool, pending)
+    @test _is_alive(pool, pending) == false
+end

test/test_world.jl

@@ -1854,7 +1854,7 @@ end
     reset!(world)
 
     @test length(_state(world)._entities) == 1
-    @test length(_state(world)._entity_pool.entities) == 1
+    @test length(_state(world)._entity_pool.gens) == 1
     @test length(_state(world)._cache.filters) == 0
     @test filter._filter.id[] == 0
     @test length(filter._filter.tables) == 0