Repository: mlange-42/arche
Pull request #358 (2024-02-02 16:29:16)

Query indexing with `EntityAt`
Adds method `Query.EntityAt()`, useful for things like random sampling of entities.

coveralls: ## Pull Request Test Coverage Report for [Build 7758763547](https://coveralls.io/builds/65469187)


* **0** of **77**   **(100.0%)**  changed or added relevant lines in **1** file are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage remained the same at **100.0%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/65469187/badge)](https://coveralls.io/builds/65469187) |
| :-- | --: |
| Change from base [Build 7754289523](https://coveralls.io/builds/65461123): |  0.0% |
| Covered Lines: | 6224 |
| Relevant Lines: | 6224 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4a4d769..ba7f1ca 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [[unpublished]](https://github.com/mlange-42/arche/compare/v0.10.1...main)
 
+### Features
+
+* Adds method `Query.EntityAt()`, useful for things like random sampling of entities (#358)
+
 ### Bugfixes
 
 * Prevents using the same component multiple times in any operations, through panic (#357)
diff --git a/ecs/query.go b/ecs/query.go
index 40beb70..5a91cfa 100644
--- a/ecs/query.go
+++ b/ecs/query.go
@@ -101,6 +101,30 @@ func (q *Query) Entity() Entity {
 	return q.access.GetEntity(q.entityIndex)
 }
 
+// EntityAt returns the entity at a given index.
+//
+// Do not use this to iterate a query! Use [Query.Next] instead.
+//
+// The method is particularly useful for random sampling of entities
+// from a query (see the example).
+// However, performance heavily depends on the number of archetypes in the world and in the query.
+// It is recommended to use a filter which is registered via the [World.Cache],
+// This increases performance at least by factor 4.
+//
+// Some numbers for comparison, with given number of archetypes in the query:
+//   - 1 archetype, filter not registered: ≈ 11ns
+//   - 1 archetype, filter registered: ≈ 3ns
+//   - 10 archetypes, filter not registered: ≈ 50ns
+//   - 10 archetypes, filter registered: ≈ 10ns
+//
+// The total number of archetypes in the world has no performance impact for registered filters,
+// while for filters that are not registered, it has.
+//
+// Panics if the index is out of range, as indicated by [Query.Count].
+func (q *Query) EntityAt(index int) Entity {
+	return q.entityAt(index)
+}
+
 // Relation returns the target entity for an entity relation.
 //
 // Panics if the entity does not have the given component, or if the component is not a [Relation].
@@ -123,10 +147,11 @@ func (q *Query) relationUnchecked(comp ID) Entity {
 }
 
 // Step advances the query iterator by the given number of entities.
+// Returns whether the query is still at a valid position.
+// Returning false indicates that the step exceeded the Query's entities.
+// The query is closed in this case.
 //
-// Query.Step(1) is equivalent to [Query.Next]().
-//
-// This method, used together with [Query.Count], can be useful for the selection of random entities.
+// Query.Step(1) is equivalent to [Query.Next](), although probably slower.
 func (q *Query) Step(step int) bool {
 	if step <= 0 {
 		panic("step size must be positive")
@@ -329,6 +354,15 @@ func (q *Query) stepArchetype(step uint32) (int, bool) {
 func (q *Query) countEntities() int {
 	var count uint32 = 0
 
+	if q.isFiltered {
+		ln := int32(len(q.archetypes))
+		var i int32
+		for i = 0; i < ln; i++ {
+			count += q.archetypes[i].Len()
+		}
+		return int(count)
+	}
+
 	if q.isBatch {
 		batch := q.nodeArchetypes.(*batchArchetypes)
 		nArch := batch.Len()
@@ -339,13 +373,70 @@ func (q *Query) countEntities() int {
 		return int(count)
 	}
 
+	for _, nd := range q.nodes {
+		if !nd.IsActive || !nd.Matches(q.filter) {
+			continue
+		}
+
+		if !nd.HasRelation {
+			// There should be at least one archetype.
+			// Otherwise, the node would be inactive.
+			arch := nd.Archetypes().Get(0)
+			count += arch.Len()
+			continue
+		}
+
+		if rf, ok := q.filter.(*RelationFilter); ok {
+			target := rf.Target
+			if arch, ok := nd.archetypeMap[target]; ok {
+				count += arch.Len()
+			}
+			continue
+		}
+
+		arches := nd.Archetypes()
+		nArch := arches.Len()
+		var j int32
+		for j = 0; j < nArch; j++ {
+			a := arches.Get(j)
+			count += a.Len()
+		}
+	}
+	return int(count)
+}
+
+func (q *Query) entityAt(index int) Entity {
+	if index < 0 {
+		panic("can't get entity at negative index")
+	}
+	var count uint32 = 0
+	idx := uint32(index)
+
 	if q.isFiltered {
 		ln := int32(len(q.archetypes))
 		var i int32
 		for i = 0; i < ln; i++ {
-			count += q.archetypes[i].Len()
+			ln := q.archetypes[i].Len()
+			if idx < count+ln {
+				return q.archetypes[i].GetEntity(idx - count)
+			}
+			count += ln
 		}
-		return int(count)
+		panic(fmt.Sprintf("query index out of range: index %d, length %d", index, count))
+	}
+
+	if q.isBatch {
+		batch := q.nodeArchetypes.(*batchArchetypes)
+		nArch := batch.Len()
+		var j int32
+		for j = 0; j < nArch; j++ {
+			ln := batch.EndIndex[j] - batch.StartIndex[j]
+			if idx < count+ln {
+				return batch.Archetype[j].GetEntity(batch.StartIndex[j] + idx - count)
+			}
+			count += ln
+		}
+		panic(fmt.Sprintf("query index out of range: index %d, length %d", index, count))
 	}
 
 	for _, nd := range q.nodes {
@@ -356,16 +447,25 @@ func (q *Query) countEntities() int {
 		if !nd.HasRelation {
 			// There should be at least one archetype.
 			// Otherwise, the node would be inactive.
-			arches := nd.Archetypes()
-			arch := arches.Get(0)
-			count += arch.Len()
+			arch := nd.Archetypes().Get(0)
+
+			ln := arch.Len()
+			if idx < count+ln {
+				return arch.GetEntity(idx - count)
+			}
+			count += ln
 			continue
 		}
 
 		if rf, ok := q.filter.(*RelationFilter); ok {
 			target := rf.Target
 			if arch, ok := nd.archetypeMap[target]; ok {
-				count += arch.Len()
+
+				ln := arch.Len()
+				if idx < count+ln {
+					return arch.GetEntity(idx - count)
+				}
+				count += ln
 			}
 			continue
 		}
@@ -374,9 +474,14 @@ func (q *Query) countEntities() int {
 		nArch := arches.Len()
 		var j int32
 		for j = 0; j < nArch; j++ {
-			a := arches.Get(j)
-			count += a.Len()
+			arch := arches.Get(j)
+
+			ln := arch.Len()
+			if idx < count+ln {
+				return arch.GetEntity(idx - count)
+			}
+			count += ln
 		}
 	}
-	return int(count)
+	panic(fmt.Sprintf("query index out of range: index %d, length %d", index, count))
 }
diff --git a/examples/random_sampling/main.go b/examples/random_sampling/main.go
index 2ffae6c..2b5ea62 100644
--- a/examples/random_sampling/main.go
+++ b/examples/random_sampling/main.go
@@ -1,10 +1,9 @@
-// Demonstrates random sampling of a fixed number of entities from a query using Query.Step().
+// Demonstrates random sampling of a fixed number of entities from a query using Query.EntityAt().
 package main
 
 import (
 	"fmt"
 	"math/rand"
-	"sort"
 
 	"github.com/mlange-42/arche/ecs"
 	"github.com/mlange-42/arche/generic"
@@ -21,31 +20,28 @@ func main() {
 
 	// Create a generic filter.
 	filter := generic.NewFilter0()
+	// Register the filter. Optional, but recommended for query random access.
+	filter.Register(&world)
 
 	// Get a fresh query iterator.
 	query := filter.Query(&world)
 	// Get the number of entities in the query
 	count := query.Count()
-	// Get sorted random indices
-	sample := sortedSample(25, count)
+	// Get random indices without replacement
+	sample := sample(25, count)
 	fmt.Println(sample)
 
-	// Iterate, only visiting the entities at the given indices.
-	last := -1
+	// Iterate over sampled indices.
 	for _, idx := range sample {
-		// Calculate the step size (can be 0)
-		step := idx - last
-		// Advance the query iterator
-		query.Step(step)
+		// Get the entity at the random index.
+		entity := query.EntityAt(idx)
 		// Do something with the entity and/or components at the current iterator position.
-		fmt.Println(query.Entity())
-		// Required for calculating the step size.
-		last = idx
+		fmt.Println(entity)
 	}
 }
 
-// Returns a sorted random sample of indices of size k, for a "thing" of length n.
-func sortedSample(k, n int) []int {
+// Returns a random sample of indices without replacement.Take k of n sample.
+func sample(k, n int) []int {
 	m := make([]int, n)
 	for i := 0; i < n; i++ {
 		j := rand.Intn(i + 1)
@@ -53,6 +49,5 @@ func sortedSample(k, n int) []int {
 		m[j] = i
 	}
 	m = m[:k]
-	sort.Ints(m)
 	return m
 }


Test patch:
diff --git a/ecs/query_examples_test.go b/ecs/query_examples_test.go
index 14f6bad..2e0e11e 100644
--- a/ecs/query_examples_test.go
+++ b/ecs/query_examples_test.go
@@ -2,6 +2,7 @@ package ecs_test
 
 import (
 	"fmt"
+	"math/rand"
 
 	"github.com/mlange-42/arche/ecs"
 )
@@ -47,3 +48,38 @@ func ExampleQuery_Close() {
 	query.Close()
 	// Output: 1
 }
+
+func ExampleQuery_EntityAt() {
+	// Set up the world.
+	world := ecs.NewWorld()
+	posID := ecs.ComponentID[Position](&world)
+
+	// Create entities.
+	builder := ecs.NewBuilder(&world, posID)
+	builder.NewBatch(100)
+
+	// Create a random generator.
+	rng := rand.New(rand.NewSource(42))
+
+	// Register a filter (optional, but recommended).
+	filter := world.Cache().Register(ecs.All(posID))
+
+	// Query some entities.
+	query := world.Query(&filter)
+	// Get the number of entities in the query
+	cnt := query.Count()
+
+	// Sample some random entities.
+	for i := 0; i < 5; i++ {
+		randomEntity := query.EntityAt(rng.Intn(cnt))
+		fmt.Println(randomEntity)
+	}
+
+	// Close the query.
+	query.Close()
+	// Output: {6 0}
+	// {88 0}
+	// {69 0}
+	// {51 0}
+	// {24 0}
+}
diff --git a/ecs/query_test.go b/ecs/query_test.go
index b26ba88..1b9c32f 100644
--- a/ecs/query_test.go
+++ b/ecs/query_test.go
@@ -1,6 +1,7 @@
 package ecs
 
 import (
+	"math/rand"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -316,6 +317,7 @@ func TestQueryStep(t *testing.T) {
 	q := w.Query(All(posID))
 	cnt := 0
 	for q.Next() {
+		assert.Equal(t, cnt+1, int(q.Entity().id))
 		cnt++
 	}
 	assert.Equal(t, 10, cnt)
@@ -459,3 +461,188 @@ func TestQueryIds(t *testing.T) {
 	query.Next()
 	assert.Equal(t, query.Ids(), []ID{id(0), id(1)})
 }
+
+func TestQueryEntityAt(t *testing.T) {
+	world := NewWorld()
+
+	rotID := ComponentID[rotation](&world)
+	posID := ComponentID[Position](&world)
+	velID := ComponentID[Velocity](&world)
+	relID := ComponentID[relationComp](&world)
+
+	bPos := NewBuilder(&world, rotID, posID)
+	bPosVelRel := NewBuilder(&world, rotID, posID, velID, relID).WithRelation(relID)
+	bPosRel := NewBuilder(&world, rotID, posID, relID).WithRelation(relID)
+	bVelRel := NewBuilder(&world, rotID, velID, relID).WithRelation(relID)
+	bRel := NewBuilder(&world, rotID, relID).WithRelation(relID)
+
+	parent1 := world.NewEntity()
+	parent2 := world.NewEntity()
+
+	bPos.NewBatch(10)
+	bPosVelRel.NewBatch(10, parent2)
+	bPosRel.NewBatch(10, parent1)
+	bVelRel.NewBatch(10, parent1)
+	bRel.NewBatch(10)
+
+	query := world.Query(All())
+
+	assert.Panics(t, func() { query.EntityAt(-1) })
+	for i := 0; i < 52; i++ {
+		assert.Equal(t, Entity{eid(i + 1), 0}, query.EntityAt(i))
+	}
+	assert.Panics(t, func() { query.EntityAt(52) })
+	query.Close()
+
+	relFilter := NewRelationFilter(All(relID), parent2)
+	query = world.Query(&relFilter)
+	for i := 0; i < 10; i++ {
+		assert.Equal(t, Entity{eid(i + 13), 0}, query.EntityAt(i))
+	}
+	assert.Panics(t, func() { query.EntityAt(10) })
+	query.Close()
+
+	query = world.Query(All(relID))
+	for i := 0; i < 40; i++ {
+		assert.Equal(t, Entity{eid(i + 13), 0}, query.EntityAt(i))
+	}
+	assert.Panics(t, func() { query.EntityAt(40) })
+	query.Close()
+
+	query = world.Batch().ExchangeQ(All(posID), nil, []ID{posID})
+	assert.Equal(t, 30, query.Count())
+
+	for i := 0; i < 30; i++ {
+		assert.Equal(t, Entity{eid(i + 3), 0}, query.EntityAt(i))
+	}
+	assert.Panics(t, func() { query.EntityAt(30) })
+	query.Close()
+
+	f := All(relID)
+	filter := world.Cache().Register(f)
+	query = world.Query(&filter)
+	assert.Equal(t, 40, query.Count())
+	for i := 0; i < 40; i++ {
+		assert.Equal(t, (i+3)%10, int(query.EntityAt(i).id)%10)
+	}
+	assert.Panics(t, func() { query.EntityAt(40) })
+	query.Close()
+}
+
+func BenchmarkQueryEntityAt_1Arch_1000(b *testing.B) {
+	b.StopTimer()
+	w := NewWorld()
+	posID := ComponentID[Position](&w)
+	builder := NewBuilder(&w, posID)
+	builder.NewBatch(1000)
+
+	indices := make([]int, 1000)
+	for i := range indices {
+		indices[i] = rand.Intn(1000)
+	}
+
+	query := w.Query(All(posID))
+	b.StartTimer()
+	var e Entity
+	for i := 0; i < b.N; i++ {
+		for _, idx := range indices {
+			e = query.EntityAt(idx)
+		}
+	}
+	_ = e
+}
+
+func BenchmarkQueryEntityAt_1Arch_1000_Registered(b *testing.B) {
+	b.StopTimer()
+	w := NewWorld()
+	posID := ComponentID[Position](&w)
+	builder := NewBuilder(&w, posID)
+	builder.NewBatch(1000)
+
+	indices := make([]int, 1000)
+	for i := range indices {
+		indices[i] = rand.Intn(1000)
+	}
+
+	f := All(posID)
+	filter := w.Cache().Register(f)
+	query := w.Query(&filter)
+	b.StartTimer()
+	var e Entity
+	for i := 0; i < b.N; i++ {
+		for _, idx := range indices {
+			e = query.EntityAt(idx)
+		}
+	}
+	_ = e
+}
+
+func BenchmarkQueryEntityAt_10Arch_1000(b *testing.B) {
+	b.StopTimer()
+	w := NewWorld()
+	id1 := ComponentID[testStruct0](&w)
+	id2 := ComponentID[testStruct1](&w)
+	id3 := ComponentID[testStruct2](&w)
+	id4 := ComponentID[testStruct3](&w)
+
+	comps := [][]ID{
+		{id1}, {id2}, {id3}, {id4},
+		{id1, id2}, {id1, id3}, {id1, id4},
+		{id2, id3}, {id2, id4}, {id3, id4},
+	}
+	for _, c := range comps {
+		b := NewBuilder(&w, c...)
+		b.NewBatch(100)
+	}
+
+	indices := make([]int, 1000)
+	for i := range indices {
+		indices[i] = rand.Intn(1000)
+	}
+
+	query := w.Query(All())
+	b.StartTimer()
+	var e Entity
+	for i := 0; i < b.N; i++ {
+		for _, idx := range indices {
+			e = query.EntityAt(idx)
+		}
+	}
+	_ = e
+}
+
+func BenchmarkQueryEntityAt_10Arch_1000_Registered(b *testing.B) {
+	b.StopTimer()
+	w := NewWorld()
+	id1 := ComponentID[testStruct0](&w)
+	id2 := ComponentID[testStruct1](&w)
+	id3 := ComponentID[testStruct2](&w)
+	id4 := ComponentID[testStruct3](&w)
+
+	comps := [][]ID{
+		{id1}, {id2}, {id3}, {id4},
+		{id1, id2}, {id1, id3}, {id1, id4},
+		{id2, id3}, {id2, id4}, {id3, id4},
+	}
+	for _, c := range comps {
+		b := NewBuilder(&w, c...)
+		b.NewBatch(100)
+	}
+
+	indices := make([]int, 1000)
+	for i := range indices {
+		indices[i] = rand.Intn(1000)
+	}
+
+	f := All()
+	filter := w.Cache().Register(f)
+	query := w.Query(&filter)
+	b.StartTimer()
+	var e Entity
+	for i := 0; i < b.N; i++ {
+		for _, idx := range indices {
+			e = query.EntityAt(idx)
+		}
+	}
+	_ = e
+}