Repository: mlange-42/arche
Pull request #119 (2023-02-24 16:16:45)

Query count and step
`Query` has methods `Count()` and `Step(int)`, primarily for effective random sampling.

Alternative for #118.
Allows only forward jumps and counts, but looks generally better.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f49ef80..c7a3ec8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,14 @@
 # Changelog
 
-## [[unpublished]](https://github.com/mlange-42/arche/compare/v0.4.3...main)
+## [[unpublished]](https://github.com/mlange-42/arche/compare/v0.4.4...main)
+
+Nothing
+
+## [[v0.4.4]](https://github.com/mlange-42/arche/compare/v0.4.3...v0.4.4)
+
+### Features
+
+* `Query` has methods `Count()` and `Step(int)`, primarily for effective random sampling (#119)
 
 ### Bugfixes
 
@@ -9,6 +17,7 @@
 ### Other
 
 * Update to [go-gameengine-ecs](https://github.com/marioolofo/go-gameengine-ecs) v0.9.0 in benchmarks (#116)
+* Adds example `random_sampling` to demonstrate usage of `Query.Count()` and `Query.Step(int)` (#119)
 * Remove internal wrapper structs in generic queries and maps (#120)
 
 ## [[v0.4.3]](https://github.com/mlange-42/arche/compare/v0.4.2...v0.4.3)
diff --git a/ecs/query.go b/ecs/query.go
index a956f11..fd93889 100644
--- a/ecs/query.go
+++ b/ecs/query.go
@@ -47,6 +47,7 @@ func newQuery(world *World, filter Filter, lockBit uint8) Query {
 			world:   world,
 			index:   -1,
 			lockBit: lockBit,
+			count:   -1,
 		},
 		filter: filter,
 	}
@@ -61,6 +62,42 @@ func (q *Query) Next() bool {
 	return q.nextArchetype()
 }
 
+// Step advances the query iterator by the given number of entities.
+//
+// Query.Step(1) is equivalent to [Query.Next]().
+//
+// This method, used together with [Query.Count], can be useful for the selection of random entities.
+func (q *Query) Step(step int) bool {
+	if step <= 0 {
+		panic("step size must be positive")
+	}
+	var ok bool
+	for {
+		step, ok = q.archetype.Step(uint32(step))
+		if ok {
+			return true
+		}
+		if !q.nextArchetype() {
+			return false
+		}
+		if step == 0 {
+			return true
+		}
+	}
+}
+
+// Count counts the entities matching this query.
+//
+// Involves a small overhead of iterating through archetypes when called the first time.
+// However, it is considerable faster than manual counting via iteration.
+func (q *Query) Count() int {
+	if q.count >= 0 {
+		return q.count
+	}
+	q.count = q.world.count(q.filter)
+	return q.count
+}
+
 func (q *Query) nextArchetype() bool {
 	index, archetype, ok := q.world.nextArchetype(q.filter, q.index)
 	q.index = index
@@ -77,6 +114,7 @@ type queryIter struct {
 	archetype archetypeIter
 	index     int
 	lockBit   uint8
+	count     int
 }
 
 // Has returns whether the current [Entity] has the given component.
@@ -127,6 +165,17 @@ func (it *archetypeIter) Next() bool {
 	return it.Index < it.Length
 }
 
+func (it *archetypeIter) Step(count uint32) (int, bool) {
+	if it.Length == 0 {
+		return int(count - 1), false
+	}
+	it.Index += count
+	if it.Index < it.Length {
+		return 0, true
+	}
+	return int(it.Index) - int(it.Length), false
+}
+
 // Has returns whether the current entity has the given component
 func (it *archetypeIter) Has(comp ID) bool {
 	return it.Archetype.HasComponent(comp)
diff --git a/ecs/world.go b/ecs/world.go
index 149828c..4ed7376 100644
--- a/ecs/world.go
+++ b/ecs/world.go
@@ -524,6 +524,19 @@ func (w *World) nextArchetype(filter Filter, index int) (int, archetypeIter, boo
 	return len, archetypeIter{}, false
 }
 
+// Counts the entities matching the filter
+func (w *World) count(filter Filter) int {
+	len := int(w.archetypes.Len())
+	count := uint32(0)
+	for i := 0; i < len; i++ {
+		a := w.archetypes.Get(i)
+		if filter.Matches(a.Mask) {
+			count += a.Len()
+		}
+	}
+	return int(count)
+}
+
 // closeQuery closes a query and unlocks the world
 func (w *World) closeQuery(query *queryIter) {
 	l := query.lockBit
diff --git a/examples/random_sampling/main.go b/examples/random_sampling/main.go
new file mode 100644
index 0000000..2ffae6c
--- /dev/null
+++ b/examples/random_sampling/main.go
@@ -0,0 +1,58 @@
+// Demonstrates random sampling of a fixed number of entities from a query using Query.Step().
+package main
+
+import (
+	"fmt"
+	"math/rand"
+	"sort"
+
+	"github.com/mlange-42/arche/ecs"
+	"github.com/mlange-42/arche/generic"
+)
+
+func main() {
+	// Create a World.
+	world := ecs.NewWorld()
+
+	// Create entities.
+	for i := 0; i < 1000; i++ {
+		world.NewEntity()
+	}
+
+	// Create a generic filter.
+	filter := generic.NewFilter0()
+
+	// 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)
+	fmt.Println(sample)
+
+	// Iterate, only visiting the entities at the given indices.
+	last := -1
+	for _, idx := range sample {
+		// Calculate the step size (can be 0)
+		step := idx - last
+		// Advance the query iterator
+		query.Step(step)
+		// 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
+	}
+}
+
+// Returns a sorted random sample of indices of size k, for a "thing" of length n.
+func sortedSample(k, n int) []int {
+	m := make([]int, n)
+	for i := 0; i < n; i++ {
+		j := rand.Intn(i + 1)
+		m[i] = m[j]
+		m[j] = i
+	}
+	m = m[:k]
+	sort.Ints(m)
+	return m
+}


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index eed3660..96ea1ef 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -80,5 +80,6 @@ jobs:
           go run ./examples/generic
           go run ./examples/locked_world
           go run ./examples/random_access
+          go run ./examples/random_sampling
           go run ./examples/readme
           go run ./examples/world_stats
diff --git a/ecs/query_test.go b/ecs/query_test.go
index 5d3fe16..1f45a62 100644
--- a/ecs/query_test.go
+++ b/ecs/query_test.go
@@ -107,6 +107,106 @@ func TestQuery(t *testing.T) {
 	assert.Equal(t, 0, cnt)
 }
 
+func TestQueryCount(t *testing.T) {
+	w := NewWorld()
+
+	posID := ComponentID[position](&w)
+	rotID := ComponentID[rotation](&w)
+
+	e0 := w.NewEntity()
+	e1 := w.NewEntity()
+	e2 := w.NewEntity()
+	e3 := w.NewEntity()
+	e4 := w.NewEntity()
+
+	w.Add(e0, posID)
+	w.Add(e1, posID, rotID)
+	w.Add(e2, posID, rotID)
+	w.Add(e3, posID, rotID)
+	w.Add(e4, rotID)
+
+	q := w.Query(All(posID))
+	assert.Equal(t, 4, q.Count())
+	assert.Equal(t, 4, q.Count())
+}
+
+func TestQueryStep(t *testing.T) {
+	w := NewWorld()
+
+	posID := ComponentID[position](&w)
+	velID := ComponentID[velocity](&w)
+	rotID := ComponentID[rotation](&w)
+
+	_ = w.NewEntity(posID)
+	_ = w.NewEntity(posID, rotID)
+	_ = w.NewEntity(posID, rotID)
+	_ = w.NewEntity(posID, rotID)
+	_ = w.NewEntity(posID, rotID)
+	_ = w.NewEntity(posID, velID)
+	_ = w.NewEntity(posID, velID)
+	_ = w.NewEntity(posID, velID)
+	_ = w.NewEntity(posID, velID, rotID)
+	_ = w.NewEntity(posID, velID, rotID)
+
+	q := w.Query(All(posID))
+	cnt := 0
+	for q.Next() {
+		cnt++
+	}
+	assert.Equal(t, 10, cnt)
+
+	q = w.Query(All(posID))
+	assert.Equal(t, 10, q.Count())
+
+	cnt = 0
+	for q.Step(1) {
+		cnt++
+	}
+	assert.Equal(t, 10, cnt)
+
+	q = w.Query(All(posID))
+	q.Next()
+	assert.Equal(t, Entity{1, 0}, q.Entity())
+	q.Step(1)
+	assert.Equal(t, Entity{2, 0}, q.Entity())
+	q.Step(2)
+	assert.Equal(t, Entity{4, 0}, q.Entity())
+	q.Step(3)
+	assert.Equal(t, Entity{7, 0}, q.Entity())
+	q.Step(3)
+	assert.Equal(t, Entity{10, 0}, q.Entity())
+
+	assert.True(t, w.IsLocked())
+
+	assert.False(t, q.Step(3))
+	assert.False(t, w.IsLocked())
+
+	q = w.Query(All(posID))
+	q.Step(1)
+	assert.Equal(t, Entity{1, 0}, q.Entity())
+
+	q = w.Query(All(posID))
+	q.Step(2)
+	assert.Equal(t, Entity{2, 0}, q.Entity())
+
+	q = w.Query(All(posID))
+	q.Step(10)
+	assert.Equal(t, Entity{10, 0}, q.Entity())
+
+	q = w.Query(All(posID))
+	assert.Panics(t, func() { q.Step(0) })
+	q.Step(2)
+	assert.Panics(t, func() { q.Step(0) })
+
+	q = w.Query(All(posID))
+	cnt = 0
+	for q.Step(2) {
+		cnt++
+	}
+	assert.Equal(t, 5, cnt)
+
+}
+
 func TestQueryClosed(t *testing.T) {
 	w := NewWorld()
 
@@ -122,7 +222,11 @@ func TestQueryClosed(t *testing.T) {
 	w.Add(e2, posID, rotID)
 
 	q := w.Query(All(posID, rotID))
+	assert.Panics(t, func() { q.Entity() })
+	assert.Panics(t, func() { q.Get(posID) })
+
 	q.Close()
-	assert.Panics(t, func() { q.Next() })
+	assert.Panics(t, func() { q.Entity() })
 	assert.Panics(t, func() { q.Get(posID) })
+	assert.Panics(t, func() { q.Next() })
 }