Repository: mlange-42/arche
Pull request #35 (2023-02-08 16:38:41)

Optimize query iteration
Change `Query` iterator to avoid allocations as well as indexing into a slice of archetypes.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fae3df7..283fef4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,10 @@
 
 ## [[unpublished]](https://github.com/mlange-42/arche/compare/v0.1.4...main)
 
-Nothing
+### Other
+
+* Optimization of `Query` iteration, avoids allocations and makes it approx. 30% faster (#35)
+* Removed method `Query.Count()`, as it was a by-product of the allocations in the above point (#35)
 
 ## [[v0.1.4]](https://github.com/mlange-42/arche/compare/v0.1.3...v0.1.4)
 
diff --git a/doc.go b/doc.go
index b1f59e4..2645ea7 100644
--- a/doc.go
+++ b/doc.go
@@ -1,4 +1,4 @@
 // Package arche is an archetype-based Entity Component System for Go.
 //
-// See sub-package [ecs] for the API docs.
+// See sub-package [arche/ecs] for the API docs.
 package arche
diff --git a/ecs/archetype.go b/ecs/archetype.go
index 09066f4..0e0f802 100644
--- a/ecs/archetype.go
+++ b/ecs/archetype.go
@@ -15,11 +15,9 @@ type archetype struct {
 
 var entityType = reflect.TypeOf(Entity{})
 
-// newArchetype creates a new Archetype for the given components
-// Component arguments must be sorted by ID!
-func newArchetype(capacityIncrement int, components ...componentType) archetype {
+func (a *archetype) init(capacityIncrement int, components ...componentType) {
 	var mask Mask
-	indices := make([]ID, len(components))
+	a.indices = make([]ID, len(components))
 	comps := [MaskTotalBits]storage{}
 
 	prev := -1
@@ -30,16 +28,15 @@ func newArchetype(capacityIncrement int, components ...componentType) archetype
 		prev = int(c.ID)
 
 		mask.Set(c.ID, true)
-		indices[i] = c.ID
-		comps[c.ID] = newStorage(c.Type, capacityIncrement)
+		a.indices[i] = c.ID
+		comps[c.ID] = storage{}
+		comps[c.ID].init(c.Type, capacityIncrement)
 	}
 
-	return archetype{
-		mask:       mask,
-		indices:    indices,
-		entities:   newStorage(entityType, capacityIncrement),
-		components: comps,
-	}
+	a.mask = mask
+	a.components = comps
+	a.entities = storage{}
+	a.entities.init(entityType, capacityIncrement)
 }
 
 // GetEntity returns the entity at the given index
diff --git a/ecs/query.go b/ecs/query.go
index edfcf04..03c970d 100644
--- a/ecs/query.go
+++ b/ecs/query.go
@@ -8,7 +8,7 @@ import (
 //
 // Create queries through the [World] using [World.Query].
 //
-// ## Example:
+// # Example:
 //
 //	query := world.Query(posID, rotID)
 //	for query.Next() {
@@ -16,66 +16,53 @@ import (
 //	    pos.X += 1.0
 //	}
 type Query struct {
-	world      *World
-	archetypes []archetypeIter
-	index      int
-	done       bool
-	count      int
-	lockBit    uint8
+	world     *World
+	mask      Mask
+	archetype archetypeIter
+	index     int
+	done      bool
+	lockBit   uint8
 }
 
 // newQuery creates a new Query
-func newQuery(world *World, arches []archetypeIter, count int, lockBit uint8) Query {
+func newQuery(world *World, mask Mask, lockBit uint8) Query {
 	return Query{
-		world:      world,
-		archetypes: arches,
-		index:      0,
-		done:       false,
-		count:      count,
-		lockBit:    lockBit,
+		world:   world,
+		mask:    mask,
+		index:   -1,
+		lockBit: lockBit,
 	}
 }
 
 // Next proceeds to the next [Entity] in the Query.
 func (q *Query) Next() bool {
-	if q.done {
-		panic("Query is used up. Create a new Query!")
+	if q.archetype.Next() {
+		return true
 	}
-	for {
-		if q.archetypes[q.index].Next() {
-			return true
-		}
-		q.index++
-		if q.index >= len(q.archetypes) {
-			q.done = true
-			q.world.closeQuery(q)
-			return false
-		}
+	i, a, ok := q.world.nextArchetype(q.mask, q.index)
+	q.index = i
+	if ok {
+		q.archetype = a
+		return true
 	}
+	q.done = true
+	q.world.closeQuery(q)
+	return false
 }
 
 // Has returns whether the current [Entity] has the given component
 func (q *Query) Has(comp ID) bool {
-	if q.done {
-		panic("Query is used up. Create a new Query!")
-	}
-	return q.archetypes[q.index].Has(comp)
+	return q.archetype.Has(comp)
 }
 
 // Get returns the pointer to the given component at the iterator's current [Entity]
 func (q *Query) Get(comp ID) unsafe.Pointer {
-	if q.done {
-		panic("Query is used up. Create a new Query!")
-	}
-	return q.archetypes[q.index].Get(comp)
+	return q.archetype.Get(comp)
 }
 
 // Entity returns the [Entity] at the iterator's position
 func (q *Query) Entity() Entity {
-	if q.done {
-		panic("Query is used up. Create a new Query!")
-	}
-	return q.archetypes[q.index].Entity()
+	return q.archetype.Entity()
 }
 
 // Close closes the Query and unlocks the world.
@@ -87,11 +74,6 @@ func (q *Query) Close() {
 	q.world.closeQuery(q)
 }
 
-// Count returns the number of matching entities
-func (q *Query) Count() int {
-	return q.count
-}
-
 type archetypeIter struct {
 	Archetype *archetype
 	Length    int
@@ -102,7 +84,6 @@ func newArchetypeIter(arch *archetype) archetypeIter {
 	return archetypeIter{
 		Archetype: arch,
 		Length:    int(arch.Len()),
-		Index:     -1,
 	}
 }
 
diff --git a/ecs/storage.go b/ecs/storage.go
index cf951e4..93700ba 100644
--- a/ecs/storage.go
+++ b/ecs/storage.go
@@ -18,21 +18,18 @@ type storage struct {
 }
 
 // newStorage creates a new ReflectStorage
-func newStorage(tp reflect.Type, increment int) storage {
+func (s *storage) init(tp reflect.Type, increment int) {
 	size := tp.Size()
 	align := uintptr(tp.Align())
 	size = (size + (align - 1)) / align * align
 
-	buffer := reflect.New(reflect.ArrayOf(increment, tp)).Elem()
-	return storage{
-		buffer:            buffer,
-		bufferAddress:     buffer.Addr().UnsafePointer(),
-		typeOf:            tp,
-		itemSize:          size,
-		len:               0,
-		cap:               uint32(increment),
-		capacityIncrement: uint32(increment),
-	}
+	s.buffer = reflect.New(reflect.ArrayOf(increment, tp)).Elem()
+	s.bufferAddress = s.buffer.Addr().UnsafePointer()
+	s.typeOf = tp
+	s.itemSize = size
+	s.len = 0
+	s.cap = uint32(increment)
+	s.capacityIncrement = uint32(increment)
 }
 
 // Get retrieves an unsafe pointer to an element
diff --git a/ecs/world.go b/ecs/world.go
index 9d15b25..65fc450 100644
--- a/ecs/world.go
+++ b/ecs/world.go
@@ -19,13 +19,15 @@ func NewWorld() World {
 
 // FromConfig creates a new [World] from a [Config]
 func FromConfig(conf Config) World {
+	arch := archetype{}
+	arch.init(conf.CapacityIncrement)
 	return World{
 		config:     conf,
 		entities:   []entityIndex{{-1, 0}},
 		entityPool: newEntityPool(conf.CapacityIncrement),
 		bitPool:    newBitPool(),
 		registry:   newComponentRegistry(),
-		archetypes: []archetype{newArchetype(conf.CapacityIncrement)},
+		archetypes: []archetype{arch},
 		locks:      Mask(0),
 	}
 }
@@ -51,10 +53,11 @@ func (w *World) NewEntity() Entity {
 
 	entity := w.entityPool.Get()
 	idx := w.archetypes[0].Add(entity)
-	if int(entity.id) == len(w.entities) {
-		if len(w.entities) == cap(w.entities) {
+	len := len(w.entities)
+	if int(entity.id) == len {
+		if len == cap(w.entities) {
 			old := w.entities
-			w.entities = make([]entityIndex, len(w.entities), len(w.entities)+w.config.CapacityIncrement)
+			w.entities = make([]entityIndex, len, len+w.config.CapacityIncrement)
 			copy(w.entities, old)
 		}
 		w.entities = append(w.entities, entityIndex{0, idx})
@@ -224,8 +227,10 @@ func (w *World) createArchetype(comps ...ID) int {
 	for i, id := range comps {
 		types[i] = componentType{id, w.registry.types[id]}
 	}
-	w.archetypes = append(w.archetypes, newArchetype(w.config.CapacityIncrement, types...))
-	return len(w.archetypes) - 1
+	w.archetypes = append(w.archetypes, archetype{})
+	idx := len(w.archetypes) - 1
+	w.archetypes[idx].init(w.config.CapacityIncrement, types...)
+	return idx
 }
 
 // Alive reports whether an entity is still alive.
@@ -238,24 +243,28 @@ func (w *World) componentID(tp reflect.Type) ID {
 	return w.registry.ComponentID(tp)
 }
 
+func (w *World) nextArchetype(mask Mask, index int) (int, archetypeIter, bool) {
+	len := len(w.archetypes)
+	if index >= len {
+		panic("exceeded end of query")
+	}
+	for i := index + 1; i < len; i++ {
+		a := &w.archetypes[i]
+		if a.Len() > 0 && a.mask.Contains(mask) {
+			return i, newArchetypeIter(a), true
+		}
+	}
+	return len, archetypeIter{}, false
+}
+
 // Query creates a [Query] iterator for the given components.
 //
 // Locks the world to prevent changes to component compositions.
 func (w *World) Query(comps ...ID) Query {
 	mask := NewMask(comps...)
-	arches := []archetypeIter{}
-	length := len(w.archetypes)
-	count := 0
-	for i := 0; i < length; i++ {
-		arch := &w.archetypes[i]
-		if arch.mask.Contains(mask) {
-			arches = append(arches, newArchetypeIter(arch))
-			count += int(arch.Len())
-		}
-	}
 	lock := w.bitPool.Get()
 	w.locks.Set(ID(lock), true)
-	return newQuery(w, arches, count, lock)
+	return newQuery(w, mask, lock)
 }
 
 // closeQuery closes a query and unlocks the world
diff --git a/profile/main/profile.go b/profile/main/profile.go
index dcde45a..a4cc3ce 100644
--- a/profile/main/profile.go
+++ b/profile/main/profile.go
@@ -22,7 +22,7 @@ type rotation struct {
 
 func main() {
 
-	count := 100
+	count := 250
 	iters := 10000
 	entities := 1000
 


Test patch:
diff --git a/benchmark/iterate/arche_test.go b/benchmark/iterate/arche_test.go
index 44eee6e..6b6b374 100644
--- a/benchmark/iterate/arche_test.go
+++ b/benchmark/iterate/arche_test.go
@@ -30,6 +30,30 @@ func runArcheQuery(b *testing.B, count int) {
 	}
 }
 
+func runArcheQueryEntity(b *testing.B, count int) {
+	b.StopTimer()
+	world := ecs.NewWorld()
+
+	posID := ecs.ComponentID[position](&world)
+	rotID := ecs.ComponentID[rotation](&world)
+
+	for i := 0; i < count; i++ {
+		entity := world.NewEntity()
+		world.Add(entity, posID, rotID)
+	}
+	b.StartTimer()
+
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		query := world.Query(posID, rotID)
+		b.StartTimer()
+		for query.Next() {
+			e := query.Entity()
+			_ = e
+		}
+	}
+}
+
 func runArcheWorld(b *testing.B, count int) {
 	b.StopTimer()
 	world := ecs.NewWorld()
@@ -67,6 +91,18 @@ func BenchmarkArcheIterQuery_100_000(b *testing.B) {
 	runArcheQuery(b, 100000)
 }
 
+func BenchmarkArcheIterQueryEntity_1_000(b *testing.B) {
+	runArcheQueryEntity(b, 1000)
+}
+
+func BenchmarkArcheIterQueryEntity_10_000(b *testing.B) {
+	runArcheQueryEntity(b, 10000)
+}
+
+func BenchmarkArcheIterQueryEntity_100_000(b *testing.B) {
+	runArcheQueryEntity(b, 100000)
+}
+
 func BenchmarkArcheIterWorld_1_000(b *testing.B) {
 	runArcheWorld(b, 1000)
 }
diff --git a/ecs/archetype_test.go b/ecs/archetype_test.go
index 8519f2b..6aac4bc 100644
--- a/ecs/archetype_test.go
+++ b/ecs/archetype_test.go
@@ -13,7 +13,8 @@ func TestArchetype(t *testing.T) {
 		{ID: 1, Type: reflect.TypeOf(rotation{})},
 	}
 
-	arch := newArchetype(32, comps...)
+	arch := archetype{}
+	arch.init(32, comps...)
 
 	arch.Add(
 		newEntity(0),
@@ -76,38 +77,41 @@ func TestNewArchetype(t *testing.T) {
 		{ID: 0, Type: reflect.TypeOf(position{})},
 		{ID: 1, Type: reflect.TypeOf(rotation{})},
 	}
-	_ = newArchetype(32, comps...)
+	arch := archetype{}
+	arch.init(32, comps...)
 
 	comps = []componentType{
 		{ID: 1, Type: reflect.TypeOf(rotation{})},
 		{ID: 0, Type: reflect.TypeOf(position{})},
 	}
 	assert.Panics(t, func() {
-		_ = newArchetype(32, comps...)
+		arch := archetype{}
+		arch.init(32, comps...)
 	})
 }
 
 func BenchmarkArchetypeAccess_1000(b *testing.B) {
 	b.StopTimer()
 	comps := []componentType{
-		{ID: 0, Type: reflect.TypeOf(position{})},
-		{ID: 1, Type: reflect.TypeOf(rotation{})},
+		{ID: 0, Type: reflect.TypeOf(testStruct{})},
 	}
 
-	arch := newArchetype(32, comps...)
+	arch := archetype{}
+	arch.init(32, comps...)
 
 	for i := 0; i < 1000; i++ {
 		arch.Add(
 			newEntity(i),
-			component{ID: 0, Component: &position{1, 2}},
-			component{ID: 1, Component: &rotation{3}},
+			component{ID: 0, Component: &testStruct{}},
 		)
 	}
 	b.StartTimer()
 
 	for i := 0; i < b.N; i++ {
-		for j := 0; j < 1000; j++ {
-			pos := (*position)(arch.Get(i, ID(0)))
+		len := int(arch.Len())
+		id := ID(0)
+		for j := 0; j < len; j++ {
+			pos := (*testStruct)(arch.Get(i, id))
 			_ = pos
 		}
 	}
diff --git a/ecs/query_test.go b/ecs/query_test.go
index 8a9e599..ee96e6c 100644
--- a/ecs/query_test.go
+++ b/ecs/query_test.go
@@ -36,7 +36,6 @@ func TestQuery(t *testing.T) {
 		cnt++
 	}
 	assert.Equal(t, 2, cnt)
-	assert.Equal(t, 2, q.Count())
 
 	q = w.Query(posID)
 	cnt = 0
@@ -48,7 +47,6 @@ func TestQuery(t *testing.T) {
 		cnt++
 	}
 	assert.Equal(t, 3, cnt)
-	assert.Equal(t, 3, q.Count())
 
 	q = w.Query(rotID)
 	cnt = 0
@@ -62,10 +60,6 @@ func TestQuery(t *testing.T) {
 		cnt++
 	}
 	assert.Equal(t, 4, cnt)
-	assert.Equal(t, 4, q.Count())
 
 	assert.Panics(t, func() { q.Next() })
-	assert.Panics(t, func() { q.Entity() })
-	assert.Panics(t, func() { q.Has(posID) })
-	assert.Panics(t, func() { q.Get(posID) })
 }
diff --git a/ecs/storage_test.go b/ecs/storage_test.go
index 602ed32..301854f 100644
--- a/ecs/storage_test.go
+++ b/ecs/storage_test.go
@@ -20,7 +20,8 @@ type simpleStruct struct {
 
 func TestReflectStorageAddGet(t *testing.T) {
 	obj1 := testStruct{}
-	s := newStorage(reflect.TypeOf(obj1), 1)
+	s := storage{}
+	s.init(reflect.TypeOf(obj1), 1)
 	storageAddGet(t, s)
 }
 
@@ -52,7 +53,8 @@ func storageAddGet(t *testing.T, s storage) {
 
 func TestReflectStorageRemove(t *testing.T) {
 	ref := simpleStruct{}
-	s := newStorage(reflect.TypeOf(ref), 32)
+	s := storage{}
+	s.init(reflect.TypeOf(ref), 32)
 
 	storageRemove(t, s)
 }
@@ -76,7 +78,8 @@ func storageRemove(t *testing.T, s storage) {
 
 func TestReflectStorageDataSize(t *testing.T) {
 	ref := simpleStruct{}
-	s := newStorage(reflect.TypeOf(ref), 1)
+	s := storage{}
+	s.init(reflect.TypeOf(ref), 1)
 
 	for i := 0; i < 5; i++ {
 		obj := simpleStruct{i}
@@ -100,13 +103,15 @@ func TestReflectStorageDataSize(t *testing.T) {
 }
 
 func TestNewReflectStorage(t *testing.T) {
-	_ = newStorage(reflect.TypeOf(simpleStruct{}), 32)
+	s := storage{}
+	s.init(reflect.TypeOf(simpleStruct{}), 32)
 }
 
 func BenchmarkIterReflectStorage_1000(b *testing.B) {
 	b.StopTimer()
 	ref := testStruct{}
-	s := newStorage(reflect.TypeOf(ref), 128)
+	s := storage{}
+	s.init(reflect.TypeOf(ref), 128)
 	for i := 0; i < 1000; i++ {
 		s.Add(&testStruct{})
 	}
@@ -114,7 +119,8 @@ func BenchmarkIterReflectStorage_1000(b *testing.B) {
 	b.StartTimer()
 
 	for i := 0; i < b.N; i++ {
-		for j := 0; j < int(s.Len()); j++ {
+		len := int(s.Len())
+		for j := 0; j < len; j++ {
 			a := (*testStruct)(s.Get(uint32(j)))
 			_ = a
 		}
@@ -159,7 +165,8 @@ func BenchmarkAddReflectStorage_1000(b *testing.B) {
 	ref := testStruct{}
 	for i := 0; i < b.N; i++ {
 		b.StopTimer()
-		s := newStorage(reflect.TypeOf(ref), 1024)
+		s := storage{}
+		s.init(reflect.TypeOf(ref), 1024)
 		b.StartTimer()
 
 		for i := 0; i < 1000; i++ {
@@ -206,7 +213,8 @@ func BenchmarkRemoveReflectStorage_1000(b *testing.B) {
 	ref := testStruct{}
 	for i := 0; i < b.N; i++ {
 		b.StopTimer()
-		s := newStorage(reflect.TypeOf(ref), 1024)
+		s := storage{}
+		s.init(reflect.TypeOf(ref), 1024)
 		for i := 0; i < 1000; i++ {
 			s.Add(&ref)
 		}