diff --git a/src/TechTreeImproved.cpp b/src/TechTreeImproved.cpp
index 84a394fd54a8d4b918f06aad6159266f8bee8314..7cdd25589f2c4343a28acec3969d4c94d6bd968e 100644
--- a/src/TechTreeImproved.cpp
+++ b/src/TechTreeImproved.cpp
@@ -26,12 +26,31 @@ void add_requirement(BuildDescription & description, json & requirement)
     }
     else if (type == "unitCount")
     {
+        // Units that needs to be completed at the unit location -> Addons
         if (requirement["state"] == "CompleteOnlyAtUnit")
         {
             for (auto & unit : requirement["unit"])
             {
                 //std::cout << "Addon: " << sc2::UnitTypeToName(static_cast<sc2::UNIT_TYPEID>(unit)) << " (" << unit << ")" << std::endl;
-                description.addons_needed.push_back(static_cast<sc2::UNIT_TYPEID>(unit));
+                sc2::UNIT_TYPEID unit_typeid = static_cast<sc2::UNIT_TYPEID>(unit);
+                if (unit_typeid == sc2::UNIT_TYPEID::TERRAN_TECHLAB)
+                {
+                    switch (description.producer_type)
+                    {
+                    case sc2::UNIT_TYPEID::TERRAN_BARRACKS:
+                        unit_typeid = sc2::UNIT_TYPEID::TERRAN_BARRACKSTECHLAB;
+                        break;
+                    case sc2::UNIT_TYPEID::TERRAN_FACTORY:
+                        unit_typeid = sc2::UNIT_TYPEID::TERRAN_FACTORYTECHLAB;
+                        break;
+                    case sc2::UNIT_TYPEID::TERRAN_STARPORT:
+                        unit_typeid = sc2::UNIT_TYPEID::TERRAN_STARPORTTECHLAB;
+                        break;
+                    default:
+                        std::cout << "Unknown producer type for TECHLAB addon: " << sc2::UnitTypeToName(unit_typeid) << std::endl;
+                    }
+                }
+                description.addons_needed.push_back(unit_typeid);
             }
         }
         else if (requirement["state"] == "CompleteOnly")