Return chainable iterator after QueryResult:without()

lib/World.lua

@@ -7,6 +7,8 @@ local assertValidComponent = Component.assertValidComponent
 local archetypeOf = archetypeModule.archetypeOf
 local areArchetypesCompatible = archetypeModule.areArchetypesCompatible
 
+local QueryResult = require(script.Parent.query)
+
 local ERROR_NO_ENTITY = "Entity doesn't exist, use world:contains to check if needed"
 
 --[=[
@@ -370,168 +372,20 @@ function World:get(id, ...)
 	return unpack(components, 1, length)
 end
 
---[=[
-	@class QueryResult
-
-	A result from the [`World:query`](/api/World#query) function.
-
-	Calling the table or the `next` method allows iteration over the results. Once all results have been returned, the
-	QueryResult is exhausted and is no longer useful.
-
-	```lua
-	for id, enemy, charge, model in world:query(Enemy, Charge, Model) do
-		-- Do something
-	end
-	```
-]=]
-local QueryResult = {}
-QueryResult.__index = QueryResult
-
-function QueryResult:__call()
-	return self._expand(self._next())
-end
-
-function QueryResult:__iter()
-	return function()
-		return self._expand(self._next())
-	end
-end
-
---[=[
-	Returns the next set of values from the query result. Once all results have been returned, the
-	QueryResult is exhausted and is no longer useful.
-
-	:::info
-	This function is equivalent to calling the QueryResult as a function. When used in a for loop, this is implicitly
-	done by the language itself.
-	:::
+local function noop() end
 
-	```lua
-	-- Using world:query in this position will make Lua invoke the table as a function. This is conventional.
-	for id, enemy, charge, model in world:query(Enemy, Charge, Model) do
-		-- Do something
-	end
-	```
-
-	If you wanted to iterate over the QueryResult without a for loop, it's recommended that you call `next` directly
-	instead of calling the QueryResult as a function.
-	```lua
-	local id, enemy, charge, model = world:query(Enemy, Charge, Model):next()
-	local id, enemy, charge, model = world:query(Enemy, Charge, Model)() -- Possible, but unconventional
-	```
-
-	@return id -- Entity ID
-	@return ...ComponentInstance -- The requested component values
-]=]
-function QueryResult:next()
-	return self._expand(self._next())
-end
-
-local snapshot = {
-	__iter = function(self): any
-		local i = 0
-		return function()
-			i += 1
-
-			local data = self[i]
-
-			if data then
-				return unpack(data, 1, data.n)
-			end
-			return
-		end
+local noopQuery = setmetatable({
+	next = noop,
+	snapshot = noop,
+	without = function(self)
+		return self
 	end,
-}
-
---[=[
-	Creates a "snapshot" of this query, draining this QueryResult and returning a list containing all of its results.
-
-	By default, iterating over a QueryResult happens in "real time": it iterates over the actual data in the ECS, so
-	changes that occur during the iteration will affect future results.
-
-	By contrast, `QueryResult:snapshot()` creates a list of all of the results of this query at the moment it is called,
-	so changes made while iterating over the result of `QueryResult:snapshot` do not affect future results of the
-	iteration.
-
-	Of course, this comes with a cost: we must allocate a new list and iterate over everything returned from the
-	QueryResult in advance, so using this method is slower than iterating over a QueryResult directly.
-
-	The table returned from this method has a custom `__iter` method, which lets you use it as you would use QueryResult
-	directly:
-
-	```lua
-		for entityId, health, player in world:query(Health, Player):snapshot() do
-
-		end
-	```
-
-	However, the table itself is just a list of sub-tables structured like `{entityId, component1, component2, ...etc}`.
-
-	@return {{entityId: number, component: ComponentInstance, component: ComponentInstance, component: ComponentInstance, ...}}
-]=]
-function QueryResult:snapshot()
-	local list = setmetatable({}, snapshot)
-
-	local function iter()
-		return self._next()
-	end
-
-	for entityId, entityData in iter do
-		if entityId then
-			table.insert(list, table.pack(self._expand(entityId, entityData)))
-		end
-	end
-
-	return list
-end
-
---[=[
-	Returns an iterator that will skip any entities that also have the given components.
-
-	:::tip
-	This is essentially equivalent to querying normally, using `World:get` to check if a component is present,
-	and using Lua's `continue` keyword to skip this iteration (though, using `:without` is faster).
-
-	This means that you should avoid queries that return a very large amount of results only to filter them down
-	to a few with `:without`. If you can, always prefer adding components and making your query more specific.
-	:::
-
-	@param ... Component -- The component types to filter against.
-	@return () -> (id, ...ComponentInstance) -- Iterator of entity ID followed by the requested component values
-
-	```lua
-	for id in world:query(Target):without(Model) do
-		-- Do something
-	end
-	```
-]=]
-function QueryResult:without(...)
-	local metatables = { ... }
-	return function(): any
-		while true do
-			local entityId, entityData = self._next()
-
-			if not entityId then
-				break
-			end
-
-			local skip = false
-			for _, metatable in ipairs(metatables) do
-				if entityData[metatable] then
-					skip = true
-					break
-				end
-			end
-
-			if skip then
-				continue
-			end
-
-			return self._expand(entityId, entityData)
-		end
-		return
-	end
-end
+	view = noop,
+}, {
+	__iter = function()
+		return noop
+	end,
+})
 
 --[=[
 	Performs a query against the entities in this World. Returns a [QueryResult](/api/QueryResult), which iterates over
@@ -552,6 +406,7 @@ end
 	@param ... Component -- The component types to query. Only entities with *all* of these components will be returned.
 	@return QueryResult -- See [QueryResult](/api/QueryResult) docs.
 ]=]
+
 function World:query(...)
 	debug.profilebegin("World:query")
 	assertValidComponent((...), 1)
@@ -571,10 +426,7 @@ function World:query(...)
 
 	if next(compatibleArchetypes) == nil then
 		-- If there are no compatible storages avoid creating our complicated iterator
-		return setmetatable({
-			_expand = function() end,
-			_next = function() end,
-		}, QueryResult)
+		return noopQuery
 	end
 
 	local queryOutput = table.create(queryLength)
@@ -625,9 +477,12 @@ function World:query(...)
 	local function nextItem()
 		local entityId, entityData
 
+		local storages = self._storages
 		repeat
-			if self._storages[storageIndex][currentCompatibleArchetype] then
-				entityId, entityData = next(self._storages[storageIndex][currentCompatibleArchetype], lastEntityId)
+			local nextStorage = storages[storageIndex]
+			local currently = nextStorage[currentCompatibleArchetype]
+			if nextStorage[currentCompatibleArchetype] then
+				entityId, entityData = next(currently, lastEntityId)
 			end
 
 			while entityId == nil do
@@ -636,7 +491,7 @@ function World:query(...)
 				if currentCompatibleArchetype == nil then
 					storageIndex += 1
 
-					local nextStorage = self._storages[storageIndex]
+					nextStorage = storages[storageIndex]
 
 					if nextStorage == nil or next(nextStorage) == nil then
 						return
@@ -649,24 +504,23 @@ function World:query(...)
 					end
 
 					continue
-				elseif self._storages[storageIndex][currentCompatibleArchetype] == nil then
+				elseif nextStorage[currentCompatibleArchetype] == nil then
 					continue
 				end
 
-				entityId, entityData = next(self._storages[storageIndex][currentCompatibleArchetype])
+				entityId, entityData = next(nextStorage[currentCompatibleArchetype])
 			end
+
 			lastEntityId = entityId
 
 		until seenEntities[entityId] == nil
 
 		seenEntities[entityId] = true
+
 		return entityId, entityData
 	end
 
-	return setmetatable({
-		_expand = expand,
-		_next = nextItem,
-	}, QueryResult)
+	return QueryResult.new(expand, nextItem)
 end
 
 local function cleanupQueryChanged(hookState)

lib/World.spec.lua

@@ -505,9 +505,10 @@ return function()
 				})
 			)
 
-			local snapshot = world:query(Health, Player):snapshot()
+			local query = world:query(Health, Player)
+			local snapshot = query:snapshot()
 
-			for entityId, health, player in world:query(Health, Player):snapshot() do
+			for entityId, health, player in snapshot do
 				expect(type(entityId)).to.equal("number")
 				expect(type(player.name)).to.equal("string")
 				expect(type(health.value)).to.equal("number")
@@ -521,6 +522,8 @@ return function()
 			else
 				expect(snapshot[2][1]).to.equal(1)
 			end
+
+			expect(#world:query(Player):without(Poison):snapshot()).to.equal(1)
 		end)
 
 		it("should not invalidate iterators", function()