From 62ac6e62c009c079259a47065c807f9a2feb99c8 Mon Sep 17 00:00:00 2001
From: Marcus <ukendio@gmail.com>
Date: Mon, 16 Sep 2024 21:33:12 +0200
Subject: [PATCH] Iter should not drain by default (#123)

---
 src/init.luau   | 21 ++++++++++++++-------
 test/tests.luau | 22 ++++++++++++++++++++--
 wally.toml      |  2 +-
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/init.luau b/src/init.luau
index b5af96f..2a8e0e9 100644
--- a/src/init.luau
+++ b/src/init.luau
@@ -956,14 +956,16 @@ local EMPTY_QUERY = {
 setmetatable(EMPTY_QUERY, EMPTY_QUERY)
 
 local function query_init(query)
-    local world_query_iter_next = query.iter_next
-    if world_query_iter_next then
-        return world_query_iter_next
+    local world_query_iter_next
+
+    if query.should_drain then
+        world_query_iter_next = query.iter_next
+        if world_query_iter_next then
+            return world_query_iter_next
+        end
     end
+
     local compatible_archetypes = query.compatible_archetypes
-    local ids = query.ids
-    local A, B, C, D, E, F, G, H, I = unpack(ids)
-    local a, b, c, d, e, f, g, h
     local lastArchetype = 1
 	local archetype = compatible_archetypes[1]
 	if not archetype then
@@ -972,8 +974,12 @@ local function query_init(query)
 	local columns = archetype.columns
 	local entities = archetype.entities
 	local i = #entities
-
 	local records = archetype.records
+
+    local ids = query.ids
+    local A, B, C, D, E, F, G, H, I = unpack(ids)
+    local a, b, c, d, e, f, g, h
+
 	if not B then
 		a = columns[records[A].column]
 	elseif not C then
@@ -1216,6 +1222,7 @@ end
 local function query_drain(query)
     local query_iter_next = query_init(query)
     query.next = query_iter_next
+    query.should_drain = true
     return query
 end
 
diff --git a/test/tests.luau b/test/tests.luau
index d95ff1f..d50b047 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -239,6 +239,24 @@ TEST("world:add()", function()
 end)
 
 TEST("world:query()", function()
+    do CASE "multiple iter"
+        local world = jecs.World.new()
+        local A = world:component()
+        local B = world:component()
+        local e = world:entity()
+        world:add(e, A, "a")
+        world:add(e, B)
+        local q = world:query(A, B)
+        local counter = 0
+        for x in q:iter() do
+            counter += 1
+        end
+        for x in q:iter() do
+            counter += 1
+        end
+        print(counter)
+        CHECK(counter == 2)
+    end
     do CASE "tag"
         local world = jecs.World.new()
         local A = world:entity()
@@ -545,8 +563,8 @@ TEST("world:query()", function()
 
 	do CASE("should error when setting invalid pair")
 		local world = jecs.World.new()
-		local Eats = world:entity()
-		local Apples = world:entity()
+		local Eats = world:component()
+		local Apples = world:component()
 		local bob = world:entity()
 
 		world:delete(Apples)
diff --git a/wally.toml b/wally.toml
index 7aab522..5631efb 100644
--- a/wally.toml
+++ b/wally.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ukendio/jecs"
-version = "0.3.0-rc.1"
+version = "0.3.0-rc.2"
 registry = "https://github.com/UpliftGames/wally-index"
 realm = "shared"
 include = [