Implement entity enumeration

Author: GeorchW

src/EntityRegistry.cs

@@ -1,9 +1,10 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 
 namespace SimpleECS
 {
-    class EntityRegistry
+    class EntityRegistry : IReadOnlyCollection<Entity>
     {
         public struct EntityLocation
         {
@@ -20,19 +21,29 @@ public EntityLocation(ArchetypeContainer archetypeContainer, int index)
         private const int InitialCapacity = 512;
         EntityLocation[] Locations = new EntityLocation[InitialCapacity];
         int[] Versions = new int[InitialCapacity];
-        Queue<int> freeIDs = new Queue<int>();
-        int count = 0;
+        HashSet<int> freeIDs = new HashSet<int>();
+        int nextId = 0;
+
+        public int Count => nextId - freeIDs.Count;
 
         public int GetVersion(int entity) => Versions[entity];
 
         public Entity RegisterEntity(ArchetypeContainer container, int index)
         {
-            if (count == Locations.Length)
+            if (nextId == Locations.Length)
                 Array.Resize(ref Locations, Locations.Length * 2);
-            if (!freeIDs.TryDequeue(out int id))
+            int id;
+            if (freeIDs.Count > 0)
             {
-                id = count;
-                count++;
+                // faster implementation of First()
+                foreach (var _id in freeIDs) { id = _id; goto done; }
+                throw new Exception();
+            done:;
+            }
+            else
+            {
+                id = nextId;
+                nextId++;
             }
             Locations[id] = new EntityLocation(container, index);
             return new Entity(id, Versions[id]);
@@ -44,7 +55,7 @@ public void UnregisterEntity(Entity entity, out EntityLocation lastLocation)
             lastLocation = Locations[entity.Id];
             Locations[entity.Id] = default;
             Versions[entity.Id]++;
-            freeIDs.Enqueue(entity.Id);
+            freeIDs.Add(entity.Id);
         }
         public void MoveEntity(int entity, ArchetypeContainer newContainer, int newIndex)
         {
@@ -63,5 +74,38 @@ public bool TryGetLocation(Entity entity, out EntityLocation location)
                 return true;
             }
         }
+
+        public IEnumerator<Entity> GetEnumerator() => new EntityEnumerator(this);
+        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+        struct EntityEnumerator : IEnumerator<Entity>
+        {
+            int currentId;
+            EntityRegistry registry;
+            public Entity Current => new Entity(currentId, registry.GetVersion(currentId));
+
+            object? IEnumerator.Current => this.Current;
+
+            public EntityEnumerator(EntityRegistry registry)
+            {
+                this.currentId = -1;
+                this.registry = registry;
+            }
+
+            public void Dispose() { }
+
+            public bool MoveNext()
+            {
+                do
+                {
+                    currentId++;
+                    if (registry.nextId == currentId)
+                        return false;
+                }
+                while (registry.freeIDs.Contains(currentId));
+                return true;
+            }
+
+            public void Reset() => currentId = -1;
+        }
     }
 }

src/Scene.cs

@@ -1,10 +1,11 @@
 using System.Linq;
 using System;
 using System.Collections.Generic;
+using System.Collections;
 
 namespace SimpleECS
 {
-    public class Scene
+    public class Scene : IReadOnlyCollection<Entity>
     {
         internal Dictionary<ComponentSet.Readonly, ArchetypeContainer> archetypes = new Dictionary<ComponentSet.Readonly, ArchetypeContainer>();
 
@@ -14,6 +15,10 @@ public class Scene
         public GlobalStorage Globals { get; } = new GlobalStorage();
         public CallbackManager Callbacks { get; } = new CallbackManager();
 
+        public int Count => EntityRegistry.Count;
+        public IEnumerator<Entity> GetEnumerator() => EntityRegistry.GetEnumerator();
+        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
         public Scene()
         {
             InitialContainer = new ArchetypeContainer(this);
@@ -41,7 +46,7 @@ public void DeleteEntity(Entity entity)
             }
             foreach (var (index, type) in location.ArchetypeContainer.additions.Keys)
             {
-                if(index == location.Index)
+                if (index == location.Index)
                 {
                     Callbacks.TryGet(type)?.OnComponentRemoved(this, entity);
                 }

test/Archetypes.cs

@@ -46,7 +46,7 @@ public void CorrectArchetypesCreated()
             Assert.That(scene.archetypes.Count == 1);
             scene.InsertNewComponents();
             Assert.That(scene.archetypes.Count == 2);
-            
+
             var newEntity = scene.CreateEntity();
             scene.InsertNewComponents();
             Assert.That(scene.archetypes.Count == 2);
@@ -55,5 +55,11 @@ public void CorrectArchetypesCreated()
             newEntity.Add<ExampleComp2>();
             Assert.That(scene.archetypes.Count == 2);
         }
+
+        [Test]
+        public void EnumerateEntities()
+        {
+            Assert.That(scene, Is.EquivalentTo(new[] { entity }));
+        }
     }
 }