Finish Collection implementation

* Depends on https://github.com/fortressforever/fortressforever/pull/134

maps/includes/util/collection.lua

@@ -1,6 +1,7 @@
 -- backwards compatibility for the Collection class
 
 local Class = require "util.class"
+require "util.utils"
 
 local function setup_items_key(self)
 	-- keep a usable reference to the base metatable
@@ -29,12 +30,6 @@ local function setup_items_key(self)
 	})
 end
 
-local Collection = Class(function(self)
-	setup_items_key(self)
-	self.entities = {}
-end)
-
-
 local function get_param_as_table(param)
 	if param == nil then return {} end
 	if type(param) == "table" then return param end
@@ -51,6 +46,67 @@ local function to_collection_index(lua_index)
 	return lua_index-1
 end
 
+local Collection = Class(function(self, entity_or_entities)
+	setup_items_key(self)
+	self.entities = get_param_as_table(entity_or_entities)
+end)
+
+Collection.filters = {
+	[CF.kNone] = function(entity) return true end,
+	[CF.kPlayers] = function(entity) return IsPlayer(entity) end,
+	[CF.kHumanPlayers] = function(entity) return IsPlayer(entity) and not CastToPlayer(entity):IsBot() end,
+	[CF.kBotPlayers] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):IsBot() end,
+	[CF.kPlayerScout] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kScout end,
+	[CF.kPlayerSniper] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kSniper end,
+	[CF.kPlayerSoldier] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kSoldier end,
+	[CF.kPlayerDemoman] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kDemoman end,
+	[CF.kPlayerMedic] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kMedic end,
+	[CF.kPlayerHWGuy] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kHwguy end,
+	[CF.kPlayerPyro] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kPyro end,
+	[CF.kPlayerSpy] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kSpy end,
+	[CF.kPlayerEngineer] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kEngineer end,
+	[CF.kPlayerCivilian] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kCivilian end,
+	[CF.kPlayerScout] = function(entity) return IsPlayer(entity) and CastToPlayer(entity):GetClass() == Player.kScout end,
+	[CF.kTeams] = function(entity) return IsEntity(entity) and entity:GetTeamId() >= Team.kSpectator and entity:GetTeamId() <= Team.kGreen end,
+	[CF.kTeamSpec] = function(entity) return IsEntity(entity) and entity:GetTeamId() == Team.kSpectator end,
+	[CF.kTeamBlue] = function(entity) return IsEntity(entity) and entity:GetTeamId() == Team.kBlue end,
+	[CF.kTeamRed] = function(entity) return IsEntity(entity) and entity:GetTeamId() == Team.kRed end,
+	[CF.kTeamYellow] = function(entity) return IsEntity(entity) and entity:GetTeamId() == Team.kYellow end,
+	[CF.kTeamGreen] = function(entity) return IsEntity(entity) and entity:GetTeamId() == Team.kGreen end,
+	[CF.kProjectiles] = function(entity) return IsProjectile(entity) end,
+	[CF.kGrenades] = function(entity) return IsGrenade(entity) end,
+	[CF.kInfoScipts] = function(entity) return IsInfoScript(entity) end,
+	[CF.kInfoScript_Carried] = function(entity) return IsInfoScript(entity) and CastToInfoScript(entity):IsCarried() end,
+	[CF.kInfoScript_Dropped] = function(entity) return IsInfoScript(entity) and CastToInfoScript(entity):IsDropped() end,
+	[CF.kInfoScript_Returned] = function(entity) return IsInfoScript(entity) and CastToInfoScript(entity):IsReturned() end,
+	[CF.kInfoScript_Active] = function(entity) return IsInfoScript(entity) and CastToInfoScript(entity):IsActive() end,
+	[CF.kInfoScript_Inactive] = function(entity) return IsInfoScript(entity) and CastToInfoScript(entity):IsInactive() end,
+	[CF.kInfoScript_Removed] = function(entity) return IsInfoScript(entity) and CastToInfoScript(entity):IsRemoved() end,
+	[CF.kTraceBlockWalls] = function(entity) return true end, -- not applicable for this type of filtering
+	[CF.kBuildables] = function(entity) return IsBuildable(entity) end,
+	[CF.kDispenser] = function(entity) return IsDispenser(entity) end,
+	[CF.kSentrygun] = function(entity) return IsSentrygun(entity) end,
+	[CF.kDetpack] = function(entity) return IsDetpack(entity) end,
+	[CF.kJumpPad] = function(entity) return IsJumpPad(entity) end,
+}
+
+function Collection.PassesFilter(entity, filter)
+	if not filter then filter = CF.kNone end
+	assert(Collection.filters[filter], "Unknown filter: "..tostring(filter))
+	return Collection.filters[filter](entity)
+end
+
+function Collection.PassesFilters(entity, filters)
+	if not filters then return true end
+	filters = get_param_as_table(filters)
+	for _,filter in ipairs(filters) do
+		if not Collection.PassesFilter(entity, filter) then
+			return false
+		end
+	end
+	return true
+end
+
 function Collection:AddItem(entity_or_entities)
 	local entities_to_add = get_param_as_table(entity_or_entities)
 
@@ -59,6 +115,16 @@ function Collection:AddItem(entity_or_entities)
 	end
 end
 
+function Collection:AddFiltered(entity_or_entities, filters)
+	local entities_to_add = get_param_as_table(entity_or_entities)
+
+	for i,entity_to_add in ipairs(entities_to_add) do
+		if Collection.PassesFilters(entity_to_add, filters) then
+			table.insert(self.entities, entity_to_add)
+		end
+	end
+end
+
 function Collection:RemoveItem(entity_or_entities)
 	local entities_to_find = get_param_as_table(entity_or_entities)
 
@@ -71,15 +137,21 @@ function Collection:RemoveItem(entity_or_entities)
 end
 
 function Collection:RemoveAllItems()
-	for k in ipairs(self.entities) do
-		self.entities[k] = nil
-	end
+	table.clear(self.entities)
 end
 
 function Collection:Count()
 	return #self.entities
 end
 
+function Collection:NumItems()
+	return self:Count()
+end
+
+function Collection:IsEmpty()
+	return self:Count() == 0
+end
+
 function Collection:HasItem(entity_or_entities)
 	local entities_to_find = get_param_as_table(entity_or_entities)
 
@@ -116,4 +188,41 @@ function Collection:Element(i)
 	return self.entities[to_lua_index(i)]
 end
 
+function Collection:GetByFilter(filters)
+	filters = get_param_as_table(filters)
+
+	-- optimization for players
+	local players_only = false
+	for _,filter in ipairs(filters) do
+		if filter >= CF.kPlayers and filter <= CF.kPlayerScout then
+			players_only = true
+			break
+		end
+	end
+
+	if players_only then
+		self:AddFiltered(GetPlayers(), filters)
+	else
+		for ent_id, entity in ipairs(GlobalEntityList) do
+			self:AddFiltered(entity, filters)
+		end
+	end
+end
+
+function Collection:GetByName(name_or_names, filters)
+	local names_to_find = get_param_as_table(name_or_names)
+	filters = get_param_as_table(filters)
+
+	for i,name_to_find in ipairs(names_to_find) do
+		self:AddFiltered(GetEntitiesByName(name_to_find), filters)
+	end
+end
+
+function Collection:GetInSphere(entity_or_origin, radius, filters)
+	filters = get_param_as_table(filters)
+	local origin = IsEntity(entity_or_origin) and entity_or_origin:GetOrigin() or entity_or_origin
+	local ignore_walls = not table.contains(filters, CF.kTraceBlockWalls)
+	self:AddFiltered(GetEntitiesInSphere(origin, radius, ignore_walls), filters)
+end
+
 return Collection