dfhack/library/lua/utils.lua

780 lines
20 KiB
Lua

local _ENV = mkmodule('utils')
local df = df
local getopt = require('3rdparty.alt_getopt')
-- Comparator function
function compare(a,b)
if a < b then
return -1
elseif a > b then
return 1
else
return 0
end
end
-- Sort strings; compare empty last
function compare_name(a,b)
if a == '' then
if b == '' then
return 0
else
return 1
end
elseif b == '' then
return -1
else
return compare(a,b)
end
end
-- Make a field comparator
function compare_field(field,cmp)
cmp = cmp or compare
if field then
return function (a,b)
return cmp(a[field],b[field])
end
else
return cmp
end
end
-- Make a comparator of field vs key
function compare_field_key(field,cmp)
cmp = cmp or compare
if field then
return function (a,b)
return cmp(a[field],b)
end
else
return cmp
end
end
function is_container(obj)
return df.isvalid(obj) == 'ref' and obj._kind == 'container'
end
-- Make a sequence of numbers in 1..size
function make_index_sequence(istart,iend)
local index = {}
for i=istart,iend do
index[i-istart+1] = i
end
return index
end
--[[
Sort items in data according to ordering.
Each ordering spec is a table with possible fields:
* key = function(value)
Computes comparison key from a data value. Not called on nil.
* key_table = function(data)
Computes a key table from the data table in one go.
* compare = function(a,b)
Comparison function. Defaults to compare above.
Called on non-nil keys; nil sorts last.
* nil_first
If true, nil keys are sorted first instead of last.
* reverse
If true, sort non-nil keys in descending order.
Returns a table of integer indices into data.
--]]
function make_sort_order(data,ordering)
-- Compute sort keys and comparators
local keys = {}
local cmps = {}
local size = data.n or #data
for i=1,#ordering do
local order = ordering[i]
if order.key_table then
keys[i] = order.key_table(data)
elseif order.key then
local kt = {}
local kf = order.key
for j=1,size do
if data[j] == nil then
kt[j] = nil
else
kt[j] = kf(data[j])
end
end
keys[i] = kt
else
keys[i] = data
end
cmps[i] = order.compare or compare
end
-- Make an order table
local index = make_index_sequence(1,size)
-- Sort the ordering table
table.sort(index, function(ia,ib)
for i=1,#keys do
local ka = keys[i][ia]
local kb = keys[i][ib]
-- Sort nil keys to the end
if ka == nil then
if kb ~= nil then
return ordering[i].nil_first
end
elseif kb == nil then
return not ordering[i].nil_first
else
local cmpv = cmps[i](ka,kb)
if ordering[i].reverse then
cmpv = -cmpv
end
if cmpv < 0 then
return true
elseif cmpv > 0 then
return false
end
end
end
return ia < ib -- this should ensure stable sort
end)
return index
end
--[[
Iterate a 'list' structure, e.g. df.global.world.job_list
--]]
local function next_df_list(s,link)
link = link.next
if link then
return link, link.item
end
end
function listpairs(list)
return next_df_list, nil, list
end
--[[
Recursively assign data into a table.
--]]
function assign(tgt,src)
if df.isvalid(tgt) == 'ref' then
df.assign(tgt, src)
elseif type(tgt) == 'table' then
for k,v in pairs(src) do
if type(v) == 'table' then
local cv = tgt[k]
if cv == nil then
cv = {}
tgt[k] = cv
end
assign(cv, v)
else
tgt[k] = v
end
end
else
error('Invalid assign target type: '..tostring(tgt))
end
return tgt
end
local function copy_field(obj,k,v,deep)
if v == nil then
return NULL
end
if deep then
local field = obj:_field(k)
if field == v then
return clone(v,deep)
end
end
return v
end
-- Copy the object as lua data structures.
function clone(obj,deep)
if type(obj) == 'table' then
if deep then
return assign({},obj)
else
return copyall(obj)
end
elseif df.isvalid(obj) == 'ref' then
local kind = obj._kind
if kind == 'primitive' then
return obj.value
elseif kind == 'bitfield' then
local rv = {}
for k,v in pairs(obj) do
rv[k] = v
end
return rv
elseif kind == 'container' then
local rv = {}
for k,v in ipairs(obj) do
rv[k+1] = copy_field(obj,k,v,deep)
end
return rv
else -- struct
local rv = {}
for k,v in pairs(obj) do
rv[k] = copy_field(obj,k,v,deep)
end
return rv
end
else
return obj
end
end
local function get_default(default,key,base)
if type(default) == 'table' then
local dv = default[key]
if dv == nil then
dv = default._default
end
if dv == nil then
dv = base
end
return dv
else
return default
end
end
-- Copy the object as lua data structures, skipping values matching defaults.
function clone_with_default(obj,default,force)
local rv = nil
local function setrv(k,v)
if v ~= nil then
if rv == nil then
rv = {}
end
rv[k] = v
end
end
if default == nil then
return nil
elseif type(obj) == 'table' then
for k,v in pairs(obj) do
setrv(k, clone_with_default(v, get_default(default,k)))
end
elseif df.isvalid(obj) == 'ref' then
local kind = obj._kind
if kind == 'primitive' then
return clone_with_default(obj.value,default,force)
elseif kind == 'bitfield' then
for k,v in pairs(obj) do
setrv(k, clone_with_default(v, get_default(default,k,false)))
end
elseif kind == 'container' then
for k,v in ipairs(obj) do
setrv(k+1, clone_with_default(v, default, true))
end
else -- struct
for k,v in pairs(obj) do
setrv(k, clone_with_default(v, get_default(default,k)))
end
end
elseif obj == default and not force then
return nil
elseif obj == nil then
return NULL
else
return obj
end
if force and rv == nil then
rv = {}
end
return rv
end
-- Parse an integer value into a bitfield table
function parse_bitfield_int(value, type_ref)
if value == 0 then
return nil
end
local res = {}
for i,v in ipairs(type_ref) do
if bit32.extract(value, i) ~= 0 then
res[v] = true
end
end
return res
end
-- List the enabled flag names in the bitfield table
function list_bitfield_flags(bitfield, list)
list = list or {}
if bitfield then
for name,val in pairs(bitfield) do
if val then
table.insert(list, name)
end
end
end
return list
end
-- Sort a vector or lua table
function sort_vector(vector,field,cmp)
local fcmp = compare_field(field,cmp)
local scmp = function(a,b)
return fcmp(a,b) < 0
end
if df.isvalid(vector) then
if vector._kind ~= 'container' then
error('Container expected: '..tostring(vector))
end
local items = clone(vector, true)
table.sort(items, scmp)
vector:assign(items)
else
table.sort(vector, scmp)
end
return vector
end
-- Linear search
function linear_index(vector,key,field)
local min,max
if df.isvalid(vector) then
min,max = 0,#vector-1
else
min,max = 1,#vector
end
if field then
for i=min,max do
local obj = vector[i]
if obj[field] == key then
return i, obj
end
end
else
for i=min,max do
local obj = vector[i]
if obj == key then
return i, obj
end
end
end
return nil
end
-- Binary search in a vector or lua table
function binsearch(vector,key,field,cmp,min,max)
if not(min and max) then
if df.isvalid(vector) then
min = -1
max = #vector
else
min = 0
max = #vector+1
end
end
local mf = math.floor
local fcmp = compare_field_key(field,cmp)
while true do
local mid = mf((min+max)/2)
if mid <= min then
return nil, false, max
end
local item = vector[mid]
local cv = fcmp(item, key)
if cv == 0 then
return item, true, mid
elseif cv < 0 then
min = mid
else
max = mid
end
end
end
-- Binary search and insert
function insert_sorted(vector,item,field,cmp)
local key = item
if field and item then
key = item[field]
end
local cur,found,pos = binsearch(vector,key,field,cmp)
if found then
return false,cur,pos
else
if df.isvalid(vector) then
vector:insert(pos, item)
else
table.insert(vector, pos, item)
end
return true,vector[pos],pos
end
end
-- Binary search, then insert or overwrite
function insert_or_update(vector,item,field,cmp)
local added,cur,pos = insert_sorted(vector,item,field,cmp)
if not added then
vector[pos] = item
cur = vector[pos]
end
return added,cur,pos
end
-- Binary search and erase
function erase_sorted_key(vector,key,field,cmp)
local cur,found,pos = binsearch(vector,key,field,cmp)
if found then
if df.isvalid(vector) then
vector:erase(pos)
else
table.remove(vector, pos)
end
end
return found,cur,pos
end
function erase_sorted(vector,item,field,cmp)
local key = item
if field and item then
key = item[field]
end
return erase_sorted_key(vector,key,field,cmp)
end
-- Calls a method with a string temporary
function call_with_string(obj,methodname,...)
return dfhack.with_temp_object(
df.new "string",
function(str,obj,methodname,...)
obj[methodname](obj,str,...)
return str.value
end,
obj,methodname,...
)
end
function getBuildingName(building)
return call_with_string(building, 'getName')
end
function getBuildingCenter(building)
return xyz2pos(building.centerx, building.centery, building.z)
end
function getItemDescription(item,mode)
return call_with_string(item, 'getItemDescription', mode or 0)
end
function getItemDescriptionPrefix(item,mode)
return call_with_string(item, 'getItemDescriptionPrefix', mode or 0)
end
-- Split the string by the given delimiter
function split_string(self, delimiter)
local result = { }
local from = 1
local delim_from, delim_to = string.find( self, delimiter, from )
while delim_from do
table.insert( result, string.sub( self, from , delim_from-1 ) )
from = delim_to + 1
delim_from, delim_to = string.find( self, delimiter, from )
end
table.insert( result, string.sub( self, from ) )
return result
end
-- Ask a yes-no question
function prompt_yes_no(msg,default)
local prompt = msg
if default == nil then
prompt = prompt..' (y/n): '
elseif default then
prompt = prompt..' (y/n)[y]: '
else
prompt = prompt..' (y/n)[n]: '
end
while true do
local rv,err = dfhack.lineedit(prompt)
if not rv then
qerror(err);
elseif string.match(rv,'^[Yy]') then
return true
elseif string.match(rv,'^[Nn]') then
return false
elseif rv == 'abort' then
qerror('User abort')
elseif rv == '' and default ~= nil then
return default
end
end
end
-- Ask for input with check function
function prompt_input(prompt,check,quit_str)
quit_str = quit_str or '~~~'
while true do
local rv,err = dfhack.lineedit(prompt)
if not rv then
qerror(err);
end
if rv == quit_str then
qerror('User abort')
end
local rtbl = table.pack(check(rv))
if rtbl[1] then
return table.unpack(rtbl,2,rtbl.n)
end
end
end
function check_number(text)
local nv = tonumber(text)
return nv ~= nil, nv
end
function invert(tab)
local result = {}
for k,v in pairs(tab) do
result[v]=k
end
return result
end
function processArgs(args, validArgs)
--[[
standardized argument processing for scripts
-argName value
-argName [list of values]
-argName [list of [nested values] -that can be [whatever] format of matched square brackets]
-arg1 \-arg3
escape sequences
--]]
local result = {}
local argName
local bracketDepth = 0
for i,arg in ipairs(args) do
if argName then
if arg == '[' then
if bracketDepth > 0 then
table.insert(result[argName], arg)
end
bracketDepth = bracketDepth+1
elseif arg == ']' then
bracketDepth = bracketDepth-1
if bracketDepth > 0 then
table.insert(result[argName], arg)
else
argName = nil
end
elseif string.sub(arg,1,1) == '\\' then
if bracketDepth == 0 then
result[argName] = string.sub(arg,2)
argName = nil
else
table.insert(result[argName], string.sub(arg,2))
end
else
if bracketDepth == 0 then
result[argName] = arg
argName = nil
else
table.insert(result[argName], arg)
end
end
elseif string.sub(arg,1,1) == '-' then
argName = string.sub(arg,2)
if validArgs and not validArgs[argName] then
error('error: invalid arg: ' .. i .. ': ' .. argName)
end
if result[argName] then
error('duplicate arg: ' .. i .. ': ' .. argName)
end
if i+1 > #args or string.sub(args[i+1],1,1) == '-' then
result[argName] = ''
argName = nil
else
result[argName] = {}
end
else
error('error parsing arg ' .. i .. ': ' .. arg)
end
end
return result
end
-- processes commandline options according to optionActions and returns all
-- argument strings that are not options. Options and non-option strings can
-- appear in any order, and single-letter options that do not take arguments
-- can be combined into a single option string (e.g. '-abc' is the same as
-- '-a -b -c' if options 'a' and 'b' do not take arguments.
--
-- Numbers cannot be options and negative numbers (e.g. -10) will be interpreted
-- as positional parameters and returned in the nonoptions list.
--
-- optionActions is a vector with elements in the following format:
-- {shortOptionName, longOptionAlias, hasArg=boolean, handler=fn}
-- shortOptionName and handler are required. If the option takes an argument,
-- it will be passed to the handler function.
-- longOptionAlias is optional.
-- hasArgument defaults to false.
--
-- example usage:
--
-- local filename = nil
-- local open_readonly = false
-- local nonoptions = processArgsGetopt(args, {
-- {'r', handler=function() open_readonly = true end},
-- {'f', 'filename', hasArg=true,
-- handler=function(optarg) filename = optarg end}
-- })
--
-- when args is {'first', '-f', 'fname', 'second'} or, equivalently,
-- {'first', '--filename', 'fname', 'second'} (note the double dash in front of
-- the long option alias), then filename will be fname and nonoptions will
-- contain {'first', 'second'}.
function processArgsGetopt(args, optionActions)
local sh_opts, long_opts = '', {}
local handlers = {}
for _,optionAction in ipairs(optionActions) do
local sh_opt,long_opt = optionAction[1], optionAction[2]
if not sh_opt or type(sh_opt) ~= 'string' or #sh_opt ~= 1 then
error('optionAction missing option letter at index 1')
end
if not optionAction.handler then
error(string.format('handler missing for option "%s"', sh_opt))
end
sh_opts = sh_opts .. sh_opt
if optionAction.hasArg then sh_opts = sh_opts .. ':' end
handlers[sh_opt] = optionAction.handler
if long_opt then
long_opts[long_opt] = sh_opt
handlers[long_opt] = optionAction.handler
end
end
local opts, optargs, nonoptions =
getopt.get_ordered_opts(args, sh_opts, long_opts)
for i,v in ipairs(opts) do
handlers[v](optargs[i])
end
return nonoptions
end
function fillTable(table1,table2)
for k,v in pairs(table2) do
table1[k] = v
end
end
function unfillTable(table1,table2)
for k,v in pairs(table2) do
table1[k] = nil
end
end
function df_shortcut_var(k)
if k == 'scr' or k == 'screen' then
return dfhack.gui.getCurViewscreen()
elseif k == 'bld' or k == 'building' then
return dfhack.gui.getSelectedBuilding()
elseif k == 'item' then
return dfhack.gui.getSelectedItem()
elseif k == 'job' then
return dfhack.gui.getSelectedJob()
elseif k == 'wsjob' or k == 'workshop_job' then
return dfhack.gui.getSelectedWorkshopJob()
elseif k == 'unit' then
return dfhack.gui.getSelectedUnit()
elseif k == 'plant' then
return dfhack.gui.getSelectedPlant()
else
for g in pairs(df.global) do
if g == k then
return df.global[k]
end
end
return _G[k]
end
end
function df_shortcut_env()
local env = {}
setmetatable(env, {__index = function(self, k) return df_shortcut_var(k) end})
return env
end
df_env = df_shortcut_env()
function df_expr_to_ref(expr)
expr = expr:gsub('%["(.-)"%]', function(field) return '.' .. field end)
:gsub('%[\'(.-)\'%]', function(field) return '.' .. field end)
:gsub('%[(%d+)]', function(field) return '.' .. field end)
local parts = split_string(expr, '%.')
local obj = df_env[parts[1]]
for i = 2, #parts do
local key = tonumber(parts[i]) or parts[i]
local cur = obj[key]
if i == #parts and ((type(cur) ~= 'userdata') or
type(cur) == 'userdata' and getmetatable(cur) == nil) then
obj = obj:_field(key)
else
obj = obj[key]
end
end
return obj
end
function addressof(obj)
return select(2, df.sizeof(obj))
end
function OrderedTable()
-- store values in a separate table to ensure that __index and __newindex
-- run on every table index operation
local t = {}
local key_to_index = {}
local index_to_key = {}
local mt = {}
function mt:__index(k)
return t[k]
end
function mt:__newindex(k, v)
if not key_to_index[k] then
table.insert(index_to_key, k)
key_to_index[k] = #index_to_key
end
t[k] = v
end
function mt:__pairs()
return function(_, k)
if k then
k = index_to_key[key_to_index[k] + 1]
else
k = index_to_key[1]
end
if k then
return k, t[k]
end
end, nil, nil
end
local self = {}
setmetatable(self, mt)
return self
end
return _ENV