diff --git a/maps/includes/util/collection.lua b/maps/includes/util/collection.lua index 5a3cb31..ea39e02 100644 --- a/maps/includes/util/collection.lua +++ b/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 \ No newline at end of file