diff --git a/src/tbtr_template_vehicle.cpp b/src/tbtr_template_vehicle.cpp index a130e1a02c..160b02214e 100644 --- a/src/tbtr_template_vehicle.cpp +++ b/src/tbtr_template_vehicle.cpp @@ -40,6 +40,8 @@ #include "newgrf_engine.h" #include "newgrf_cargo.h" +#include "3rdparty/robin_hood/robin_hood.h" + #include "safeguards.h" TemplatePool _template_pool("TemplatePool"); @@ -48,7 +50,10 @@ INSTANTIATE_POOL_METHODS(Template) TemplateReplacementPool _template_replacement_pool("TemplateReplacementPool"); INSTANTIATE_POOL_METHODS(TemplateReplacement) -btree::btree_map _template_replacement_index; +robin_hood::unordered_flat_map _template_replacement_index; +robin_hood::unordered_flat_map _template_replacement_index_recursive; + +static void ReindexTemplateReplacementsRecursive(); void TemplateVehicleImageDimensions::SetFromTrain(const Train *t) { @@ -141,14 +146,16 @@ TemplateReplacement::~TemplateReplacement() if (CleaningPool()) return; _template_replacement_index.erase(this->Group()); + ReindexTemplateReplacementsRecursive(); } void TemplateReplacement::PreCleanPool() { _template_replacement_index.clear(); + _template_replacement_index_recursive.clear(); } -TemplateReplacement* GetTemplateReplacementByGroupID(GroupID gid) +TemplateReplacement *GetTemplateReplacementByGroupID(GroupID gid) { if (GetTemplateIDByGroupID(gid) == INVALID_TEMPLATE) return nullptr; @@ -169,14 +176,9 @@ TemplateID GetTemplateIDByGroupID(GroupID gid) TemplateID GetTemplateIDByGroupIDRecursive(GroupID gid) { - while (gid != INVALID_GROUP) { - auto iter = _template_replacement_index.find(gid); - if (iter != _template_replacement_index.end()) return iter->second; - const Group *g = Group::GetIfValid(gid); - if (g == nullptr) break; - gid = Group::Get(gid)->parent; - } - return INVALID_TEMPLATE; + auto iter = _template_replacement_index_recursive.find(gid); + if (iter == _template_replacement_index_recursive.end()) return INVALID_TEMPLATE; + return iter->second; } bool IssueTemplateReplacement(GroupID gid, TemplateID tid) @@ -187,10 +189,12 @@ bool IssueTemplateReplacement(GroupID gid, TemplateID tid) /* Then set the new TemplateVehicle and return */ tr->SetTemplate(tid); _template_replacement_index[gid] = tid; + ReindexTemplateReplacementsRecursive(); return true; } else if (TemplateReplacement::CanAllocateItem()) { tr = new TemplateReplacement(gid, tid); _template_replacement_index[gid] = tid; + ReindexTemplateReplacementsRecursive(); return true; } else { return false; @@ -228,4 +232,24 @@ void ReindexTemplateReplacements() for (const TemplateReplacement *tr : TemplateReplacement::Iterate()) { _template_replacement_index[tr->group] = tr->sel_template; } + ReindexTemplateReplacementsRecursive(); +} + +static void ReindexTemplateReplacementsRecursive() +{ + _template_replacement_index_recursive.clear(); + for (const Group *group : Group::Iterate()) { + if (group->vehicle_type != VEH_TRAIN) continue; + + const Group *g = group; + while (true) { + auto iter = _template_replacement_index.find(g->index); + if (iter != _template_replacement_index.end()) { + _template_replacement_index_recursive[group->index] = iter->second; + break; + } + if (g->parent == INVALID_GROUP) break; + g = Group::Get(g->parent); + } + } }