545 lines
18 KiB
Lua

--- Extends Lua 5.2 table.
-- @module table
-- @see table
local M = {}
local insert = table.insert
--- Given a mapping function, creates a transformed copy of the table
--- by calling the function for each element in the table, and using
--- the result as the new value for the key. Passes the index as second argument to the function.
--- @usage a= { 1, 2, 3, 4, 5}
---table.map(a, function(v) return v * 10 end) --produces: { 10, 20, 30, 40, 50 }
--- @usage a = {1, 2, 3, 4, 5}
---table.map(a, function(v, k, x) return v * k + x end, 100) --produces { 101, 104, 109, 116, 125}
-- @tparam table tbl the table to be mapped to the transform
-- @tparam function func the function to transform values
-- @param[opt] ... additional arguments passed to the function
-- @treturn table a new table containing the keys and mapped values
function M.map(tbl, func, ...)
local newtbl = {}
for i, v in pairs(tbl) do
newtbl[i] = func(v, i, ...)
end
return newtbl
end
--- Given a filter function, creates a filtered copy of the table
--- by calling the function for each element in the table, and
--- filtering out any key-value pairs for non-true results. Passes the index as second argument to the function.
--- @usage a= { 1, 2, 3, 4, 5}
---table.filter(a, function(v) return v % 2 == 0 end) --produces: { 2, 4 }
--- @usage a = {1, 2, 3, 4, 5}
---table.filter(a, function(v, k, x) return k % 2 == 1 end) --produces: { 1, 3, 5 }
-- @tparam table tbl the table to be filtered
-- @tparam function func the function to filter values
-- @param[opt] ... additional arguments passed to the function
-- @treturn table a new table containing the filtered key-value pairs
function M.filter(tbl, func, ...)
local newtbl = {}
local add = #tbl > 0
for k, v in pairs(tbl) do
if func(v, k, ...) then
if add then
insert(newtbl, v)
else
newtbl[k] = v
end
end
end
return newtbl
end
--- Given a candidate search function, iterates over the table, calling the function
--- for each element in the table, and returns the first element the search function returned true.
--- Passes the index as second argument to the function.
--- @usage a= { 1, 2, 3, 4, 5}
---table.find(a, function(v) return v % 2 == 0 end) --produces: 2
--- @usage a = {1, 2, 3, 4, 5}
---table.find(a, function(v, k, x) return k % 2 == 1 end) --produces: 1
-- @tparam table tbl the table to be searched
-- @tparam function func the function to use to search for any matching element
-- @param[opt] ... additional arguments passed to the function
-- @treturn ?|nil|Mixed the first found value, or nil if none was found
function M.find(tbl, func, ...)
for k, v in pairs(tbl) do
if func(v, k, ...) then
return v, k
end
end
return nil
end
--- Given a candidate search function, iterates over the table, calling the function
-- for each element in the table, and returns true if search function returned true.
-- Passes the index as second argument to the function.
-- @see table.find
--- @usage a= { 1, 2, 3, 4, 5}
---table.any(a, function(v) return v % 2 == 0 end) --produces: true
--- @usage a = {1, 2, 3, 4, 5}
---table.any(a, function(v, k, x) return k % 2 == 1 end) --produces: true
-- @tparam table tbl the table to be searched
-- @tparam function func the function to use to search for any matching element
-- @param[opt] ... additional arguments passed to the function
-- @treturn boolean true if an element was found, false if none was found
function M.any(tbl, func, ...)
return M.find(tbl, func, ...) ~= nil
end
--- Given a function, apply it to each element in the table.
-- Passes the index as the second argument to the function.
-- <p>Iteration is aborted if the applied function returns true for any element during iteration.
-- @usage
-- a = {10, 20, 30, 40}
-- table.each(a, function(v) game.print(v) end) --prints 10, 20, 30, 40, 50
-- @tparam table tbl the table to be iterated
-- @tparam function func the function to apply to elements
-- @param[opt] ... additional arguments passed to the function
-- @treturn table the table where the given function has been applied to its elements
function M.each(tbl, func, ...)
for k, v in pairs(tbl) do
if func(v, k, ...) then
break
end
end
return tbl
end
--- Returns a new array that is a one-dimensional recursive flattening of the given array.
-- For every element that is an array, extract its elements into the new array.
-- <p>The optional level argument determines the level of recursion to flatten.
--> This function flattens an integer-indexed array, but not an associative array.
-- @tparam array tbl the array to be flattened
-- @tparam[opt] uint level recursive levels, or no limit to recursion if not supplied
-- @treturn array a new array that represents the flattened contents of the given array
function M.flatten(tbl, level)
local flattened = {}
M.each(
tbl,
function(value)
if type(value) == 'table' and #value > 0 then
if level then
if level > 0 then
M.merge(flattened, M.flatten(value, level - 1), true)
else
insert(flattened, value)
end
else
M.merge(flattened, M.flatten(value), true)
end
else
insert(flattened, value)
end
end
)
return flattened
end
--- Given an array, returns the first element or nil if no element exists.
-- @tparam array tbl the array
-- @treturn ?|nil|Mixed the first element
function M.first(tbl)
return tbl[1]
end
--- Given an array, returns the last element or nil if no elements exist.
-- @tparam array tbl the array
-- @treturn ?|nil|Mixed the last element or nil
function M.last(tbl)
local size = #tbl
if size == 0 then
return nil
end
return tbl[size]
end
--- Given an array of only numeric values, returns the minimum or nil if no element exists.
-- @tparam {number,...} tbl the array with only numeric values
-- @treturn ?|nil|number the minimum value
function M.min(tbl)
if #tbl == 0 then
return nil
end
local min = tbl[1]
for _, num in pairs(tbl) do
min = num < min and num or min
end
return min
end
---Given an array of only numeric values, returns the maximum or nil if no element exists.
-- @tparam {number,...} tbl the array with only numeric values
-- @treturn ?|nil|number the maximum value
function M.max(tbl)
if #tbl == 0 then
return nil
end
local max = tbl[1]
for _, num in pairs(tbl) do
max = num > max and num or max
end
return max
end
--- Given an array of only numeric values, return the sum of all values, or 0 for empty arrays.
-- @tparam {number,...} tbl the array with only numeric values
-- @treturn number the sum of the numbers or zero if the given array was empty
function M.sum(tbl)
local sum = 0
for _, num in pairs(tbl) do
sum = sum + num
end
return sum
end
--- Given an array of only numeric values, returns the average or nil if no element exists.
-- @tparam {number,...} tbl the array with only numeric values
-- @treturn ?|nil|number the average value
function M.avg(tbl)
local cnt = #tbl
return cnt ~= 0 and M.sum(tbl) / cnt or nil
end
--- Merges two tables, values from first get overwritten by the second.
-- @usage
-- function some_func(x, y, args)
-- args = table.merge({option1=false}, args)
-- if opts.option1 == true then return x else return y end
-- end
-- some_func(1,2) -- returns 2
-- some_func(1,2,{option1=true}) -- returns 1
-- @tparam table tblA first table
-- @tparam table tblB second table
-- @tparam[opt=false] boolean array_merge set to true to merge the tables as an array or false for an associative array
-- @tparam[opt=false] boolean raw use rawset for associated array
-- @treturn array|table an array or an associated array where tblA and tblB have been merged
function M.merge(tblA, tblB, array_merge, raw)
if not tblB then
return tblA
end
if array_merge then
for _, v in pairs(tblB) do
table.insert(tblA, v)
end
else
for k, v in pairs(tblB) do
if raw then
rawset(tblA, k, v)
else
tblA[k] = v
end
end
end
return tblA
end
--- Creates a new merged dictionary, if the values in tbl_b are in tbl_a they are not overwritten.
-- @usage
-- local a = {one = A}
-- local b = {one = Z, two = B}
-- local merged = table.dictionary_merge(tbl_a, tbl_b)
-- --merged = {one = A, two = B}
-- @tparam table tbl_a
-- @tparam table tbl_b
-- @treturn table with a and b merged together
function M.dictionary_merge(tbl_a, tbl_b)
local meta_a = getmetatable(tbl_a)
local meta_b = getmetatable(tbl_b)
setmetatable(tbl_a, nil)
setmetatable(tbl_b, nil)
local new_t = {}
for k, v in pairs(tbl_a) do
new_t[k] = v
end
for k, v in pairs(tbl_b or {}) do
if not new_t[k] then
new_t[k] = v
end
end
setmetatable(tbl_a, meta_a)
setmetatable(tbl_b, meta_b)
return new_t
end
--- Compares 2 tables for inner equality.
-- copied from factorio/data/core/lualib/util.lua
-- @tparam table tbl_a
-- @tparam table tbl_2
-- @treturn boolean if the tables are the same
function M.compare(tbl_a, tbl_b)
if tbl_a == tbl_b then
return true
end
for k, v in pairs(tbl_a) do
if type(v) == 'table' and type(tbl_b[k]) == 'table' then
if not M.compare(v, tbl_b[k]) then
return false
end
else
if (v ~= tbl_b[k]) then
return false
end
end
end
for k, v in pairs(tbl_b) do
if type(v) == 'table' and type(tbl_a[k]) == 'table' then
if not M.compare(v, tbl_a[k]) then
return false
end
else
if v ~= tbl_a[k] then
return false
end
end
end
return true
end
--- Creates a deep copy of table without copying Factorio objects.
-- copied from factorio/data/core/lualib/util.lua
-- @usage local copy = table.deepcopy[data.raw.["stone-furnace"]["stone-furnace"]] -- returns a copy of the stone furnace entity
-- @tparam table object the table to copy
-- @treturn table a copy of the table
function M.deepcopy(object)
local lookup_table = {}
local function _copy(this_object)
if type(this_object) ~= 'table' then
return this_object
elseif this_object.__self then
return this_object
elseif lookup_table[this_object] then
return lookup_table[this_object]
end
local new_table = {}
lookup_table[this_object] = new_table
for index, value in pairs(this_object) do
new_table[_copy(index)] = _copy(value)
end
return setmetatable(new_table, getmetatable(this_object))
end
return _copy(object)
end
--- Creates a flexible deep copy of an object, recursively copying sub-objects
-- @usage local copy = table.flexcopy(data.raw.["stone-furnace"]["stone-furnace"]) -- returns a copy of the stone furnace entity
-- @tparam table object the table to copy
-- @treturn table a copy of the table
function M.flexcopy(object)
local lookup_table = {}
local function _copy(this_object)
if type(this_object) ~= 'table' then
return this_object
elseif this_object.__self then
return this_object
elseif lookup_table[this_object] then
return lookup_table[this_object]
elseif type(this_object._copy_with) == 'function' then
lookup_table[this_object] = this_object:_copy_with(_copy)
return lookup_table[this_object]
end
local new_table = {}
lookup_table[this_object] = new_table
for index, value in pairs(this_object) do
new_table[_copy(index)] = _copy(value)
end
return setmetatable(new_table, getmetatable(this_object))
end
return _copy(object)
end
--- Returns a copy of all of the values in the table.
-- @tparam table tbl the table to copy the keys from, or an empty table if tbl is nil
-- @tparam[opt] boolean sorted whether to sort the keys (slower) or keep the random order from pairs()
-- @tparam[opt] boolean as_string whether to try and parse the values as strings, or leave them as their existing type
-- @treturn array an array with a copy of all the values in the table
function M.values(tbl, sorted, as_string)
if not tbl then
return {}
end
local valueset = {}
local n = 0
if as_string then --checking as_string /before/ looping is faster
for _, v in pairs(tbl) do
n = n + 1
valueset[n] = tostring(v)
end
else
for _, v in pairs(tbl) do
n = n + 1
valueset[n] = v
end
end
if sorted then
table.sort(
valueset,
function(x, y) --sorts tables with mixed index types.
local tx = type(x) == 'number'
local ty = type(y) == 'number'
if tx == ty then
return x < y and true or false --similar type can be compared
elseif tx == true then
return true --only x is a number and goes first
else
return false --only y is a number and goes first
end
end
)
end
return valueset
end
--- Returns a copy of all of the keys in the table.
-- @tparam table tbl the table to copy the keys from, or an empty table if tbl is nil
-- @tparam[opt] boolean sorted whether to sort the keys (slower) or keep the random order from pairs()
-- @tparam[opt] boolean as_string whether to try and parse the keys as strings, or leave them as their existing type
-- @treturn array an array with a copy of all the keys in the table
function M.keys(tbl, sorted, as_string)
if not tbl then
return {}
end
local keyset = {}
local n = 0
if as_string then --checking as_string /before/ looping is faster
for k, _ in pairs(tbl) do
n = n + 1
keyset[n] = tostring(k)
end
else
for k, _ in pairs(tbl) do
n = n + 1
keyset[n] = k
end
end
if sorted then
table.sort(
keyset,
function(x, y) --sorts tables with mixed index types.
local tx = type(x) == 'number'
local ty = type(y) == 'number'
if tx == ty then
return x < y and true or false --similar type can be compared
elseif tx == true then
return true --only x is a number and goes first
else
return false --only y is a number and goes first
end
end
)
end
return keyset
end
--- Removes keys from a table by setting the values associated with the keys to nil.
-- @usage local a = {1, 2, 3, 4}
--table.remove_keys(a, {1,3}) --returns {nil, 2, nil, 4}
-- @usage local b = {k1 = 1, k2 = 'foo', old_key = 'bar'}
--table.remove_keys(b, {'old_key'}) --returns {k1 = 1, k2 = 'foo'}
-- @tparam table tbl the table to remove the keys from
-- @tparam {Mixed,...} keys an array of keys that exist in the given table
-- @treturn table tbl without the specified keys
function M.remove_keys(tbl, keys)
for i = 1, #keys do
tbl[keys[i]] = nil
end
return tbl
end
--- Returns the number of keys in a table, if func is passed only count keys when the function is true.
-- @tparam table tbl to count keys
-- @tparam[opt] function func to increment counter
-- @param[optchain] ... additional arguments passed to the function
-- @treturn number The number of keys matching the function or the number of all keys if func isn't passed
-- @treturn number The total number of keys
-- @usage local a = { 1, 2, 3, 4, 5}
-- table.count_keys(a) -- produces: 5, 5
-- @usage local a = {1, 2, 3, 4, 5}
-- table.count_keys(a, function(v, k) return k % 2 == 1 end) -- produces: 3, 5
function M.count_keys(tbl, func, ...)
local count, total = 0, 0
if type(tbl) == 'table' then
for k, v in pairs(tbl) do
total = total + 1
if func then
if func(v, k, ...) then
count = count + 1
end
else
count = count + 1
end
end
end
return count, total
end
--- Returns an inverted (***{[value] = key,...}***) copy of the given table. If the values are not unique, the assigned key depends on the order of pairs().
-- @usage local a = {k1 = 'foo', k2 = 'bar'}
--table.invert(a) --returns {'foo' = k1, 'bar' = k2}
-- @usage local b = {k1 = 'foo', k2 = 'bar', k3 = 'bar'}
--table.invert(b) --returns {'foo' = k1, 'bar' = ?}
-- @tparam table tbl the table to invert
-- @treturn table a new table with inverted mapping
function M.invert(tbl)
local inverted = {}
for k, v in pairs(tbl) do
inverted[v] = k
end
return inverted
end
local function _size(tbl)
local count = 0
for _ in pairs(tbl or {}) do
count = count + 1
end
return count
end
--- Return the size of a table using the factorio built in table_size function
-- @function size
-- @tparam table table to use
-- @treturn int size of the table
M.size = _G.table_size or _size
--- For all string or number values in an array map them to a value = value table
-- @usage local a = {"v1", "v2"}
-- table.array_to_bool(a) -- return {["v1"] = "v1", ["v2"]= "v2"}
-- @tparam table tbl the table to convert
-- @treturn table the converted table
function M.array_to_dictionary(tbl)
local newtbl = {}
for _, v in ipairs(tbl) do
if type(v) == 'string' or type(v) == 'number' then
newtbl[v] = true
end
end
return newtbl
end
--- Does the table contain any elements
-- @tparam table tbl
-- @treturn boolean
function M.is_empty(tbl)
return _G.table_size and _G.table_size(tbl) == 0 or next(tbl) == nil
end
--- Clear all elements in a table
-- @tparam table tbl the table to clear
-- @treturn table the cleared table
function M.clear(tbl)
for k in pairs(tbl) do
tbl[k] = nil
end
return tbl
end
-- Extend into table
for k, v in pairs(M) do
table[k] = v -- luacheck: globals table (Allow mutating global table)
end
return M