diff --git a/src/TechTree.cpp b/src/TechTree.cpp index 8f7cccfbf85b5f196eb8ae2438bfbfd866e7450a..62df02425088577508c09cc53ee09220a3d189ee 100644 --- a/src/TechTree.cpp +++ b/src/TechTree.cpp @@ -41,16 +41,45 @@ void TechTree::onStart() for (const BuildDescription & description : howToBuild) { data.whatBuilds.push_back(UnitType(description.producer_type, m_bot)); + for (sc2::UNIT_TYPEID unit_typeid : description.buildings_needed) { data.requiredUnits.push_back(UnitType(unit_typeid, m_bot)); } + for (sc2::UNIT_TYPEID unit_typeid : description.addons_needed) { data.requiredAddons.push_back(UnitType(unit_typeid, m_bot)); } } } + + for (std::pair <const sc2::UpgradeID, TypeData> & pair : m_upgradeData) + { + const sc2::UpgradeID id = pair.first; + TypeData & data = pair.second; + + data.whatBuilds.clear(); + data.requiredUnits.clear(); + data.requiredUpgrades.clear(); + + + for (const ResearchDescription & description : tree.HowToResearch(id)) + { + data.buildAbility = description.ability_used; + data.whatBuilds.push_back(UnitType(description.producer_type, m_bot)); + + for (sc2::UNIT_TYPEID unit_typeid : description.buildings_needed) + { + data.requiredUnits.push_back(UnitType(unit_typeid, m_bot)); + } + + for (sc2::UPGRADE_ID upgrade_id : description.upgrades_needed) + { + data.requiredUpgrades.push_back(upgrade_id); + } + } + } } diff --git a/src/TechTreeImproved.cpp b/src/TechTreeImproved.cpp index 6b07ff929a46efad1f31e41a297181b800ed89e6..14ef7d55b3918ef3663bae426f41caf122be0ebc 100644 --- a/src/TechTreeImproved.cpp +++ b/src/TechTreeImproved.cpp @@ -224,6 +224,19 @@ const std::vector<BuildDescription> & TechTreeImproved::HowToBuild(sc2::UnitType else { std::cout << "No information about unit type " << sc2::UnitTypeToName(unit) << " (" << static_cast<int>(unit) << ")" << std::endl; - return empty; + return empty_build; } -} \ No newline at end of file +} + +const std::vector<ResearchDescription> & TechTreeImproved::HowToResearch(sc2::UpgradeID upgrade) const +{ + if (upgrade_to_data.count(upgrade)) + { + return upgrade_to_data.at(upgrade); + } + else + { + std::cout << "No information about upgrade type " << sc2::UpgradeIDToName(upgrade) << " (" << static_cast<int>(upgrade) << ")" << std::endl; + return empty_research; + } +} diff --git a/src/TechTreeImproved.h b/src/TechTreeImproved.h index 344a8094183d7c6df49fc147f986c292fa19fbd8..5e9710a86dbfd096f84492665658914524425bf0 100644 --- a/src/TechTreeImproved.h +++ b/src/TechTreeImproved.h @@ -26,8 +26,8 @@ struct ResearchDescription sc2::UNIT_TYPEID producer_type; sc2::AbilityID ability_used; - std::vector<sc2::UPGRADE_ID> upgrades_needed; std::vector<sc2::UNIT_TYPEID> buildings_needed; + std::vector<sc2::UPGRADE_ID> upgrades_needed; }; class TechTreeImproved @@ -36,7 +36,8 @@ class TechTreeImproved std::map<sc2::UPGRADE_ID, std::vector<ResearchDescription>> upgrade_to_data; // If there is no BuildDescription for a given type, a reference to tihs list is returned. - const std::vector<BuildDescription> empty {}; + const std::vector<BuildDescription> empty_build {}; + const std::vector<ResearchDescription> empty_research {}; void parse_unit(nlohmann::json::iterator it); public: @@ -44,4 +45,5 @@ public: bool LoadData(); // Given a unit, how can we build it? const std::vector<BuildDescription> & HowToBuild(sc2::UnitTypeID unit) const; + const std::vector<ResearchDescription> & HowToResearch(sc2::UpgradeID upgrade) const; }; \ No newline at end of file