z_mirror/Krastorio2

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

Caleb HeuerChange to four-space indentation07409ca62026-01-05

local flib_position = require("__flib__.position")
local flib_table = require("__flib__.table")

local data_util = {}

-- To enable verbose mode, either hook the debug adapter or add a DEBUG.lua file to the mod root.
local verbose = __DebugAdapter ~= nil or pcall(require, "DEBUG")

--- Prints the given error message to the log when in debug mode.
--- @param message string
function data_util.error(message)
    if not verbose then
        return
    end
    local output = "\n      KRASTORIO 2 ERROR: " .. message
    local i = 2
    while true do
        local info = debug.getinfo(i, "nSl")
        if not info then
            break
        end

        output = output .. "\n        " .. info.source .. ":" .. info.currentline
        if info.name then
            output = output .. " in function '" .. info.name .. "'"
        else
            output = output .. " in main chunk"
        end

        i = i + 1
    end
    log(output)
end

--- Calls `data_util.error` with the given `message` if `obj` is falsey.
--- @param obj any?
--- @param message string
--- @return boolean passed "Whether the assert passed."
function data_util.assert(obj, message)
    if not obj then
        data_util.error(message)
        return false
    end
    return true
end

--- Overwrites or adds ammo effects for the given technology.
--- @param tech data.TechnologyPrototype
--- @param effects data.Modifier[]
function data_util.update_ammo_effects(tech, effects)
    if not tech.effects then
        tech.effects = {}
    end

    for _, updated_effect in pairs(effects) do
        local exists = false
        for _, effect in pairs(tech.effects) do
            if
                effect.type == updated_effect.type
                and (
                    (effect.ammo_category ~= nil and effect.ammo_category == updated_effect.ammo_category)
                    or (effect.turret_id ~= nil and effect.turret_id == updated_effect.turret_id)
                )
            then
                effect.modifier = updated_effect.modifier --- @diagnostic disable-line:assign-type-mismatch
                exists = true
                break
            end
        end
        if not exists then
            table.insert(
                tech.effects,
                {
                    type = updated_effect.type,
                    ammo_category = updated_effect.ammo_category,
                    modifier = updated_effect.modifier,
                } --- @diagnostic disable-line:assign-type-mismatch
            )
        end
    end
end

--- @param technology_name data.TechnologyID
--- @param new_effect data.Modifier
function data_util.add_effect(technology_name, new_effect)
    local technology = data.raw.technology[technology_name]
    if not technology then
        data_util.error("Technology " .. technology_name .. " does not exist.")
        return
    end
    local effects = technology.effects
    if not effects or not next(effects) then
        technology.effects = { new_effect }
        return
    end

    local found = false
    for _, effect in pairs(effects) do
        if effect.type == new_effect.type then
            -- TODO: Update for new effect types.
            if effect.type == "gun-speed" then
                if effect.ammo_category == new_effect.ammo_category and effect.modifier == new_effect.modifier then
                    found = true
                    break
                end
            elseif effect.type == "ammo-damage" then
                if effect.ammo_category == new_effect.ammo_category and effect.modifier == new_effect.modifier then
                    found = true
                    break
                end
            elseif effect.type == "give-item" then
                if effect.item == new_effect.item and effect.count == new_effect.count then
                    found = true
                    break
                end
            elseif effect.type == "turret-attack" then
                if effect.turret_id == new_effect.turret_id and effect.modifier == new_effect.modifier then
                    found = true
                    break
                end
            elseif effect.type == "unlock-recipe" then
                if effect.recipe == new_effect.recipe then
                    found = true
                    break
                end
            elseif effect.type == "nothing" then
                if effect.effect_description == new_effect.effect_description then
                    found = true
                    break
                end
            else
                if effect.modifier == new_effect.modifier then
                    found = true
                    break
                end
            end
        end
    end

    if not found then
        table.insert(effects, new_effect)
    end
end

--- Adds the given recipe as an unlock of the given technology.
--- @param technology_name data.TechnologyID
--- @param recipe_name data.RecipeID
function data_util.add_recipe_unlock(technology_name, recipe_name)
    local technology = data.raw.technology[technology_name]
    if not technology then
        data_util.error("Technology " .. technology_name .. " does not exist.")
        return
    end
    for _, effect in pairs(technology.effects) do
        if effect.type == "unlock-recipe" and effect.recipe == recipe_name then
            data_util.error("Technology " .. technology_name .. " already unlocks recipe " .. recipe_name .. ".")
            return
        end
    end
    table.insert(technology.effects, { type = "unlock-recipe", recipe = recipe_name })
end

--- Removes the given recipe from the unlocks of the given technology.
--- @param technology_name data.TechnologyID
--- @param recipe_name data.RecipeID
function data_util.remove_recipe_unlock(technology_name, recipe_name)
    local technology = data.raw.technology[technology_name]
    if not technology then
        data_util.error("Technology " .. technology_name .. " does not exist.")
        return
    end
    for i, effect in pairs(technology.effects) do
        if effect.type == "unlock-recipe" and effect.recipe == recipe_name then
            table.remove(technology.effects, i)
            return
        end
    end
    data_util.error("Technology " .. technology_name .. " does not unlock recipe " .. recipe_name .. ".")
end

--- Adds the given prerequisite to the technology.
--- @param technology_id data.TechnologyID
--- @param prerequisite_id data.TechnologyID
function data_util.add_prerequisite(technology_id, prerequisite_id)
    local technology = data.raw.technology[technology_id]
    if not technology then
        data_util.error("Technology " .. technology_id .. " does not exist.")
        return
    end
    local prerequisite = data.raw.technology[prerequisite_id]
    if not prerequisite then
        data_util.error("Technology prerequisite " .. prerequisite_id .. " does not exist.")
        return
    end
    if not technology.prerequisites then
        technology.prerequisites = { prerequisite_id }
        return
    end
    if flib_table.find(technology.prerequisites, prerequisite_id) then
        data_util.error("Technology " .. technology_id .. " already has prerequisite " .. prerequisite_id)
        return
    end
    table.insert(technology.prerequisites, prerequisite_id)
end

--- Removes the given prerequisite from the technology.
--- @param technology_id data.TechnologyID
--- @param prerequisite_id data.TechnologyID
function data_util.remove_prerequisite(technology_id, prerequisite_id)
    local technology = data.raw.technology[technology_id]
    if not technology then
        data_util.error("Technology " .. technology_id .. " does not exist.")
        return
    end
    local prerequisite = data.raw.technology[prerequisite_id]
    if not prerequisite then
        data_util.error("Technology prerequisite " .. prerequisite_id .. " does not exist.")
        return
    end
    local i = flib_table.find(technology.prerequisites or {}, prerequisite_id)
    if not i then
        data_util.error("Technology " .. technology_id .. " does not have prerequisite " .. prerequisite_id)
        return
    end
    table.remove(technology.prerequisites, i)
end

--- Converts one research unit ingredient to another, updating science pack prerequisites if needed.
--- @param technology_id data.TechnologyID
--- @param from_id data.ItemID
--- @param to_id data.ItemID
function data_util.convert_research_unit_ingredient(technology_id, from_id, to_id)
    local technology = data.raw.technology[technology_id]
    if not technology then
        data_util.error("Technology " .. technology_id .. " does not exist.")
        return
    end
    if not technology.unit or not technology.unit.ingredients then
        data_util.error("Technology " .. technology_id .. " has no research unit ingredients.")
        return
    end
    local converted = false
    for _, ingredient in pairs(technology.unit.ingredients) do
        if ingredient[1] == from_id then
            converted = true
            ingredient[1] = to_id
        elseif ingredient[1] == to_id then
            data_util.error("Technology " .. technology_id .. " already has research unit ingredient " .. to_id)
            return
        end
    end
    if not converted then
        data_util.error("Technology " .. technology_id .. " does not have research unit ingredient " .. from_id)
        return
    end
    -- XXX: Hardcoding the science pack name like this feels bad, but it's what the old code did...
    local from_index = flib_table.find(technology.prerequisites, from_id)
    if from_index then
        table.remove(technology.prerequisites, from_index)
    end
    local new_tech = data.raw.technology[to_id]
    if not new_tech then
        return
    end
    if not technology.prerequisites then
        technology.prerequisites = { to_id }
        return
    end
    if not flib_table.find(technology.prerequisites, to_id) then
        table.insert(technology.prerequisites, to_id)
    end
end

--- Adds the given research ingredient to the technology.
--- @param technology_id data.TechnologyID
--- @param ingredient_id data.ItemID
--- @param ingredient_count uint? Defaults to 1.
function data_util.add_research_unit_ingredient(technology_id, ingredient_id, ingredient_count)
    ingredient_count = ingredient_count or 1
    local technology = data.raw.technology[technology_id]
    if not technology then
        data_util.error("Technology " .. technology_id .. " does not exist.")
        return
    end
    if technology.research_trigger then
        data_util.error("Technology " .. technology_id .. " has a research trigger, so cannot have ingredients.")
        return
    end
    local unit = technology.unit
    if not unit then
        data_util.error("Technology " .. technology_id .. " has no research unit.")
        return
    end
    for _, ingredient in pairs(technology.unit.ingredients) do
        if ingredient[1] == ingredient_id then
            data_util.error("Technology " .. technology_id .. " already has research unit ingredient" .. ingredient_id)
            return
        end
    end
    table.insert(technology.unit.ingredients, { ingredient_id, ingredient_count })
end

--- Removes the given research ingredient from the technology.
--- @param technology_id data.TechnologyID
--- @param ingredient_id data.ItemID
function data_util.remove_research_unit_ingredient(technology_id, ingredient_id)
    local technology = data.raw.technology[technology_id]
    if not technology then
        data_util.error("Technology " .. technology_id .. " does not exist.")
        return
    end
    if not technology.unit then
        data_util.error("Technology " .. technology_id .. " has no research units.")
        return
    end
    for i, ingredient in pairs(technology.unit.ingredients) do
        if ingredient[1] == ingredient_id then
            table.remove(technology.unit.ingredients, i)
            return
        end
    end
    data_util.error(
        "Technology " .. technology_id .. " does not have research unit ingredient " .. ingredient_id .. "."
    )
end

--- Returns a technology that unlocks the given recipe. The technology returned will be the one that was added earliest.
--- @param recipe_id data.RecipeID
--- @return data.TechnologyPrototype?
function data_util.get_technology_for_recipe(recipe_id)
    for _, technology in pairs(data.raw.technology) do
        for _, effect in pairs(technology.effects or {}) do
            if effect.type == "unlock-recipe" and effect.recipe == recipe_id then
                return technology
            end
        end
    end
end

--- Converts a furnace prototype into an assembling machine prototype.
--- @param furnace_name data.EntityID
--- @return data.AssemblingMachinePrototype?
function data_util.furnace_to_assembler(furnace_name)
    local furnace = data.raw.furnace[furnace_name]
    if not furnace then
        data_util.error("Furnace " .. furnace_name .. " does not exist.")
        return
    end

    local assembler = table.deepcopy(furnace) --[[@as data.AssemblingMachinePrototype]]
    assembler.type = "assembling-machine"
    assembler.source_inventory_size = nil --- @diagnostic disable-line
    assembler.energy_source.emissions_per_minute = { pollution = 2 }
    assembler.energy_usage = "0.2MW"
    data.raw.furnace[furnace_name] = nil
    data:extend({ assembler })
    return assembler
end

--- Returns a copy of the given prototype's icons in the standard format.
--- @param prototype k2.PrototypeWithIcons
--- @return data.IconData[]
function data_util.get_icons(prototype)
    local icons = prototype.icons -- LuaLS isn't inferring the dropped optional unless we assign it to a local first.
    if icons then
        return table.deepcopy(icons)
    end
    if prototype.icon then
        return { { icon = prototype.icon, icon_size = prototype.icon_size } }
    end
    return {}
end

--- Converts the given prototype's icons to the standard format and returns a reference to them.
--- The game will throw an error if you provide an empty icons array, so if using the function, care must be taken to
--- ensure that at least one icon layer exists at tne end of loading.
--- @param prototype k2.PrototypeWithIcons
--- @return data.IconData[]
function data_util.standardize_icons(prototype)
    if not prototype.icons then
        prototype.icons = {}
        if prototype.icon then
            table.insert(prototype.icons, { icon = prototype.icon, icon_size = prototype.icon_size })
            prototype.icon = nil
            prototype.icon_size = nil
        end
    end
    return prototype.icons
end

--- Overrides the prototype's icon.
--- @param prototype k2.PrototypeWithIcons
--- @param icon data.FileName
--- @param icon_size double?
function data_util.set_icon(prototype, icon, icon_size)
    data_util.assert(prototype, "Provided prototype was nil.")
    if prototype.icons then
        prototype.icons = nil
    end
    prototype.icon = icon
    prototype.icon_size = icon_size
end

--- Overrides the prototype's icons array.
--- @param prototype k2.PrototypeWithIcons
--- @param icons data.IconData[]
function data_util.set_icons(prototype, icons)
    data_util.assert(prototype, "Provided prototype was nil.")
    prototype.icons = icons
    prototype.icon = nil
    prototype.icon_size = nil
end

--- Fields are analagous to `data.IconData`.
--- @class k2.IconTransformOptions
--- @field shift Vector?
--- @field scale double?

--- Applies the given transformations to each layer in the icons array in-place and returns the array.
--- @param icons data.IconData[]
--- @param transforms k2.IconTransformOptions
--- @return data.IconData[]
function data_util.transform_icons(icons, transforms)
    for _, icon in pairs(icons) do
        if transforms.shift then
            icon.shift = flib_position.add((icon.shift or { 0, 0 }), transforms.shift)
        end
        if transforms.scale then
            icon.scale = (icon.scale or 1) * transforms.scale / (icon.icon_size or 64) * 64
        end
    end
    return icons
end

--- Merges or adds the given ingredient to the recipe matching the given name.
--- @param recipe_name data.RecipeID
--- @param new_ingredient data.IngredientPrototype
function data_util.merge_or_add_ingredient(recipe_name, new_ingredient)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist")
        return
    end
    recipe.ingredients = recipe.ingredients or {}
    for _, ingredient in pairs(recipe.ingredients) do
        if ingredient.name == new_ingredient.name then
            ingredient.amount = ingredient.amount + new_ingredient.amount
            return
        end
    end
    table.insert(recipe.ingredients, new_ingredient)
end

--- Overwrites or adds the given ingredient to the recipe matching the given name.
--- @param recipe_name data.RecipeID
--- @param new_ingredient data.IngredientPrototype
function data_util.overwrite_or_add_ingredient(recipe_name, new_ingredient)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist")
        return
    end
    recipe.ingredients = recipe.ingredients or {}
    for i, ingredient in pairs(recipe.ingredients) do
        if ingredient.name == new_ingredient.name then
            recipe.ingredients[i] = new_ingredient
            return
        end
    end
    table.insert(recipe.ingredients, new_ingredient)
end

--- Removes old_ingredient_name from the recipe matching the given name, then overwrites any ingredient with the
--- same name as the new ingredient with the given new_ingredient.
--- @param recipe_name data.RecipeID
--- @param old_ingredient_name data.FluidID|data.ItemID
--- @param new_ingredient data.IngredientPrototype
function data_util.replace_and_overwrite_or_add_ingredient(recipe_name, old_ingredient_name, new_ingredient)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist.")
        return
    end
    recipe.ingredients = recipe.ingredients or {}
    data_util.remove_ingredient(recipe_name, old_ingredient_name)
    data_util.overwrite_or_add_ingredient(recipe_name, new_ingredient)
end

--- Removes old_ingredient_name from the recipe matching the given name, then merges or adds the new ingredient.
--- @param recipe_name data.RecipeID
--- @param old_ingredient_name data.FluidID|data.ItemID
--- @param new_ingredient data.IngredientPrototype
function data_util.replace_and_merge_or_add_ingredient(recipe_name, old_ingredient_name, new_ingredient)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist")
        return
    end
    recipe.ingredients = recipe.ingredients or {}
    data_util.remove_ingredient(recipe_name, old_ingredient_name)
    data_util.merge_or_add_ingredient(recipe_name, new_ingredient)
end

--- Converts the given ingredient to the new ingredient, carrying over other properties.
--- @param recipe_name data.RecipeID
--- @param old_ingredient_name data.FluidID|data.ItemID
--- @param new_ingredient_name data.FluidID|data.ItemID
function data_util.convert_ingredient(recipe_name, old_ingredient_name, new_ingredient_name)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist.")
        return
    end
    if not recipe.ingredients then
        data_util.error("Recipe " .. recipe_name .. " has no ingredients.")
        return
    end
    for _, ingredient in pairs(recipe.ingredients) do
        if ingredient.name == old_ingredient_name then
            ingredient.name = new_ingredient_name
            break
        end
    end
end

--- Remove the given ingredient from the recipe.
--- @param recipe_name data.RecipeID
--- @param ingredient_name data.FluidID|data.ItemID
function data_util.remove_ingredient(recipe_name, ingredient_name)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist.")
        return
    end
    if not recipe.ingredients then
        data_util.error("Recipe " .. recipe_name .. " has no ingredients.")
        return
    end
    for i, ingredient in pairs(recipe.ingredients) do
        if ingredient.name == ingredient_name then
            table.remove(recipe.ingredients, i)
            return
        end
    end
    data_util.error("Recipe " .. recipe_name .. " does not have ingredient " .. ingredient_name .. ".")
end

--- Merges or adds the given prodct to the recipe matching the given name.
--- @param recipe_name data.RecipeID
--- @param new_product data.ProductPrototype
function data_util.merge_or_add_product(recipe_name, new_product)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist")
        return
    end
    recipe.results = recipe.results or {}
    for _, product in pairs(recipe.results) do
        if product.name == new_product.name then
            product.amount = product.amount + new_product.amount
            return
        end
    end
    table.insert(recipe.results, new_product)
end

--- Overwrites or adds the given product to the recipe matching the given name.
--- @param recipe_name data.RecipeID
--- @param new_product data.ProductPrototype
function data_util.overwrite_or_add_product(recipe_name, new_product)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist")
        return
    end
    recipe.results = recipe.results or {}
    for i, product in pairs(recipe.results) do
        if product.name == new_product.name then
            recipe.results[i] = new_product
            return
        end
    end
    table.insert(recipe.results, new_product)
end

--- Removes old_product_name from the recipe matching the given name, then overwrites any product with the
--- same name as the new product with the given new_product.
--- @param recipe_name data.RecipeID
--- @param old_product_name data.FluidID|data.ItemID
--- @param new_product data.ProductPrototype
function data_util.replace_and_overwrite_or_add_product(recipe_name, old_product_name, new_product)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist.")
        return
    end
    recipe.results = recipe.results or {}
    data_util.remove_product(recipe_name, old_product_name)
    data_util.overwrite_or_add_product(recipe_name, new_product)
end

--- Removes old_product_name from the recipe matching the given name, then merges or adds the new product.
--- @param recipe_name data.RecipeID
--- @param old_product_name data.FluidID|data.ItemID
--- @param new_product data.ProductPrototype
function data_util.replace_and_merge_or_add_product(recipe_name, old_product_name, new_product)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist")
        return
    end
    recipe.results = recipe.results or {}
    data_util.remove_product(recipe_name, old_product_name)
    data_util.merge_or_add_product(recipe_name, new_product)
end

--- Remove the given product from the recipe.
--- @param recipe_name data.RecipeID
--- @param product_name data.FluidID|data.ItemID
function data_util.remove_product(recipe_name, product_name)
    local recipe = data.raw.recipe[recipe_name]
    if not recipe then
        data_util.error("Recipe " .. recipe_name .. " does not exist.")
        return
    end
    if not recipe.results then
        data_util.error("Recipe " .. recipe_name .. " has no products.")
        return
    end
    for i, product in pairs(recipe.results) do
        if product.name == product_name then
            table.remove(recipe.results, i)
            return
        end
    end
    data_util.error("Recipe " .. recipe_name .. " does not have product " .. product_name .. ".")
end

--- @param lab_name data.EntityID
--- @param input_name data.ItemID
function data_util.remove_lab_input(lab_name, input_name)
    local lab = data.raw.lab[lab_name]
    data_util.assert(lab, "Lab " .. lab_name .. " does not exist.")
    local inputs = lab.inputs
    data_util.assert(inputs, "Lab inputs are nil.")
    for i = 1, #inputs do
        if inputs[i] == input_name then
            table.remove(inputs, i)
            return
        end
    end
    data_util.error("Lab " .. lab_name .. " does not contain input " .. input_name .. ".")
end

--- @param energy_source data.EnergySource
--- @param category_name data.FuelCategoryID
function data_util.add_fuel_category(energy_source, category_name)
    if energy_source.type ~= "burner" then
        data_util.error("Energy source must be a burner to add a fuel category.")
        return
    end
    if not energy_source.fuel_categories then
        energy_source.fuel_categories = {}
    end
    table.insert(energy_source.fuel_categories, category_name)
end

return data_util