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,169 +372,6 @@ 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.
-	:::
-
-	```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
-	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
-
 --[=[
 	Performs a query against the entities in this World. Returns a [QueryResult](/api/QueryResult), which iterates over
 	the results of the query.
@@ -553,99 +392,7 @@ end
 	@return QueryResult -- See [QueryResult](/api/QueryResult) docs.
 ]=]
 function World:query(...)
-	debug.profilebegin("World:query")
-	assertValidComponent((...), 1)
-
-	local metatables = { ... }
-	local queryLength = select("#", ...)
-
-	local archetype = archetypeOf(...)
-
-	if self._queryCache[archetype] == nil then
-		self:_newQueryArchetype(archetype)
-	end
-
-	local compatibleArchetypes = self._queryCache[archetype]
-
-	debug.profileend()
-
-	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)
-	end
-
-	local queryOutput = table.create(queryLength)
-
-	local function expand(entityId, entityData)
-		if not entityId then
-			return
-		end
-
-		for i, metatable in ipairs(metatables) do
-			queryOutput[i] = entityData[metatable]
-		end
-
-		return entityId, unpack(queryOutput, 1, queryLength)
-	end
-
-	local currentCompatibleArchetype = next(compatibleArchetypes)
-	local lastEntityId
-	local storageIndex = 1
-
-	if self._pristineStorage == self._storages[1] then
-		self:_markStorageDirty()
-	end
-
-	local seenEntities = {}
-
-	local function nextItem()
-		local entityId, entityData
-
-		repeat
-			if self._storages[storageIndex][currentCompatibleArchetype] then
-				entityId, entityData = next(self._storages[storageIndex][currentCompatibleArchetype], lastEntityId)
-			end
-
-			while entityId == nil do
-				currentCompatibleArchetype = next(compatibleArchetypes, currentCompatibleArchetype)
-
-				if currentCompatibleArchetype == nil then
-					storageIndex += 1
-
-					local nextStorage = self._storages[storageIndex]
-
-					if nextStorage == nil or next(nextStorage) == nil then
-						return
-					end
-
-					currentCompatibleArchetype = nil
-
-					if self._pristineStorage == nextStorage then
-						self:_markStorageDirty()
-					end
-
-					continue
-				elseif self._storages[storageIndex][currentCompatibleArchetype] == nil then
-					continue
-				end
-
-				entityId, entityData = next(self._storages[storageIndex][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(self, ...)
 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()