z_mirror/Krastorio2

git clone https://git.raiguard.me/z_mirror/Krastorio2

Caleb HeuerFix crashes related to minable/active entity flag changes1e9ba1882026-06-25

local flib_bounding_box = require("__flib__.bounding-box")

--- @alias k2.ResourceRestrictions.Forces table<ForceIndex, k2.ResourceRestrictions.ForceData?>

--- @class k2.ResourceRestrictions.ForceData
--- @field disabled_mining_drills table<UnitNumber, k2.ResourceRestrictions.DrillData>
--- @field force LuaForce
--- @field restricted_resources Set<string>

--- @class k2.ResourceRestrictions.DrillData
--- @field drill LuaEntity
--- @field icon LuaRenderObject

--- @return k2.ResourceRestrictions.Forces
local function get_forces()
  if not storage.resource_restrictions then
    --- @type k2.ResourceRestrictions.Forces
    storage.resource_restrictions = {}
  end
  return storage.resource_restrictions
end

--- @param force LuaForce
local function get_force_data(force)
  local forces = get_forces()
  local force_index = force.index
  local force_data = forces[force_index]
  if not force_data then
    force_data = {
      disabled_mining_drills = {},
      force = force,
      restricted_resources = {},
    }
    forces[force_index] = force_data
  end
  return force_data
end

--- @param force_data k2.ResourceRestrictions.ForceData
--- @param drill LuaEntity
--- @return LuaEntityPrototype? restricted_resource
local function get_first_restricted_resource(force_data, drill)
  local radius = drill.prototype.get_mining_drill_radius(drill.quality)
  --- @cast radius -?
  local area = flib_bounding_box.from_dimensions(drill.position, radius * 2, radius * 2)
  local targets = drill.surface.find_entities_filtered({ type = "resource", area = area })
  local drill_categories = drill.prototype.resource_categories or {}
  local restrictions = force_data.restricted_resources
  for i = 1, #targets do
    local target = targets[i]
    if drill_categories[target.prototype.resource_category] and restrictions[target.name] then
      return target.prototype
    end
  end
end

--- @param force_data k2.ResourceRestrictions.ForceData
--- @param drill LuaEntity
local function update_drill(force_data, drill)
  local disabled_mining_drills = force_data.disabled_mining_drills

  local unit_number = drill.unit_number
  --- @cast unit_number -?

  local drill_data = disabled_mining_drills[unit_number]

  local restriction = get_first_restricted_resource(force_data, drill)
  if restriction then
    if not drill_data then
      drill.disabled_by_script = true
      drill.custom_status = {
        diode = defines.entity_status_diode.red,
        label = { "entity-status.kr-unable-to-mine", restriction.localised_name },
      }
      disabled_mining_drills[unit_number] = {
        drill = drill,
        icon = rendering.draw_sprite({
          sprite = "utility/resources_depleted_icon",
          target = drill,
          surface = drill.surface,
          x_scale = 0.5,
          y_scale = 0.5,
          blink_interval = 30,
        }),
      }
    end
  else
    if drill_data then
      drill.disabled_by_script = false
      drill.custom_status = nil
      if drill_data.icon.valid then
        drill_data.icon.destroy()
      end
      disabled_mining_drills[unit_number] = nil
    end
  end
end

--- @param force_data k2.ResourceRestrictions.ForceData
local function refresh_all_drills(force_data)
  for _, drill_data in pairs(force_data.disabled_mining_drills) do
    if drill_data.icon.valid then
      drill_data.icon.destroy()
    end
  end
  force_data.disabled_mining_drills = {}
  for _, surface in pairs(game.surfaces) do
    for _, drill in pairs(surface.find_entities_filtered({ type = "mining-drill", force = force_data.force })) do
      update_drill(force_data, drill)
    end
  end
end

local unlocks = prototypes.mod_data["kr-resource-restrictions"].data --[[ @as table<data.TechnologyID, data.EntityID[]?> ]]

--- @param force_data k2.ResourceRestrictions.ForceData
--- @param technology LuaTechnology
--- @param restrict boolean
local function process_effects(force_data, technology, restrict)
  local technology_unlocks = unlocks[technology.name]
  if not technology_unlocks then
    return
  end

  local restricted_resources = force_data.restricted_resources
  for _, resource_name in pairs(technology_unlocks) do
    restricted_resources[resource_name] = restrict or nil
  end

  if restrict then
    -- The on_research_reversed event is very rare, and never happens during regular gameplay, so it is acceptable to have a lag spike in this case.
    refresh_all_drills(force_data)
  else
    for _, drill_data in pairs(force_data.disabled_mining_drills) do
      update_drill(force_data, drill_data.drill)
    end
  end
end

--- @param force LuaForce
local function refresh_force(force)
  force.mining_with_fluid = true
  local force_data = get_force_data(force)
  for _, technology in pairs(force.technologies) do
    if not technology.researched then
      -- TODO: This is doing a full scan of all surfaces every time!
      process_effects(force_data, technology, true)
    end
  end
end

--- @param e EntityBuiltEvent
local function on_entity_built(e)
  local entity = e.entity or e.destination
  if not entity.valid or entity.type ~= "mining-drill" then
    return
  end
  update_drill(get_force_data(entity.force --[[@as LuaForce]]), entity)
end

--- @param e EntityDestroyedEvent
local function on_entity_destroyed(e)
  local entity = e.entity
  if not entity.valid or entity.type ~= "mining-drill" then
    return
  end

  -- The icon will get destroyed automatically because its target is the mining drill.
  get_force_data(e.entity.force --[[@as LuaForce]]).disabled_mining_drills[entity.unit_number] = nil
end

--- @param e EventData.on_research_finished
local function on_research_finished(e)
  process_effects(get_force_data(e.research.force), e.research, false)
end

--- @param e EventData.on_research_finished
local function on_research_reversed(e)
  process_effects(get_force_data(e.research.force), e.research, true)
end

--- @param e EventData.on_force_created
local function on_force_created(e)
  refresh_force(e.force)
end

--- @param e EventData.on_force_reset
local function on_force_reset(e)
  refresh_force(e.force)
end

--- @param e EventData.on_forces_merging
local function on_forces_merging(e)
  local source_data = get_force_data(e.source)
  for _, drill_data in pairs(source_data.disabled_mining_drills) do
    if drill_data.icon.valid then
      drill_data.icon.destroy()
    end
  end
  get_forces()[e.source.index] = nil
  refresh_force(e.destination)
end

--- @param e EventData.on_technology_effects_reset
local function on_technology_effects_reset(e)
  refresh_force(e.force)
end

local function refresh()
  storage.resource_restrictions = nil
  for _, force in pairs(game.forces) do
    refresh_force(force)
    refresh_all_drills(get_force_data(force))
  end
end

--- @class MiningDrillRestrictionsHandler: event_handler
local handler = {}

handler.on_init = refresh
handler.on_configuration_changed = refresh

handler.events = {
  [defines.events.on_built_entity] = on_entity_built,
  [defines.events.on_entity_cloned] = on_entity_built,
  [defines.events.on_entity_died] = on_entity_destroyed,
  [defines.events.on_force_created] = on_force_created,
  [defines.events.on_force_reset] = on_force_reset,
  [defines.events.on_forces_merging] = on_forces_merging,
  [defines.events.on_player_mined_entity] = on_entity_destroyed,
  [defines.events.on_research_finished] = on_research_finished,
  [defines.events.on_research_reversed] = on_research_reversed,
  [defines.events.on_robot_built_entity] = on_entity_built,
  [defines.events.on_robot_mined_entity] = on_entity_destroyed,
  [defines.events.on_space_platform_built_entity] = on_entity_built,
  [defines.events.on_space_platform_mined_entity] = on_entity_destroyed,
  [defines.events.on_technology_effects_reset] = on_technology_effects_reset,
  [defines.events.script_raised_built] = on_entity_built,
  [defines.events.script_raised_destroy] = on_entity_destroyed,
  [defines.events.script_raised_revive] = on_entity_built,
}

return handler