From 4f5fd506602945a9470fc3e1bc0589c9f9a01c54 Mon Sep 17 00:00:00 2001
From: bloeys <bloeys@outlook.com>
Date: Fri, 6 Oct 2023 08:42:02 +0400
Subject: [PATCH] Fix iterator bug

---
 registry/iterator.go | 22 ++++++++++++++++++++--
 registry/registry.go |  4 ++--
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/registry/iterator.go b/registry/iterator.go
index c9ec2fd..c7944c5 100755
--- a/registry/iterator.go
+++ b/registry/iterator.go
@@ -11,9 +11,15 @@ package registry
 //   - If items were both added and removed, the iterator might follow either of the previous 2 cases or a combination of them
 //
 // To summarize: The iterator will *never* return more items than were alive at the time of its creation, and will *never* return freed items
+//
+// Example usage:
+//
+//	for item, handle := it.Next(); !it.IsDone(); item, handle = it.Next() {
+//		// Do stuff
+//	}
 type Iterator[T any] struct {
 	registry       *Registry[T]
-	remainingItems uint64
+	remainingItems uint
 	currIndex      int
 }
 
@@ -23,6 +29,17 @@ func (it *Iterator[T]) Next() (*T, Handle) {
 		return nil, 0
 	}
 
+	// This does two things:
+	//
+	// First is if IsDone() only checked 'remainingItems', then when Next() returns the last item IsDone() will immediately be true which will cause loops to exit before processing that last item!
+	// However, with this check IsDone will remain false until Next() is called at least one more time after returning the last item which ensures the last item is processed in the loop.
+	//
+	// Secondly, if the iterator is created on an empty registry, the IsDone() check above won't pass, however the check here will correctly handle the case and make IsDone start returning true
+	if it.remainingItems == 0 {
+		it.currIndex = -1
+		return nil, 0
+	}
+
 	for ; it.currIndex < len(it.registry.Handles); it.currIndex++ {
 
 		handle := it.registry.Handles[it.currIndex]
@@ -39,10 +56,11 @@ func (it *Iterator[T]) Next() (*T, Handle) {
 	// means that the registry changed since we were created, and that remainingItems is not accurate.
 	//
 	// As such, we zero remaining items so that this iterator is considered done
+	it.currIndex = -1
 	it.remainingItems = 0
 	return nil, 0
 }
 
 func (it *Iterator[T]) IsDone() bool {
-	return it.remainingItems == 0
+	return it.currIndex == -1 && it.remainingItems == 0
 }
diff --git a/registry/registry.go b/registry/registry.go
index efd3ea5..07aeae1 100755
--- a/registry/registry.go
+++ b/registry/registry.go
@@ -18,7 +18,7 @@ type freeListitem struct {
 //
 // It is NOT safe to concurrently create or free items. However, it is SAFE to concurrently get items
 type Registry[T any] struct {
-	ItemCount uint64
+	ItemCount uint
 	Handles   []Handle
 	Items     []T
 
@@ -32,7 +32,7 @@ type Registry[T any] struct {
 
 func (r *Registry[T]) New() (*T, Handle) {
 
-	assert.T(r.ItemCount < uint64(len(r.Handles)), "Can not add more entities to registry because it is full")
+	assert.T(r.ItemCount < uint(len(r.Handles)), "Can not add more entities to registry because it is full")
 
 	var index uint64 = math.MaxUint64