From d5414f1bc453394e0bf045a12e262c06984debdd Mon Sep 17 00:00:00 2001
From: Marcus <ukendio@gmail.com>
Date: Sun, 5 May 2024 15:22:02 +0200
Subject: [PATCH] Add iter method (#20)

---
 lib/init.lua      | 33 ++++++++++++++++++++++++++++++++-
 lib/init.spec.lua | 33 +++++++++++++++++++++++++++++++++
 tests/test1.lua   | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/lib/init.lua b/lib/init.lua
index 6d9c1fe..8c0b5bd 100644
--- a/lib/init.lua
+++ b/lib/init.lua
@@ -178,7 +178,7 @@ local World = {}
 World.__index = World
 function World.new()
 	local self = setmetatable({
-		archetypeIndex = {};
+    archetypeIndex = {};
 		archetypes = {};
 		componentIndex = {};
 		entityIndex = {};
@@ -651,6 +651,37 @@ function World.observer(world: World, ...)
 	}
 end
 
+function World.__iter(world: World): () -> (number?, unknown?)
+	local entityIndex = world.entityIndex
+	local last
+
+	return function() 
+		local entity, record = next(entityIndex, last)
+		if not entity then 
+			return
+		end
+		last = entity
+
+		local archetype = record.archetype
+		if not archetype then 
+			-- Returns only the entity id as an entity without data should not return
+			-- data and allow the user to get an error if they don't handle the case.
+			return entity 
+		end
+
+		local row = record.row
+		local types = archetype.types
+		local columns = archetype.columns
+		local entityData = {}
+		for i, column in columns do
+			-- We use types because the key should be the component ID not the column index
+			entityData[types[i]] = column[row]
+		end
+		
+		return entity, entityData
+	end
+end
+
 return table.freeze({
 	World = World;
 	ON_ADD = ON_ADD;
diff --git a/lib/init.spec.lua b/lib/init.spec.lua
index fdc8331..98f485b 100644
--- a/lib/init.spec.lua
+++ b/lib/init.spec.lua
@@ -299,5 +299,38 @@ return function()
 			expect(world:get(id, Poison)).to.never.be.ok()
 			expect(world:get(id, Health)).to.never.be.ok()
 		end)
+
+		it("should allow iterating the whole world", function() 
+			local world = jecs.World.new()
+
+			local A, B = world:entity(), world:entity()
+
+			local eA = world:entity()
+			world:set(eA, A, true)
+			local eB = world:entity()
+			world:set(eB, B, true)
+			local eAB = world:entity()
+			world:set(eAB, A, true)
+			world:set(eAB, B, true)
+
+			local count = 0
+			for id, data in world do
+				count += 1
+				if id == eA then
+					expect(data[A]).to.be.ok()
+					expect(data[B]).to.never.be.ok()
+				elseif id == eB then
+					expect(data[B]).to.be.ok()
+					expect(data[A]).to.never.be.ok()
+				elseif id == eAB then
+					expect(data[A]).to.be.ok()
+					expect(data[B]).to.be.ok()
+				else
+					error("unknown entity", id)
+				end
+			end
+
+			expect(count).to.equal(3)
+		end)
 	end)
 end
\ No newline at end of file
diff --git a/tests/test1.lua b/tests/test1.lua
index 0b031d3..7ff3b5a 100644
--- a/tests/test1.lua
+++ b/tests/test1.lua
@@ -110,6 +110,39 @@ TEST("world:query", function()
         CHECK(world:get(id, Health) == nil)
     end
 
+    do CASE "Should allow iterating the whole world" 
+        local world = jecs.World.new()
+
+        local A, B = world:entity(), world:entity()
+
+			local eA = world:entity()
+			world:set(eA, A, true)
+			local eB = world:entity()
+			world:set(eB, B, true)
+			local eAB = world:entity()
+			world:set(eAB, A, true)
+			world:set(eAB, B, true)
+
+			local count = 0
+			for id, data in world do
+				count += 1
+				if id == eA then
+					CHECK(data[A] == true)
+                    CHECK(data[B] == nil)
+				elseif id == eB then
+                    CHECK(data[B] == true)
+                    CHECK(data[A] == nil)
+				elseif id == eAB then
+					CHECK(data[A] == true)
+					CHECK(data[B] == true)
+				else
+					error("unknown entity", id)
+				end
+			end
+
+			CHECK(count == 3)
+    end
+
 end)
 
 FINISH()
\ No newline at end of file