Add support for parsing information about upgrades from JSON

src/TechTreeImproved.cpp

@@ -14,6 +14,49 @@ std::string id_to_string(int id)
     return sc2::UnitTypeToName(static_cast<sc2::UNIT_TYPEID>(id));
 }
 
+void add_requirement(ResearchDescription & description, json & requirement)
+{
+    std::string type = requirement["type"];
+
+    if (type == "and")
+    {
+        for (auto & subrequirement : requirement["operands"])
+        {
+            add_requirement(description, subrequirement);
+        }
+    }
+    else if (type == "not")
+    {
+        // Ignore this. This is mostly used for: "We cannot do upgrade X if upgrade X is already running somewhere."
+    }
+    else if (type == "unitCount")
+    {
+        if (requirement["state"] == "CompleteOnly")
+        {
+            for (auto & unit : requirement["unit"])
+            {
+                //std::cout << "Just building: " << sc2::UnitTypeToName(static_cast<sc2::UNIT_TYPEID>(unit)) << " (" << unit << ")" << std::endl;
+                description.buildings_needed.push_back(static_cast<sc2::UNIT_TYPEID>(unit));
+            }
+        }
+        else
+        {
+            std::cout << "Unexpected state: " << requirement["state"] << std::endl;
+        }
+    }
+    else if (type == "eq")
+    {
+        // TODO: Should we be more careful here?
+        sc2::UPGRADE_ID id = static_cast<sc2::UPGRADE_ID>(requirement["operands"][0]["upgrade"]);
+        int count = static_cast<int>(requirement["operands"][1]["value"]);
+
+        if (count == 1)
+        {
+            description.upgrades_needed.push_back(id);
+        }
+    }
+}
+
 void add_requirement(BuildDescription & description, json & requirement)
 {
     std::string type = requirement["type"];
@@ -47,7 +90,7 @@ void add_requirement(BuildDescription & description, json & requirement)
                         unit_typeid = sc2::UNIT_TYPEID::TERRAN_STARPORTTECHLAB;
                         break;
                     default:
-                        std::cout << "Unknown producer type for TECHLAB addon: " << sc2::UnitTypeToName(unit_typeid) << std::endl;
+                        std::wcout << "Unknown producer type for TECHLAB addon: " << sc2::UnitTypeToName(unit_typeid) << std::endl;
                     }
                 }
                 description.addons_needed.push_back(unit_typeid);
@@ -81,6 +124,20 @@ void parse_build_description(BuildDescription & description, json & build_item)
     }
 }
 
+void parse_research_item(ResearchDescription & description, json & research_item)
+{
+    //std::cout << "Parsing " << research_item["upgradeName"] << std::endl;
+
+    description.result_type = static_cast<sc2::UPGRADE_ID>(research_item["upgrade"]);
+    description.ability_used = static_cast<sc2::ABILITY_ID>(research_item["ability"]);
+
+    if (research_item.find("requires") != research_item.end() && research_item["requires"].size() > 0)
+    {
+        auto & requires = research_item["requires"][0];
+        add_requirement(description, requires);
+    }
+}
+
 void TechTreeImproved::parse_unit(json::iterator it)
 {
     sc2::UNIT_TYPEID producer_id = string_to_id(it.key());
@@ -109,7 +166,28 @@ void TechTreeImproved::parse_unit(json::iterator it)
         }
     }
 
-    producer_to_data[producer_id] = build_descriptions;
+    // TODO: Use the result from the call to find for actually looking up data later, instead of searching twice
+    if (it.value().find("researches") != it.value().end())
+    {
+        //std::cout << "Found upgrades on unit " << it.value()["name"] << std::endl;
+
+        for (auto & research_item : it.value()["researches"])
+        {
+            ResearchDescription description;
+            description.producer_type = producer_id;
+
+            parse_research_item(description, research_item);
+
+            if (upgrade_to_data.count(description.result_type))
+            {
+                upgrade_to_data[description.result_type].push_back(description);
+            }
+            else
+            {
+                upgrade_to_data[description.result_type] = { description };
+            }
+        }
+    }
 }
 
 bool TechTreeImproved::LoadData() {

src/TechTreeImproved.h

@@ -12,6 +12,7 @@ struct BuildDescription
 {
     sc2::UNIT_TYPEID producer_type;
     sc2::UNIT_TYPEID result_type;
+    // TODO: Are these 2 used?
     sc2::AbilityID ability_used;
     float time;
 
@@ -19,11 +20,20 @@ struct BuildDescription
     std::vector<sc2::UNIT_TYPEID> addons_needed;
 };
 
+struct ResearchDescription
+{
+    sc2::UPGRADE_ID result_type;
+    sc2::UNIT_TYPEID producer_type;
+    sc2::AbilityID ability_used;
+
+    std::vector<sc2::UPGRADE_ID> upgrades_needed;
+    std::vector<sc2::UNIT_TYPEID> buildings_needed;
+};
+
 class TechTreeImproved
 {
-    std::vector<BuildDescription> data;
-    std::map<sc2::UNIT_TYPEID, std::vector<BuildDescription>> producer_to_data;
     std::map<sc2::UNIT_TYPEID, std::vector<BuildDescription>> result_to_data;
+    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 {};