[#144] Fix ordered query (#145)

examples/bunnymark_ecs/component/position.go

@@ -6,10 +6,11 @@ import (
 
 type PositionData struct {
 	X, Y float64
+	ID   int
 }
 
 func (p PositionData) Order() int {
-	return int(p.Y * 600)
+	return -p.ID
 }
 
 var Position = donburi.NewComponentType[PositionData]()

examples/bunnymark_ecs/system/spawn.go

@@ -16,6 +16,7 @@ var UsePositionOrdering bool
 
 type Spawn struct {
 	settings *component.SettingsData
+	nextID   int
 }
 
 func NewSpawn() *Spawn {
@@ -73,8 +74,10 @@ func (s *Spawn) addBunnies(ecs *ecs.ECS) {
 		entry := ecs.World.Entry(entity)
 		position := component.Position.Get(entry)
 		*position = component.PositionData{
-			X: float64(i % 2), // Alternate screen edges
+			ID: s.nextID,
+			X:  float64(i % 2), // Alternate screen edges
 		}
+		s.nextID++
 		donburi.SetValue(
 			entry, component.Velocity, component.VelocityData{
 				X: helper.RangeFloat(0, 0.005),

query.go

@@ -2,6 +2,8 @@ package donburi
 
 import (
 	"iter"
+	"sort"
+	"sync"
 
 	"github.com/yohamta/donburi/filter"
 	"github.com/yohamta/donburi/internal/storage"
@@ -40,6 +42,8 @@ func NewQuery(filter filter.LayoutFilter) *Query {
 // when running ordered queries using `EachOrdered`.
 type OrderedQuery[T IOrderable] struct {
 	Query
+	entries []*Entry
+	lock    sync.Mutex
 }
 
 // NewOrderedQuery creates a new ordered query.
@@ -59,27 +63,41 @@ func NewOrderedQuery[T IOrderable](filter filter.LayoutFilter) *OrderedQuery[T]
 // ordered by the specified component.
 func (q *OrderedQuery[T]) IterOrdered(w World, orderBy *ComponentType[T]) iter.Seq[*Entry] {
 	return func(yield func(*Entry) bool) {
+		q.lock.Lock()
+		defer q.lock.Unlock()
+
 		accessor := w.StorageAccessor()
 		iter := storage.NewEntityIterator(0, accessor.Archetypes, q.evaluateQuery(w, &accessor))
 
+		// Clear the slice while keeping the underlying array
+		q.entries = q.entries[:0]
+
 		for iter.HasNext() {
 			archetype := iter.Next()
 			archetype.Lock()
 
 			ents := archetype.Entities()
-			entrIter := NewOrderedEntryIterator(0, w, ents, orderBy)
-			for entrIter.HasNext() {
-				e := entrIter.Next()
-				if e.entity.IsReady() {
-					if !yield(e) {
-						archetype.Unlock()
-						return
-					}
+			for _, entity := range ents {
+				entry := w.Entry(entity)
+				if entry.entity.IsReady() {
+					q.entries = append(q.entries, entry)
 				}
 			}
 
 			archetype.Unlock()
 		}
+
+		// Sort all entries
+		sort.Slice(q.entries, func(i, j int) bool {
+			return orderBy.GetValue(q.entries[i]).Order() < orderBy.GetValue(q.entries[j]).Order()
+		})
+
+		// Yield sorted entries
+		for _, entry := range q.entries {
+			if !yield(entry) {
+				return
+			}
+		}
 	}
 }
 

query_test.go

@@ -44,6 +44,40 @@ func TestQuery(t *testing.T) {
 	}
 }
 
+var _ donburi.IOrderable = orderableData{}
+
+type orderableData struct {
+	Index int
+}
+
+func (o orderableData) Order() int {
+	return o.Index
+}
+
+var orderable = donburi.NewComponentType[orderableData]()
+
+func TestOrderedQuery(t *testing.T) {
+	orderedEntitiesQuery := donburi.NewOrderedQuery[orderableData](
+		filter.Contains(orderable),
+	)
+
+	world := donburi.NewWorld()
+	for _, i := range []int{3, 1, 2} {
+		e := world.Create(orderable)
+		entr := world.Entry(e)
+		donburi.SetValue(entr, orderable, orderableData{i})
+	}
+
+	var i int
+	for e := range orderedEntitiesQuery.IterOrdered(world, orderable) {
+		o := orderable.GetValue(e)
+		i += 1
+		if o.Index != i {
+			t.Errorf("expected %d, but got %d", i, o.Index)
+		}
+	}
+}
+
 func BenchmarkQuery_EachOrdered(b *testing.B) {
 	world := donburi.NewWorld()
 	for i := 0; i < 30000; i++ {