init.lua

-- NOTE: everything must be in this file, unless you want the profiler_record
-- hook to be slower filtering out jitprofiler files.

if not core.global_exists("jit") then
    core.log("warning", "[jitprofiler] " ..
    "Since LuaJIT is not being used, this mod is disabled.")
    return
end

local MP = core.get_modpath(core.get_current_modname()):gsub('\\', '/') ---@type string
local STORAGE = core.get_mod_storage()

-- mapping between frame and the actual function
local frame_map = {
    -- it is actually fast_new, but engine code often use this through push_vector so yeah
    -- deactivated for now
    -- [';@B/common/vector.lua:416'] = ';@B/common/vector.lua:push_vector'
}

-- filter out noisy and unhelpful stack dumps
local stack_filter = {
    'LUANTI;@B/common/register%.lua:core%.run_callbacks %d+\n',
    'LUANTI;@B/common/vector%.lua:push_vector %d+\n',
}

local function escape_patterns(s)
    return s:gsub("([%^$()%%.%[%]*+?-])", "%%%1")
end

-- these are used to remove clutter in the <module> part of each stack frame
-- most of these are not needed for RUN_IN_PLACE builds (primarily windows)
-- but needed for system installs (primarily linux)
-- won't work for added mod paths

-- *_PATH/*_PARENT are forward slash normalized, but unresolved and preserve duped path separators
-- also they have trailing slash

-- <luanti-sys>/builtin/common/vector.lua
---@type string
local SYSTEM_PATH = debug.getinfo(vector.new, "S").source
    :gsub('\\', '/'):match("@(.+/)builtin/common/vector%.lua$")
local SYSTEM_PATH_PAT = ';' ..escape_patterns(SYSTEM_PATH)

-- <luanti-mod>/mods/jitprofiler
local MODS_PARENT = MP:match('(.+/)mods/jitprofiler$') ---@type string
local MODS_PARENT_PAT = ';' ..escape_patterns(MODS_PARENT)

-- <luanti-game>/games/<game>
---@type string
local GAMES_PARENT = core.get_game_info().path:gsub('\\', '/'):match('(.+/)games/')
local GAMES_PARENT_PAT = ';' ..escape_patterns(GAMES_PARENT)

-- <luanti-world>/worlds/<worldname>/
local WORLD_PATH = core.get_worldpath():gsub('\\', '/') ..'/'
local WORLD_PATH_PAT = ';' ..escape_patterns(WORLD_PATH)

---@type string
local GAME_MODS_PATH = ';games/' ..core.get_game_info().id ..'/mods'

local function ms_clock()
    return core.get_us_time() / 1e3
end

---@param ms number
local function pretty_time(ms)
    if ms >= 1e5 then
        return ('%ds'):format(math.floor(ms / 1e3))
    elseif ms >= 1e3 then
        return string.format(
            '%.' ..(5 - math.max(0, math.floor(math.log10(math.abs(ms))))) ..'fs',
            ms / 1e3)
    else
        return ('%dms'):format(math.floor(ms))
    end
end

local pretty_mem
do
-- it's going to do a table lookup anyways
local byte_div = {
    [0] = 1, 1, 1,
    1e3, 1e3, 1e3,
    1e6, 1e6, 1e6
}
local byte_sizes = {
    [0] = '', '', '',
    'K', 'K', 'K',
    'M', 'M', 'M',
}
---@param sz integer
function pretty_mem(sz)
    if sz < 1e3 then
        return tostring(sz) ..'B'
    end
    local i = math.floor(math.log10(sz))
    return ('%.2f' ..byte_sizes[i] ..'B'):format(sz / byte_div[i])
end
end

local function norm_path(path)
    return path:gsub('\\', '/'):gsub('/+', '/')
end

---@param path string
local function resolve_path(path)
    -- TODO very fragile stuff
    if not (path:sub(1,1) == '/' or path:sub(2,3) == ':/' or path:sub(2,3) == ':\\') then
        error('cannot resolve relative paths')
    end
    local path = path
        :gsub('/%.$', ''):gsub('/%./', '/')
        :gsub('[^/]+/%.%.$', ''):gsub('[^/]+/%.%./', '')
    assert(path:find('/../', nil, true) == nil, 'could not resolve  .. inside a path/filename')
    return path
end

local DATA_DIR = resolve_path(norm_path(core.get_worldpath())) .. "/jitprofiles"
assert(core.mkdir(DATA_DIR), "Could not create profile directory " .. DATA_DIR)
local MAPPINGS_DIR = resolve_path(core.get_mod_data_path()) .. "/mappings"
assert(core.mkdir(MAPPINGS_DIR), "Could not create profile directory " .. MAPPINGS_DIR)

---@param path string
---@param open function
local function read_file(path, open, reason)
    local file = open(path, 'r') ---@type file*?
    if not file then error('unable to read file ' ..reason ..': ' ..path) end
    local data = file:read("*a") ---@type string?
    if not data then error('unable to read file ' ..reason ..': ' ..path) end
    file:close()
    if data:sub(#data) ~= '\n' then data = data ..'\n' end
    return data
end

---@param path string
---@param open function
---@param data string|string.buffer
local function write_file(path, open, reason, data)
    local file = open(path, 'w') ---@type file*?
    if not file then error('unable to write file ' ..reason ..': ' ..path) end
    ---@cast data string
    if not file:write(data) then error('unable to write file ' ..reason ..': ' ..path) end
    file:close()
end

-- these are parallel lists. they are implicitly linked.

local zone_queue = false
local zones = {}
local zones_hashmap = {}
local zone_points = {} -- (check)point being the stack when zone is called
-- NOTE: It is important that zones are fast
local function zone_noop() end
--- @param name string
local function zone_start(name)
    local count = #zones
    if name:find('[^A-Za-z0-9_ :.-]') then
        error(('bad characters in zone %q. Allowed chars are "-A-Za-z0-9_ :."'):format(name))
    end
    if zones_hashmap[name] then
        error(('attempt to start zone %q twice! are you missing a zone_stop()?'):format(name))
    end

    count = count + 1
    zones_hashmap[name] = true
    zones[count] = name
    zone_queue = true

    -- while zone_queue do end -- force the profiler to record

    -- local info = debug.getinfo(2, "Sn")
    -- local name = info.name
    -- if not name or name == '' then name = tostring(info.linedefined) end
    -- zone_points[count] = (';%s:%s'):format(info.source:sub(2), name)
    -- core.chat_send_all(zone_points[count])
end
--- @param name string
local function zone_stop(name)
    local count = #zones
    if not zones_hashmap[name] then
        error(('attempt to end the zone %q without starting it! are you missing a zone_start()?')
            :format(name))
    end
    if zones[count] ~= name then
        error(('attempting to stop the zone %q without stopping its children'):format(name))
    end
    zones_hashmap[name] = nil
    zones[count] = nil
    zone_points[count] = nil
    count = count - 1
end

local profiler_start, profiler_stop, profiler_learn
local raw_data_buffer ---@type string.buffer?
do
    local ie = core.request_insecure_environment()
    if not ie then error("To use jitprofiler, add it to secure.trusted_mods") end

    local require = ie.require
    local strbuf = require('string.buffer')
    local profile = require("jit.profile")

    local jpstart, jpstop, jpdumpstack = profile.start, profile.stop, profile.dumpstack
    local danger_open = ie.io.open
    local clock = core.get_us_time ---@type fun(): number
    local time_point ---@type number

    local scorer
    local function default_scorer(time, samples)
        return time / samples * 0.01
    end

    -- this shows when the sample is collected, not how much time spent inside.
    local vmstate_mapping = {
        N = ';(recorded during JIT compiled trace)',
        I = ';(recorded while interpreting)',
        C = ';(recorded via C/C++ call)',
        J = ';(recorded while compiling trace)',
        G = ';(recorded while collecting garbage)',
    }

    local profiler_record
    -- NOTE: It is *VERY* important that profiler_record do as less as possible
    -- in order to allow for lower intervals/higher frequencies of sampling.
    function profiler_record(thread, samples, vmstate)
        local stack = jpdumpstack(thread, "pFZ;", -100)
        local _, found_mod = string.find(norm_path(stack), 'mods/jitprofiler/init.lua', nil, true)
        if found_mod then
            time_point = clock()
            return
        end

        -- TODO figure out a way to make zones work

        local now = clock()
        ---@cast raw_data_buffer -nil
        raw_data_buffer:putf('LUANTI;%s%s %d\n',
            stack, vmstate_mapping[vmstate],
            math.floor(scorer(now - time_point, samples)))
        time_point = now
    end

    local duration_start ---@type number
    local data_file ---@type string?
    local open_data ---@type function

    ---@param period number interval between sampling in milliseconds
    ---@param fname string data file to update
    ---@param open function? io.open, for data file
    function profiler_start(period, fname, open)
        duration_start = clock()
        data_file = resolve_path(norm_path(fname))
        open_data = open or io.open
        raw_data_buffer = strbuf.new()
        scorer = jitprofiler.scorer or default_scorer

        -- zones = {}
        -- zones_hashmap = {}
        -- zone_points = {}
        -- jitprofiler.zone_start = zone_start
        -- jitprofiler.zone_stop = zone_stop

        jpstart("vfi" ..tostring(math.max(1, tonumber(period))), profiler_record)
        time_point = clock()
    end

    function profiler_stop()
        jpstop()
        local data_size = #raw_data_buffer
        local duration = (clock() - duration_start) / 1e3

        ---@cast raw_data_buffer -nil
        ---@cast data_file -nil
        write_file(data_file, open_data, 'storing raw profiling data', raw_data_buffer)

        -- zones = {}
        -- zones_hashmap = {}
        -- zone_points = {}
        -- jitprofiler.zone_start = zone_noop
        -- jitprofiler.zone_stop = zone_noop

        data_file = nil
        raw_data_buffer = nil
        return duration, data_size
    end

    -- TODO decide what to report
    ---@param data_path string pretty data file to read
    ---@param open_data function? io.open, for pretty data file
    ---@param map_path string mapping file to write or update
    ---@param open_map function? io.open, for mapping file
    function profiler_learn(data_path, map_path, open_data, open_map)
        if raw_data_buffer then
            error('attempt to learn profiler data while a profiler is running. not allowed')
        end

        local data_path = resolve_path(norm_path(data_path))
        local map_path = resolve_path(norm_path(map_path))

        local open_data = open_data or io.open
        local open_map = open_map or io.open

        local data = read_file(data_path, open_data, 'getting pretty profiling data')
            :gsub(';@B', ';' ..SYSTEM_PATH ..'builtin')
            :gsub(';@W', ';' ..WORLD_PATH ..'worldmods')
            :gsub(';@G', ';' ..GAMES_PARENT ..GAME_MODS_PATH:sub(2))
            :gsub(';@M', ';' ..MODS_PARENT ..'mods')

        local mappings = {}
        local rfile = open_map(map_path, 'r')
        if rfile then
            local mappings_raw = rfile:read('*a') ---@type string?
            assert(mappings_raw, 'unable to read file getting existing frame mappings: ' ..map_path)
            rfile:close()
            local f = loadstring(mappings_raw
                :gsub(';@B', ';' ..SYSTEM_PATH ..'builtin')
                :gsub(';@W', ';' ..WORLD_PATH ..'worldmods')
                :gsub(';@G', ';' ..GAMES_PARENT ..GAME_MODS_PATH:sub(2))
                :gsub(';@M', ';' ..MODS_PARENT ..'mods'), map_path)
            assert(f, 'unable to load file getting existing frame mappings' ..map_path)
            setfenv(f, {})
            mappings = f()
        end

        local mapping_keys = {}
        local done_frame = {}
        for entry in data:gmatch('([^\n]+)\n') do
            local frames = entry:match('^LUANTI;(.*) %d+$'):split(';')
            for f=1, #frames do
                local frame = frames[f]
                if done_frame[frame] then goto continue end
                done_frame[frame] = true

                ---@type string?, integer?
                local module_path, lnnum = frame:match('^(.+%.lua):(%d+)$')
                lnnum = tonumber(lnnum)
                if not module_path or not lnnum then goto continue end
                local module_raw = read_file(module_path, danger_open, 'getting frame info')

                local start = 1
                for c=1, lnnum-1 do -- consume all the previous lines
                    local nx = module_raw:find('\n', start, true)
                    if not nx and c < lnnum-1 then
                        core.log('warning', 'EOF. potential change outside of Luanti. Unable to find the line in this frame: ' ..frame)
                        goto continue
                    end
                    start = nx+1
                end

                local line = module_raw:sub(start, module_raw:find('\n', start, true) - 1)

                ---@type string?
                local fn_name = line:match('function%s+([A-Za-z_][A-Za-z0-9_.:]*)%s*%(')
                if not fn_name then
                    fn_name = line:match('([A-Za-z_][A-Za-z0-9_.:]*)%s*=%s*function%s*%(')
                end
                if not fn_name then
                    local registrar = line:match('([A-Za-z0-9_.:]*register_[A-Za-z0-9_.:]*)%s*%(.*function%s*%(')
                    if registrar then fn_name = ('%s HOOK'):format(registrar) end
                end

                if not fn_name then
                    core.log('info', 'failed to find function in this frame: ' ..frame)
                    goto continue
                end

                local key = ';' ..frame
                local value = (';%s:%d (%s)'):format(module_path, lnnum, fn_name)
                mappings[key] = value
                mapping_keys[#mapping_keys+1] = key
                ::continue::
            end
        end
---@diagnostic disable-next-line: cast-local-type
        data = nil

        table.sort(mapping_keys)
        local mapping_databuf = strbuf.new():set('return {\n')
        for i=1, #mapping_keys do
            mapping_databuf:put((' [%q] = %q,\n'):format(mapping_keys[i], mappings[mapping_keys[i]]))
        end
        mapping_databuf:put('}\n')

        write_file(map_path, open_map, 'storing frame mappings', mapping_databuf:tostring()
            :gsub(WORLD_PATH_PAT, ';')
            :gsub(GAMES_PARENT_PAT, ';')
            :gsub(MODS_PARENT_PAT, ';')
            :gsub(SYSTEM_PATH_PAT, ';')
            :gsub(';builtin', ';@B')
            :gsub(';worldmods', ';@W')
            :gsub(GAME_MODS_PATH, ';@G')
            :gsub(';mods', ';@M'))
        return mappings
    end
end

---@param data_path string data file to prettify
---@param open_data function? io.open, for data file
local function profiler_pretty(data_path, open_data)
    if raw_data_buffer then
        error('attempt to pretty profiler data while a profiler is running. not allowed')
    end

    local data_path = resolve_path(norm_path(data_path))
    local open_data = open_data or io.open

    -- norm_path normalizes potential weird paths
    local data = norm_path(read_file(data_path, open_data, 'getting raw profiling data'))
        :gsub(WORLD_PATH_PAT, ';')
        :gsub(GAMES_PARENT_PAT, ';')
        :gsub(MODS_PARENT_PAT, ';')
        :gsub(SYSTEM_PATH_PAT, ';')
        :gsub(';builtin', ';@B')
        :gsub(';worldmods', ';@W')
        :gsub(GAME_MODS_PATH, ';@G')
        :gsub(';mods', ';@M')

    if next(frame_map) then data = data:gsub(';[^;\n]+\n', frame_map) end

    for f=1, #stack_filter do
        data = data:gsub(stack_filter[f], '')
    end

    if jitprofiler.trim_zeroes then
        data = data:gsub('LUANTI;[^\n]+ 0\n', '')
    end

    local data_size = #data
    write_file(data_path, open_data, 'storing pretty profiling data', data)
    return data_size
end

---@param data_path string pretty data file to read and update
---@param open_data function? io.open, for pretty data file
---@param map_path string mapping file to read
---@param open_map function? io.open, for mapping file
local function profiler_remap(data_path, map_path, open_data, open_map)
    if raw_data_buffer then
        error('attempt to pretty profiler data while a profiler is running. not allowed')
    end

    local data_path = resolve_path(norm_path(data_path))
    local map_path = resolve_path(norm_path(map_path))

    local open_data = open_data or io.open
    local open_map = open_map or io.open

    local f = loadstring(read_file(map_path, open_map, 'getting frame mapping'), map_path)
    assert(f, 'unable to load file getting existing frame mappings' ..map_path)
    setfenv(f, {})
    local mappings = f()

    for k, v in pairs(jitprofiler.frame_map) do
        mappings[k] = v
    end

    local data = read_file(data_path, open_data, 'getting pretty profiler data')
    write_file(data_path, open_data, 'storing pretty profiler data', data
        :gsub(';[^;]+', mappings))
end

local cmd_period = STORAGE:get('last_period')
local cmd_data_fname = STORAGE:get('last_data_fname')
local cmd_data_path

core.register_chatcommand("jp_start", {
    description = "Start profiler",
    params = "<period> <data filename>",
    privs = {server = true},

    ---@param args string
    func = function(_, args)
        if raw_data_buffer then
            return false, "Profiler already running"
        end

        local args = args:trim()
        if args:find('[/\\]') then return false, 'bad characters "/\\"' end

        ---@type string?, string?
        local period, fname = args:match("^(%d+)%s+(.*)$")
        if not period and not fname then fname = args end

        local period_num = tonumber(period)
        if period and not period_num then
            return false, 'bad argument to <period>, expecting integer got ' ..period
        end
        if not period_num and not cmd_period then
            return false, 'required <period> argument as no profiling was conducted previously'
        end

        if fname == '' then fname = nil end
        if not fname and not cmd_data_fname then
            return false, 'required <data filename> argument as no profiling was conducted previously'
        end

        cmd_period = period_num and math.floor(period_num) or cmd_period
        cmd_data_fname = fname or cmd_data_fname

        STORAGE:set_int('last_period', cmd_period)
        STORAGE:set_string('last_data_fname', cmd_data_fname)

        cmd_data_path = DATA_DIR .. "/" .. cmd_data_fname
        profiler_start(cmd_period, cmd_data_path)
        return true, "Profiler started. Writing data to " .. cmd_data_path
    end,
})

core.register_chatcommand("jp_stop", {
    description = "Stop profiler",
    privs = {server = true},
    func = function()
        if not raw_data_buffer then
            return false, "Profiler not running"
        end

        local duration, raw_size = profiler_stop()

        local t0 = ms_clock()
        local pretty_size = profiler_pretty(cmd_data_path)
        return true, ('Profiler stopped after %s. Raw data size was %s. Took %s to prettify data into %s'):format(
            pretty_time(duration), pretty_mem(raw_size),
            pretty_time(ms_clock() - t0), pretty_mem(pretty_size))
    end,
})

---@param args string
local function data_mapping_params(args)
    local comma = args:find(',', nil, true)
    if not comma then return false, 'missing comma separating <data filename> and <mapping filename>' end
    local data_fname = args:sub(1, comma-1):trim()
    local map_fname = args:sub(comma+1):trim()
    return DATA_DIR ..'/' ..data_fname, MAPPINGS_DIR ..'/' ..map_fname ..'.lua'
end

core.register_chatcommand('jp_learn', {
    description = 'Produce frame mapping from profiled data',
    params = "<data filename>, <mapping name>",
    privs = {server = true},
    ---@param args string
    func = function(_, args)
        if raw_data_buffer then
            return false, "Profiler is running. Not allowed"
        end

        local args = args:trim()
        if args == '' then return false, 'missing arguments' end
        if args:find('[/\\]') then return false, 'bad characters "/\\"' end

        local data_path, map_path = data_mapping_params(args)
        if not data_path then return false, map_path end

        local t0 = ms_clock()
        profiler_learn(data_path, map_path) -- TODO report more things
        return true, ('Learnt new frame mappings in %s'):format(
            pretty_time(ms_clock() - t0))
    end
})

core.register_chatcommand('jp_remap', {
    description = 'Apply frame mapping to profiled data',
    params = "<data filename>, <mapping name>",
    privs = {server = true},
    ---@param args string
    func = function(_, args)
        if raw_data_buffer then
            return false, "Profiler is running. Not allowed"
        end

        local args = args:trim()
        if args == '' then return false, 'missing arguments' end
        if args:find('[/\\]') then return false, 'bad characters "/\\"' end

        local data_path, map_path = data_mapping_params(args)
        if not data_path then return false, map_path end

        local t0 = ms_clock()
        profiler_remap(data_path, map_path) -- TODO report more things
        return true, ('Remapped frames in %s'):format(
            pretty_time(ms_clock() - t0))
    end
})

jitprofiler = {
    start = profiler_start,
    stop = profiler_stop,
    pretty = profiler_pretty,
    learn = profiler_learn,
    remap = profiler_remap,

    scorer = nil,

    -- when the profiler is on, it pushes a new zone in between the stacks
    -- zone_start = zone_noop,
    -- when the profiler is on, it pops a zone in between the stacks
    -- zone_stop = zone_noop,

    frame_map = frame_map,
    stack_filter = stack_filter,

    -- trim away stack dumps with 0 score
    trim_zeroes = true,
}