From a5302a14fb9f7f85b38a4ffbdf697e8f2cc019a6 Mon Sep 17 00:00:00 2001
From: Ukendio <ukendio@gmail.com>
Date: Wed, 24 Apr 2024 01:14:43 +0200
Subject: [PATCH] Creator specific iterators for returns

---
 lib/init.lua | 102 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 76 insertions(+), 26 deletions(-)

diff --git a/lib/init.lua b/lib/init.lua
index 08b501e..d12e83b 100644
--- a/lib/init.lua
+++ b/lib/init.lua
@@ -322,8 +322,6 @@ function World.archetypesWith(world: World, componentId: i53)
 	return compatibleArchetypes
 end
 
-
-
 function World.query(world: World, ...: i53): () -> (number, ...any)
 	local compatibleArchetypes = {}
 	local components = { ... }
@@ -332,6 +330,61 @@ function World.query(world: World, ...: i53): () -> (number, ...any)
 	local a, b, c, d, e = ...
 	local firstArchetypeMap = world.componentIndex[components[1]]
 
+	if queryLength == 1 then 
+		local function single() 
+            local id = next(firstArchetypeMap)
+            local archetype = archetypes[id :: number]
+            local lastRow
+
+            return function(): any
+                local row, entity = next(archetype.entities, lastRow)
+                while row == nil do 
+                    id = next(firstArchetypeMap, id)
+                    if id == nil then 
+                        return
+                    end
+                    archetype = archetypes[id]
+                    row = next(archetype.entities, row)
+                end 
+                lastRow = row
+
+                return entity, archetype.columns[archetype.records[a]]
+            end
+        end
+
+        return single() 
+	elseif queryLength == 2 then 
+		for id in firstArchetypeMap do
+			local archetype = archetypes[id]
+			if archetype.records[b] then
+				table.insert(compatibleArchetypes, archetype)	
+			end
+		end
+
+		local function double(): any
+			local lastArchetype, archetype = next(compatibleArchetypes)
+			local lastRow
+	
+			return function(): any
+				local row = next(archetype.entities, lastRow)
+				while row == nil do 
+					lastArchetype, archetype = next(compatibleArchetypes, lastArchetype)
+					if lastArchetype == nil then 
+						return
+					end
+					row = next(archetype.entities, row)
+				end 
+				lastRow = row
+	
+				local entity = archetype.entities[row::number]
+				local columns = archetype.columns
+				local archetypeRecords = archetype.records
+				return entity, columns[archetypeRecords[a]], columns[archetypeRecords[b]]
+			end
+		end
+		return double()
+	end
+
 	for id in firstArchetypeMap do
 		local archetype = archetypes[id]
 		local archetypeRecords = archetype.records
@@ -346,30 +399,27 @@ function World.query(world: World, ...: i53): () -> (number, ...any)
 			table.insert(compatibleArchetypes, archetype)	
 		end
 	end
-    local lastArchetype, archetype = next(compatibleArchetypes)
 	
-    local lastRow 
-
-    local function queryNext(): (...any)
-        local row = next(archetype.entities, lastRow)
-        while row == nil do 
-            lastArchetype, archetype = next(compatibleArchetypes, lastArchetype)
-            if lastArchetype == nil then 
-                return 
-            end
-            row = next(archetype.entities, row)
-        end
-        lastRow = row
-        
+	local lastArchetype, archetype = next(compatibleArchetypes)
+	
+	local lastRow 
+	
+	local function queryNext(): (...any)
+		local row = next(archetype.entities, lastRow)
+		while row == nil do 
+			lastArchetype, archetype = next(compatibleArchetypes, lastArchetype)
+			if lastArchetype == nil then 
+				return 
+			end
+			row = next(archetype.entities, row)
+		end
+		lastRow = row
+		
 		local columns = archetype.columns	
 		local entityId = archetype.entities[row :: number]
 		local archetypeRecords = archetype.records
 
-		if queryLength == 1 then 
-			return entityId, columns[archetypeRecords[a]]
-		elseif queryLength == 2 then 
-			return entityId, columns[archetypeRecords[a]], columns[archetypeRecords[b]]
-		elseif queryLength == 3 then 
+		if queryLength == 3 then 
 			return entityId, 
 				columns[archetypeRecords[a]], 
 				columns[archetypeRecords[b]], 
@@ -395,12 +445,12 @@ function World.query(world: World, ...: i53): () -> (number, ...any)
 		end
 
 		return entityId, unpack(queryOutput, 1, queryLength)
-    end
-
-	return function()
-		-- consider this to be the iterator that gets invoked each iteration step
-		return queryNext()
 	end
+
+	return function() 
+		return queryNext() 
+	end
+	
 end
 
 return {