From 52cf53176493b01d5edb91dac0e5e45b9aa82dcc Mon Sep 17 00:00:00 2001
From: Ukendio <ukendio@gmail.com>
Date: Sun, 28 Apr 2024 20:33:45 +0200
Subject: [PATCH] Prepared query

---
 lib/init.lua | 86 ++++++++++++++++++++++++++++++----------------------
 1 file changed, 50 insertions(+), 36 deletions(-)

diff --git a/lib/init.lua b/lib/init.lua
index ee62f82..e7e6ba8 100644
--- a/lib/init.lua
+++ b/lib/init.lua
@@ -361,44 +361,14 @@ local function getSmallestMap(componentIndex, components)
 	return s.sparse
 end
 
-function World.query(world: World, ...: i53): (() -> (number, ...any)) | () -> ()
-	
-	local compatibleArchetypes = {}
-	local components = { ... }
-	local archetypes = world.archetypes
-	local queryLength = #components
-	local firstArchetypeMap = getSmallestMap(world.componentIndex, components)
+local PreparedQuery = {}
+PreparedQuery.__index = PreparedQuery
 
-	if not firstArchetypeMap then 
-		return noop()
-	end
+function PreparedQuery:__iter() 
+	local compatibleArchetypes = self.compatibleArchetypes
+	local queryLength = self.queryLength
+	local components = self.components
 
-	for id in firstArchetypeMap do
-		local archetype = archetypes[id]
-        local columns = archetype.columns
-		local archetypeRecords = archetype.records
-        local indices = {}
-		local skip = false
-		
-		for i, componentId in components do 
-			local index = archetypeRecords[componentId]
-			if not index then 
-				skip = true
-				break
-			end
-			indices[i] = columns[index]
-		end
-
-		if skip then 
-			continue
-		end
-
-		table.insert(compatibleArchetypes, {
-			archetype = archetype,
-			indices = indices
-		})
-	end
-	
 	local lastArchetype, compatibleArchetype = next(compatibleArchetypes)
 	if not compatibleArchetype then 
 		return noop()
@@ -482,6 +452,50 @@ function World.query(world: World, ...: i53): (() -> (number, ...any)) | () -> (
 	end
 end
 
+function World.query(world: World, ...: i53)
+	local compatibleArchetypes = {}
+	local components = { ... }
+	local archetypes = world.archetypes
+	local queryLength = #components
+	local firstArchetypeMap = getSmallestMap(world.componentIndex, components)
+
+	if not firstArchetypeMap then 
+		return noop()
+	end
+
+	for id in firstArchetypeMap do
+		local archetype = archetypes[id]
+        local columns = archetype.columns
+		local archetypeRecords = archetype.records
+        local indices = {}
+		local skip = false
+		
+		for i, componentId in components do 
+			local index = archetypeRecords[componentId]
+			if not index then 
+				skip = true
+				break
+			end
+			indices[i] = columns[index]
+		end
+
+		if skip then 
+			continue
+		end
+
+		table.insert(compatibleArchetypes, {
+			archetype = archetype,
+			indices = indices
+		})
+	end
+	
+	return setmetatable({
+		queryLength = queryLength,
+		compatibleArchetypes = compatibleArchetypes,
+		components = components 
+	}, PreparedQuery)
+end
+
 function World.component(world: World) 
 	local id = world.nextId + 1
 	if id > HI_COMPONENT_ID then