sc2ai 0.0.4 → 0.0.6

data/lib/sc2ai/player/debug.rb

@@ -14,6 +14,7 @@ module Sc2
       # @param debug_command [Api::DebugCommand]
       # @return [void]
       def queue_debug_command(debug_command)
+        @debug_command_queue ||= []
         @debug_command_queue << debug_command
       end
 

data/lib/sc2ai/player/game_state.rb

@@ -16,9 +16,12 @@ module Sc2
 
       extend Forwardable
 
-      # @!attribute game_loop
-      #   @return [Integer] current game loop
-      def_delegators :observation, :game_loop
+      attr_writer :game_loop
+
+      # @return [Integer] current game loop
+      def game_loop
+        @game_loop || 0
+      end
 
       # @!attribute game_info [rw]
       # Access useful game information. Used in parsed pathing grid, terrain height, placement grid.
@@ -38,7 +41,7 @@ module Sc2
       attr_accessor :game_info_loop
 
       # Determines if your game_info will be refreshed at this moment
-      # Has a hard-capped refresh of only ever 2 steps
+      # Has a hard-capped refresh of only ever 4 steps
       # In general game_info is only refreshed Player::Bot reads from pathing_grid or placement_grid
       # @return [Boolean]
       def game_info_stale?
@@ -47,7 +50,7 @@ module Sc2
 
         # Note: No minimum step count set anymore
         # We can do something like, only updating every 2+ frames:
-        game_info_loop + 2 <= game_loop
+        game_info_loop + 4 <= game_loop
       end
 
       # @!attribute data

data/lib/sc2ai/player/geometry.rb

@@ -20,7 +20,7 @@ module Sc2
       # @return [Integer]
       def map_width
         # bot.bot.game_info
-        bot.game_info.start_raw.map_size.x
+        @map_width ||= bot.game_info.start_raw.map_size.x
       end
 
       # Gets the map tile height. Range is 1-255.
@@ -28,7 +28,7 @@ module Sc2
       # @return [Integer]
       def map_height
         # bot.bot.game_info
-        bot.game_info.start_raw.map_size.y
+        @map_height ||= bot.game_info.start_raw.map_size.y
       end
 
       # Returns zero to map_width as range
@@ -83,6 +83,15 @@ module Sc2
         @parsed_placement_grid
       end
 
+      # Whether this tile is where an expansion is supposed to be placed.
+      # To see if a unit/structure is blocking an expansion, pass their coordinates to this method.
+      # @param x [Float, Integer]
+      # @param y [Float, Integer]
+      # @return [Boolean] true if location has creep on it
+      def expo_placement?(x:, y:)
+        expo_placement_grid[y.to_i, x.to_i] == 1
+      end
+
       # Returns a grid where ony the expo locations are marked
       # @return [Numo::Bit]
       def expo_placement_grid
@@ -259,6 +268,13 @@ module Sc2
         parsed_terrain_height[y.to_i, x.to_i]
       end
 
+      # Returns the terrain height (z) at position x and y for a point
+      # @param position [Sc2::Position]
+      # @return [Float] z axis position between -16 and 16
+      def terrain_height_for_pos(position)
+        terrain_height(x: position.x, y: position.y)
+      end
+
       # Returns a parsed terrain_height from bot.game_info.start_raw.
       # Each value in [row][column] holds a float value which is the z height
       # @return [Numo::SFloat] Numo array
@@ -334,15 +350,17 @@ module Sc2
       end
 
       # Provides parsed minimap representation of creep spread
+      # Caches for 4 frames
       # @return [Numo::Bit] Numo array
       def parsed_creep
-        if @parsed_creep.nil?
+        if @parsed_creep.nil? || @parsed_creep[1] + 4 < bot.game_loop
           image_data = bot.observation.raw_data.map_state.creep
           # Fix endian for Numo bit parser
           data = image_data.data.unpack("b*").pack("B*")
-          @parsed_creep = ::Numo::Bit.from_binary(data, [image_data.size.y, image_data.size.x])
+          result = ::Numo::Bit.from_binary(data, [image_data.size.y, image_data.size.x])
+          @parsed_creep = [result, bot.game_loop]
         end
-        @parsed_creep
+        @parsed_creep[0]
       end
 
       # TODO: Removing. Better name or more features for this? Maybe check nearest units.
@@ -537,16 +555,75 @@ module Sc2
       # @param base [Api::Unit, Sc2::Position] base Unit or Position
       # @return [Sc2::UnitGroup] UnitGroup of minerals for the base
       def minerals_for_base(base)
-        # resources_for_base contains what we need, but slice neutral.minerals,
-        # so that only active patches remain
-        bot.neutral.minerals.slice(*resources_for_base(base).minerals.tags)
+        base_resources = resources_for_base(base)
+        cached_tags = base_resources.minerals.tags
+        observed_tags = bot.neutral.minerals.tags
+
+        # BACK-STORY: Mineral id's are fixed when in vision.
+        # Snapshots get random id's every time an object leaves vision.
+        # At game launch when we calculate and save minerals, which are mostly snapshot.
+
+        # Currently, we might have moved vision over minerals, so that their id's have changed.
+        # The alive object share a Position with our cached one, so we can get the correct id and update our cache.
+
+        # PERF: Fix takes 0.70ms, cache takes 0.10ms - we mostly call cached. This is the way.
+        # PERF: In contrast, repeated calls to neutral.minerals.units_in_circle? always costs 0.22ms
+
+        missing_tags = cached_tags - observed_tags
+        unless missing_tags.empty?
+          other_alive_minerals = bot.neutral.minerals.slice(*(observed_tags - cached_tags))
+          # For each missing calculated mineral patch...
+          missing_tags.each do |tag|
+            missing_resource = base_resources.delete(tag)
+
+            # Find an alive mineral at that position
+            new_resource = other_alive_minerals.find { |live_mineral| live_mineral.pos == missing_resource.pos }
+            base_resources.add(new_resource) unless new_resource.nil?
+          end
+        end
+
+        base_resources.minerals
       end
 
       # Gets geysers for a base or base position
       # @param base [Api::Unit, Sc2::Position] base Unit or Position
       # @return [Sc2::UnitGroup] UnitGroup of geysers for the base
       def geysers_for_base(base)
-        resources_for_base(base).geysers
+        # @see #minerals_for_base for backstory on these fixes
+        base_resources = resources_for_base(base)
+        cached_tags = base_resources.geysers.tags
+        observed_tags = bot.neutral.geysers.tags
+
+        missing_tags = cached_tags - observed_tags
+        unless missing_tags.empty?
+          other_alive_geysers = bot.neutral.geysers.slice(*(observed_tags - cached_tags))
+          # For each missing calculated geyser patch...
+          missing_tags.each do |tag|
+            missing_resource = base_resources.delete(tag)
+
+            # Find an alive geyser at that position
+            new_resource = other_alive_geysers.find { |live_geyser| live_geyser.pos == missing_resource.pos }
+            base_resources.add(new_resource) unless new_resource.nil?
+          end
+        end
+
+        base_resources.geysers
+      end
+
+      # Gets geysers which have not been taken for a base or base position
+      # @param base [Api::Unit, Sc2::Position] base Unit or Position
+      # @return [Sc2::UnitGroup] UnitGroup of geysers for the base
+      def geysers_open_for_base(base)
+        geysers = geysers_for_base(base)
+
+        # Mineral-only base, return nothing
+        return UnitGroup.new if geysers.size == 0
+
+        # Reject all which have a gas structure on-top
+        gas_positions = bot.structures.gas.map { |gas| gas.pos }
+        geysers.reject do |geyser|
+          gas_positions.include?(geyser.pos)
+        end
       end
 
       # @private
@@ -554,6 +631,7 @@ module Sc2
       # @return [Sc2::UnitGroup] UnitGroup of resources (minerals+geysers)
       private def resources_for_base(base)
         pos = base.is_a?(Api::Unit) ? base.pos : base
+        pos = pos.to_p2d if base.is_a?(Api::Point)
 
         # If we have a base setup for this exact position, use it
         if expansions.has_key?(pos)
@@ -593,12 +671,17 @@ module Sc2
       def build_coordinates(length:, on_creep: false, in_power: false)
         length = 1 if length < 1
         @_build_coordinates ||= {}
-        cache_key = [length, on_creep].hash
+        cache_key = [length, on_creep, in_power].hash
         return @_build_coordinates[cache_key] if !@_build_coordinates[cache_key].nil? && !bot.game_info_stale?
 
         result = []
         input_grid = parsed_pathing_grid & parsed_placement_grid & ~expo_placement_grid
-        input_grid = parsed_creep & input_grid if on_creep
+        input_grid = if on_creep
+          parsed_creep & input_grid
+        else
+          ~parsed_creep & input_grid
+        end
+
         input_grid = parsed_power_grid & input_grid if in_power
 
         # Dimensions
@@ -651,16 +734,17 @@ module Sc2
       # @param length [Integer] length of the building, 2 for depot/pylon, 3 for rax/gate
       # @param target [Api::Unit, Sc2::Position] near where to find a placement
       # @param random [Integer] number of nearest points to randomly choose from. 1 for nearest point.
+      # @param in_power [Boolean] whether this must be on a power field
       # @return [Api::Point2D, nil] buildable location, nil if no buildable location found
-      def build_placement_near(length:, target:, random: 1)
+      def build_placement_near(length:, target:, random: 1, in_power: false)
         target = target.pos if target.is_a? Api::Unit
         random = 1 if random.to_i.negative?
         length = 1 if length < 1
         on_creep = bot.race == Api::Race::Zerg
 
-        coordinates = build_coordinates(length:, on_creep:)
+        coordinates = build_coordinates(length:, on_creep:, in_power:)
+        cache_key = coordinates.hash
         @_build_coordinate_tree ||= {}
-        cache_key = [length, on_creep].hash
         if @_build_coordinate_tree[cache_key].nil?
           @_build_coordinate_tree[cache_key] = Kdtree.new(
             coordinates.each_with_index.map { |coords, index| coords + [index] }
@@ -813,14 +897,14 @@ module Sc2
       # @example
       #   Randomly randomly adjust both x and y by a range of -3.5 or +3.5
       #   geo.point_random_near(point: structures.hq.first, offset: 3.5)
-      # @param pos [Sc2::Location]
+      # @param pos [Sc2::Position]
       # @param offset [Float]
       # @return [Api::Point2D]
       def point_random_near(pos:, offset: 1.0)
         pos.random_offset(offset)
       end
 
-      # @param pos [Sc2::Location]
+      # @param pos [Sc2::Position]
       # @param radius [Float]
       # @return [Api::Point2D]
       def point_random_on_circle(pos:, radius: 1.0)

data/lib/sc2ai/player/previous_state.rb

@@ -24,7 +24,7 @@ module Sc2
         @status = bot.status
         @game_info = bot.game_info
         @observation = bot.observation
-
+        @game_loop = bot.observation.game_loop
         @spent_minerals = bot.spent_minerals
         @spent_vespene = bot.spent_vespene
         @spent_supply = bot.spent_supply
@@ -36,13 +36,13 @@ module Sc2
       # Override to modify the previous frame before being set to current
       # @param bot [Sc2::Player::Bot]
       def before_reset(bot)
-        # pp "### before_reset"
+        # no op
       end
 
       # Override to modify previous frame after reset is complete
       # @param bot [Sc2::Player::Bot]
       def after_reset(bot)
-        # pp "### after_reset"
+        # no op
       end
     end
   end

data/lib/sc2ai/player/units.rb

@@ -30,6 +30,94 @@ module Sc2
       #   @return [Sc2::UnitGroup] a group of neutral units
       attr_accessor :effects # not a unit
 
+      # Returns the upgrade ids you have acquired such as weapon upgrade and armor upgrade ids.
+      # Shorthand for observation.raw_data.player.upgrade_ids
+      # @!attribute [r] upgrades_completed
+      #   @return [Array<Integer>] a group of neutral units
+      def upgrades_completed = observation&.raw_data&.player&.upgrade_ids.to_a || [] # not a unit
+
+      # Returns the upgrade ids which are researching or queued
+      # Not set for enemy.
+      # @return [Array<Integer>]
+      def upgrades_in_progress
+        # We need to scan every structure which performs upgrades for any order with an upgrade ability
+
+        result = []
+        # Loop every upgrade structure
+        structures
+          .select_type(Api::TechTree.upgrade_structure_unit_type_ids)
+          .each do |structure|
+          next unless structure.is_active? # Skip idle
+
+          # Check if any order at a structure contains an upgrade ability
+          structure.orders.each do |order|
+            Api::TechTree.upgrade_ability_data(structure.unit_type).each do |upgrade_id, update_info|
+              if update_info[:ability] == order.ability_id
+                # Save the upgrade_id
+                result << upgrade_id
+              end
+            end
+          end
+        end
+
+        # If the API told use it's complete, but an order still lingers, trust the API
+        result - upgrades_completed
+      end
+
+      # Returns the upgrade ids which are researching or queued
+      # @return [Boolean]
+      def upgrade_in_progress?(upgrade_id)
+        structure_unit_type_id = Api::TechTree.upgrade_researched_from(upgrade_id: upgrade_id)
+        research_ability_id = Api::TechTree.upgrade_research_ability_id(upgrade_id: upgrade_id)
+        structures.select_type(structure_unit_type_id).each do |structure|
+          structure.orders.each do |order|
+            return true if order.ability_id == research_ability_id
+          end
+        end
+        false
+      end
+
+      # For this unit type, tells you how many are in progress by checking orders for all it's sources.
+      # @return [Integer]
+      def units_in_progress(unit_type_id)
+        source_unit_types = Api::TechTree.unit_created_from(unit_type_id: unit_type_id)
+
+        # When building from LARVA, check the intermediate models
+        if source_unit_types.include?(Api::UnitTypeId::LARVA)
+          source_unit_types << Api::UnitTypeId::EGG
+        elsif source_unit_types.include?(Api::UnitTypeId::BANELING)
+          # For certain Zerg types, return the count of specific intermediate egg/cocoon
+          return units.select_type(Api::UnitTypeId::BANELINGCOCOON).size
+        elsif source_unit_types.include?(Api::UnitTypeId::RAVAGER)
+          return units.select_type(Api::UnitTypeId::RAVAGERCOCOON).size
+        elsif source_unit_types.include?(Api::UnitTypeId::OVERSEER)
+          return units.select_type(Api::UnitTypeId::OVERLORDCOCOON).size
+        elsif source_unit_types.include?(Api::UnitTypeId::LURKERMP)
+          return units.select_type(Api::UnitTypeId::LURKERMPEGG).size
+        elsif source_unit_types.include?(Api::UnitTypeId::BROODLORD)
+          return units.select_type(Api::UnitTypeId::BROODLORDCOCOON).size
+        end
+
+        unit_create_ability = Api::TechTree.unit_type_creation_ability_id(
+          source: source_unit_types.first,
+          target: unit_type_id
+        )
+
+        origin = if unit_data(source_unit_types.first).attributes.include?(:Structure)
+          structures
+        else
+          units
+        end
+        total_in_progress = origin.select_type(source_unit_types).sum do |source|
+          source.orders.count do |order|
+            true if order.ability_id == unit_create_ability
+          end
+        end
+        total_in_progress *= 2 if unit_type_id == Api::UnitTypeId::ZERGLING
+
+        total_in_progress
+      end
+
       # An array of Protoss power sources, which have a point, radius and unit tag
       # @!attribute power_sources
       #   @return [Array<Api::PowerSource>] an array of power sources
@@ -107,6 +195,13 @@ module Sc2
         data.abilities[ability_id]
       end
 
+      # Returns static [Api::UpgradeData] for an upgrade id
+      # @param upgrade_id [Integer] Api::UpgradeId::*
+      # @return [Api::UpgradeData]
+      def upgrade_data(upgrade_id)
+        data.upgrades[upgrade_id]
+      end
+
       # Checks unit data for an attribute value
       # @param unit [Integer,Api::Unit] Api::UnitTypeId or Api::Unit
       # @param attribute [Symbol] Api::Attribute, i.e. Api::Attribute::Mechanical or :Mechanical
@@ -141,17 +236,13 @@ module Sc2
       def subtract_cost(unit_type_id)
         unit_type_data = unit_data(unit_type_id)
 
-        # food_required is a float. ensure half units are counted as full
-        # TODO: Extend UnitTypeData message. def food_required = unit_id ==  Api::UnitTypeId::ZERGLING ? 1 : send("method_missing", :food_required)
-        supply_cost = unit_type_data.food_required
-        supply_cost = 1 if unit_type_id == Api::UnitTypeId::ZERGLING
-
         @spent_minerals += unit_type_data.mineral_cost
         @spent_vespene += unit_type_data.vespene_cost
-        @spent_supply += supply_cost
+        @spent_supply += unit_type_data.food_required
       end
 
       # Checks whether you have the resources to construct quantity of unit type
+      # @return [Boolean]
       def can_afford?(unit_type_id:, quantity: 1)
         unit_type_data = unit_data(unit_type_id)
         return false if unit_type_data.nil?
@@ -178,6 +269,25 @@ module Sc2
         true
       end
 
+      # Checks whether you have the resources to
+      # @return [Boolean]
+      def can_afford_upgrade?(upgrade_id)
+        unit_type_data = upgrade_data(upgrade_id)
+        return false if unit_type_data.nil?
+
+        mineral_cost = unit_type_data.mineral_cost
+        if common.minerals - spent_minerals < mineral_cost
+          return false # not enough minerals
+        end
+
+        vespene_cost = unit_type_data.vespene_cost
+        if common.vespene - spent_vespene < vespene_cost
+          return false # you require more vespene gas
+        end
+
+        true
+      end
+
       private
 
       # @private

data/lib/sc2ai/player.rb

@@ -226,13 +226,14 @@ module Sc2
         # Callback for step 0
         on_step
 
+        puts ""
         # Step 1 to n
         loop do
           r = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
           perform_actions
           perform_debug_commands # TODO: Detect IS_LADDER? -> unless IS_LADDER?
           step_forward
-          puts (::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - r) * 1000
+          print "\e[2K#{@step_count} Steps Took (ms): #{(::Process.clock_gettime(::Process::CLOCK_MONOTONIC) - r) * 1000}\n\e[1A\r"
           return @result unless @result.nil?
           break if @status != :in_game
         end
@@ -525,6 +526,7 @@ module Sc2
       @previous.reset(self)
       # Reset
       self.observation = response_observation.observation
+      self.game_loop = observation.game_loop
       self.chats_received = response_observation.chat
       self.spent_minerals = 0
       self.spent_vespene = 0

data/lib/sc2ai/protocol/_meta_documentation.rb

@@ -21,6 +21,8 @@
 #     class Point < Google::Protobuf::AbstractMessage; end;
 #     # Protobuf virtual class.
 #     class Unit < Google::Protobuf::AbstractMessage; end;
+#     # Protobuf virtual class.
+#     class UnitTypeData < Google::Protobuf::AbstractMessage; end;
 #   end
 
 # Protobuf enums ---

data/lib/sc2ai/protocol/extensions/ability_remapable.rb

@@ -0,0 +1,22 @@
+module Api
+  # This module make sure that a read from method ability_id always returns the proper source id
+  module AbilityRemapable
+    # Ability Id. The generic id or "remapid".
+    # i.e. Api::AbilityId::ATTACK_BATTLECRUISER returns generic Api::AbilityId::ATTACK
+    # @return [Integer]
+    def ability_id
+      @ability_id ||= Api::AbilityId.generic_id(send(:method_missing, :ability_id))
+    end
+  end
+end
+
+# AbilityData should not include, since it holds exact info and contains the remap id as attr
+# Similarly Request* methods only do requests and ew supply correct id's.
+Api::AvailableAbility.include Api::AbilityRemapable
+Api::UnitOrder.include Api::AbilityRemapable
+Api::ActionRawUnitCommand.include Api::AbilityRemapable
+Api::ActionRawToggleAutocast.include Api::AbilityRemapable
+Api::ActionError.include Api::AbilityRemapable
+Api::ActionSpatialUnitCommand.include Api::AbilityRemapable
+Api::BuildItem.include Api::AbilityRemapable
+Api::ActionToggleAutocast.include Api::AbilityRemapable

data/lib/sc2ai/protocol/extensions/point.rb

@@ -3,7 +3,9 @@ module Api
   module PointExtension
     # @private
     def hash
-      [x, y, z].hash
+      # Only one plane is ever used. Ignore hashing on z
+      # [x, y, z].hash
+      [x, y].hash
     end
 
     def eql?(other)

data/lib/sc2ai/protocol/extensions/point_2_d.rb

@@ -10,6 +10,12 @@ module Api
       self.class == other.class && hash == other.hash
     end
 
+    # Create a new 3d Point, by adding a y axis.
+    # @return [Api::Point]
+    def to_3d(z:)
+      Api::Point[x, y, z]
+    end
+
     # Adds additional functionality to message class Api::Point2D
     module ClassMethods
       # Shorthand for creating an instance for [x, y]

data/lib/sc2ai/protocol/extensions/position.rb

@@ -1,12 +1,23 @@
 module Sc2
   # A unified construct that tames Api::* messages which contain location data
-  # Items which are of type Sc2::Location will have #x and #y property at the least.
+  # Items which are of type Sc2::Position will have #x and #y property at the least.
   module Position
     # Tolerance for floating-point comparisons.
     TOLERANCE = 1e-9
 
     # Basic operations
 
+    # Loose equality matches on floats x and y.
+    # We never check z-axis, because the map is single-level.
+    # TODO: We should almost certainly introduce TOLERANCE here, but verify it's cost first.
+    def ==(other)
+      if other.is_a? Position
+        x == other.x && y == other.y
+      else
+        false
+      end
+    end
+
     # A new point representing the sum of this point and the other point.
     # @param other [Api::Point2D, Numeric] The other point/number to add.
     # @return [Api::Point2D]
@@ -50,12 +61,38 @@ module Sc2
     # @see #divide
     alias_method :/, :divide
 
-    # Bug: Psych implements method 'y' on Kernel, but protobuf uses method_missing to read AbstractMethod
-    # We send method missing ourselves when y to fix this chain.
+    # Returns x coordinate
+    # @return [Float]
+    def x
+      # Perf: Memoizing attributes which are hit hard, show gain
+      @x ||= send(:method_missing, :x)
+    end
+
+    # Sets x coordinate
+    # @return [Float]
+    def x=(x)
+      send(:method_missing, :x=, x)
+      @x = x
+    end
+
+    # Returns y coordinate
+    # @return [Float]
     def y
+      # Bug: Psych implements method 'y' on Kernel, but protobuf uses method_missing to read AbstractMethod
+      # We send method missing ourselves when y to fix this chain.
       # This is correct, but an unnecessary conditional:
       # raise NoMethodError unless location == self
-      send(:method_missing, :y)
+
+      # Perf: Memoizing attributes which are hit hard, show gain
+
+      @y ||= send(:method_missing, :y)
+    end
+
+    # Sets y coordinate
+    # @return [Float]
+    def y=(y)
+      send(:method_missing, :y=, y)
+      @y = y
     end
 
     # Randomly adjusts both x and y by a range of: -offset..offset

data/lib/sc2ai/protocol/extensions/unit.rb

@@ -9,6 +9,22 @@ module Api
       tag || super
     end
 
+    # Returns an integer unique identifier
+    # If the unit goes out of vision and is snapshot-able, they get a random id
+    # - Such a unit gets the same unit tag when it re-enters vision
+    # @return [Integer]
+    def tag
+      # Perf: This speeds up hash and therefore common UnitGroup operations. Sometimes 3x!
+      @tag ||= send(:method_missing, :tag)
+    end
+
+    # Sets unit tag
+    # @return [Integer]
+    def tag=(tag)
+      send(:method_missing, :tag=, tag)
+      @tag = tag
+    end
+
     # Every unit gets access back to the bot to allow api access.
     # For your own units, this allows API access.
     # @return [Sc2::Player] player with active connection
@@ -37,9 +53,8 @@ module Api
     # Checks unit data for an attribute value
     # @return [Boolean] whether unit has attribute
     # @example
-    #   has_attribute?(Api::UnitTypeId::SCV, Api::Attribute::Mechanical)
-    #   has_attribute?(units.workers.first, :Mechanical)
-    #   has_attribute?(Api::UnitTypeId::SCV, :Mechanical)
+    #   unit.has_attribute?(Api::Attribute::Mechanical)
+    #   unit.has_attribute?(:Mechanical)
     def has_attribute?(attribute)
       attributes.include? attribute
     end
@@ -162,6 +177,15 @@ module Api
 
     # @!endgroup Virtual properties
 
+    # Whether unit is effected by buff_id
+    # @example
+    #   unit.has_buff??(Api::BuffId::QUEENSPAWNLARVATIMER)
+    # @param [Integer] buff_id
+    # @return [Boolean]
+    def has_buff?(buff_id)
+      buff_ids.include?(buff_id)
+    end
+
     # @!group Actions
 
     # Performs action on this unit
@@ -233,10 +257,18 @@ module Api
 
     # Issues repair command on target
     # @param target [Api::Unit, Integer] is a unit or unit tag
+    # @param queue_command [Boolean] shift+command
     def repair(target:, queue_command: false)
       action(ability_id: Api::AbilityId::EFFECT_REPAIR, target:, queue_command:)
     end
 
+    # Research a specific upgrade
+    # @param upgrade_id [Integer] Api::UnitTypeId the unit type which will do the creation
+    # @param queue_command [Boolean] shift+command
+    def research(upgrade_id:, queue_command: false)
+      @bot.research(units: self, upgrade_id:, queue_command:)
+    end
+
     # @!endgroup Actions
     #
     # Debug ----

data/lib/sc2ai/protocol/extensions/unit_type_data.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Api
+  # Adds additional functionality to message object Api::UnitTypeData
+  module UnitTypeDataExtension
+    # @!attribute mineral_cost_sum
+    #   Sum of all morphs mineral cost
+    #   i.e. 550M Orbital command = 400M CC + 150M Upgrade
+    #   i.e. 350M Hatchery = 50M Drone + 300M Build
+    #   @return [Integer] sum of mineral costs
+    attr_accessor :mineral_cost_sum
+
+    # @!attribute vespene_cost_sum
+    #   Sum of all morphs vespene gas cost
+    #   i.e. 250G Broodlord = 100G Corruptor + 150G Morph
+    #   @return [Integer] sum of vespene gas costs
+    attr_accessor :vespene_cost_sum
+  end
+end
+Api::UnitTypeData.include Api::UnitTypeDataExtension