pretty/pretty.lua

-- Ensure loading library, if it exists, no matter where pretty.lua was loaded from.

-- Load the library component
local format_number, format_function, analyze_structure, get_table_info, TABLE_TYPE
do
    local thispath = ... and select('1', ...):match('.+%.') or ''
    local function import (name, ignore_failure)
        local was_loaded, lib_or_error      =  pcall(require, thispath..name)
        if not was_loaded then
            if ignore_failure then  return nil  end
            error('[pretty]: Could not load vital library: '..name..'.lua:\n\t'..lib_or_error)
        end
        return lib_or_error
    end

    -- Load number and function formatting
    -- Will use a very simple number formatting, if number.lua is not available.
    -- Will use a very simple function formatting, if function.lua is not available.
    format_number      =  import('number',   true)  or  function (value, _, l)  l[#l+1]  =  tostring(value)  end
    format_function    =  import('function', true)  or  function (value, _, l)  l[#l+1]  =  'function (...) --[['..tostring(value):sub(11)..']] end'  end

    -- Load other stuff
    analyze_structure  =  import 'analyze_structure'
    analyze_structure, get_table_info  =  analyze_structure[1], analyze_structure[2]
    TABLE_TYPE         =  import 'table_type'
end
--

local ERROR_UNKNOWN_TYPE = [[
[pretty]: Attempting to format unsupported value of type "%s".
    A native formatting of the value is: %s

We are attempting to cover all Lua features, so please report this bug, so we can improve.
]]

local NR_CHARS_IN_LONG_STRING         = 40
local MAX_WIDTH_FOR_SINGLE_LINE_TABLE = 38

local KEY_TYPE_SORT_ORDER = {
    ['number']   = 0,
    ['string']   = 1,
    ['boolean']  = 2,
    ['table']    = 3,
    ['userdata'] = 4,
    ['thread']   = 5,
    ['function'] = 6,
}

local VALUE_TYPE_SORT_ORDER = {
    ['nil']      = 0,
    ['boolean']  = 1,
    ['number']   = 2,
    ['string']   = 3,
    ['table']    = 4,
    ['userdata'] = 5,
    ['thread']   = 6,
    ['function'] = 7,
}

local CHAR_TO_STR_REPR = {}

do
    for i = 00, 031 do  CHAR_TO_STR_REPR[i] = '\\0'..(i < 10 and '0' or '')..i  end
    for i = 32, 255 do  CHAR_TO_STR_REPR[i] = string.char(i)  end
    CHAR_TO_STR_REPR[7]   = '\\a'
    CHAR_TO_STR_REPR[8]   = '\\b'
    CHAR_TO_STR_REPR[9]   = '\t'
    CHAR_TO_STR_REPR[10]  = '\n'
    CHAR_TO_STR_REPR[11]  = '\\v'
    CHAR_TO_STR_REPR[12]  = '\\f'
    CHAR_TO_STR_REPR[13]  = '\\r'
    CHAR_TO_STR_REPR[92]  = '\\\\'
    CHAR_TO_STR_REPR[127] = '\\127'
end

--------------------------------------------------------------------------------
-- Util

local function padnum(d)
    local dec, n = string.match(d, "(%.?)0*(.+)")
    return #dec > 0 and ("%.12f"):format(d) or ("%s%03d%s"):format(dec, #n, n)
end

local function alphanum_compare_strings (a, b)
    return tostring(a):gsub("%.?%d+", padnum)..("%3d"):format(#b)
         < tostring(b):gsub("%.?%d+", padnum)..("%3d"):format(#a)
end

local function compare_key_value_pairs (a, b)
    -- Get types
    local type_key_a,   type_key_b   = type(a[1]), type(b[1])
    local type_value_a, type_value_b = type(a[2]), type(b[2])

    -- Tons of compare
    if (1 == (type_key_a == 'number' and 1 or 0) + (type_key_b == 'number' and 1 or 0)) then
        return type_key_a == 'number'
    elseif (type_key_a == 'number' and type_key_b == 'number') then
        return a[1] < b[1]
    elseif (type_value_a ~= type_value_b) then
        return VALUE_TYPE_SORT_ORDER[type_value_a] < VALUE_TYPE_SORT_ORDER[type_value_b]
    elseif (type_key_a == 'string' and type_key_b == 'string') then
        return alphanum_compare_strings(a[1], b[1])
    elseif (type_key_a ~= type_key_b) then
        return KEY_TYPE_SORT_ORDER[type_value_a] < KEY_TYPE_SORT_ORDER[type_value_b]
    end
end

local function get_key_value_pairs_in_proper_order (t)
    -- Generates a sequence of key value pairs, in proper order.
    -- Proper order is:
    -- 1. All integer keys are first, in order
    -- 2. Then by value type, as defined in  VALUE_TYPE_SORT_ORDER in the top.
    -- 3. Then by key type.
    -- 3.1. String in alphanumeric order
    -- 3.2. Other wierdness.

    -- Error checking
    assert(type(t) == 'table')

    -- Do stuff
    local key_value_pairs = {}

    for key, value in pairs(t) do
        key_value_pairs[#key_value_pairs+1] = { key, value }
    end

    table.sort(key_value_pairs, compare_key_value_pairs)

    return key_value_pairs
end

local function fill_holes_in_key_value_pairs (key_value_pairs)
    -- NOTE: Assumes that all keys are numbers

    -- Error checking
    assert(type(key_value_pairs) == 'table')

    -- Do stuff
    for i = 2, #key_value_pairs do
        for j = key_value_pairs[i-1][1] + 1, key_value_pairs[i][1] - 1 do
            key_value_pairs[#key_value_pairs+1] = { j, nil }
        end
    end
    table.sort(key_value_pairs, compare_key_value_pairs)
end

local function smallest_secure_longform_string_level (str)
    -- Determines the level a longform string needs to use, to avoid "code"
    -- injection. For example, if we want to use longform on the string
    -- 'Hello ]] World', we cannot use level-0 as this would result in
    -- '[[Hello ]] World]]', which could be an issue in certain applications.

    -- Error checking
    assert(type(str) == 'string')

    -- Do stuff
    local levels = { [1] = 1 }
    str:gsub('%]=*%]', function (m) levels[m:len()] = true end)
    return #levels - 1
end

local function escape_string (str)

    -- Error checking
    assert(type(str) == 'string')

    -- Do stuff
    local l = {}
    for i = 1, #str do
        l[#l+1] = CHAR_TO_STR_REPR[str:byte(i)]
    end
    return table.concat(l, '')
end

local function width_of_strings_in_l (l, start_i, stop_i)

    -- Argument fixing and Error Checking
    assert(type(l) == 'table')

    local start_i, stop_i = start_i or 1, stop_i or #l

    assert(type(start_i) == 'number')
    assert(type(start_i) == 'number')

    -- Do stuff
    local  width = 0
    for i = start_i, stop_i do
        width = width + ((type(l[i]) ~= 'string') and 1 or #l[i])
    end
    return width
end

local function ignore_alignment_info (l, start_i, stop_i)

    -- Argument fixing and Error Checking
    assert(type(l) == 'table')

    local start_i, stop_i = start_i or 1, stop_i or #l

    assert(type(start_i) == 'number')
    assert(type(start_i) == 'number')

    -- Do stuff
    for i = start_i, stop_i do
        if type(l[i]) == 'table' and l[i][1] == 'align' then
            l[i] = ''
        end
    end
end

local function fix_alignment (l, start_i, stop_i)

    -- Argument fixing and Error Checking
    assert(type(l) == 'table')

    local start_i, stop_i = start_i or 1, stop_i or #l

    assert(type(start_i) == 'number')
    assert(type(start_i) == 'number')
    -- Do stuff

    -- Find maximums
    local max  =  {}
    for i = start_i, stop_i do
        if type(l[i]) == 'table' and l[i][1] == 'align' then
            max[ l[i][2] ]  =  math.max( l[i][3], max[ l[i][2] ] or 0 )
        end
    end

    -- Insert the proper whitespace
    for i = start_i, stop_i do
        if type(l[i]) == 'table' and l[i][1] == 'align' then
            l[i] = string.rep(' ', max[ l[i][2] ] - l[i][3])
        end
    end
end

local function fix_seperator_info (l, indent_char, max_depth)

    -- Error Checking
    assert(type(l) == 'table')
    assert(type(indent_char) == 'string')
    assert(type(max_depth) == 'number')
    -- Do stuff

    local depth, inline_depth = 0, nil
    for i = 1, #l do
        if type(l[i]) ~= 'table' then
            -- Do nothing
        elseif l[i][1] == 'seperator' then
            assert(l[i][2] == nil or type(l[i][2]) == 'string')
            l[i]  =  (l[i][2] or '') .. (inline_depth and ' ' or ('\n' .. indent_char:rep(depth)))
        elseif l[i][1] == 'indent' then
            depth, inline_depth  =  depth + 1, inline_depth or l[i][3] == 'inline' and depth + 1 or nil
            l[i] = l[i][2] .. (inline_depth and ' ' or ('\n' .. indent_char:rep(depth)))
        elseif l[i][1] == 'unindent' then
            l[i] = (inline_depth and ' ' or ('\n' .. indent_char:rep(depth-1))) .. l[i][2]
            depth, inline_depth  =  depth - 1, (depth ~= inline_depth) and inline_depth or nil
        end
    end
end

--------------------------------------------------------------------------------
-- Formatting stuff

local format_value

-- Ways to format keys

local function format_key_and_value_string_map (l, key, value, depth)
    l[#l+1] = key
    l[#l+1] = { 'align', 'key', #key }
    l[#l+1] = ' = '
    return format_value(value, depth, l)
end

local function format_key_and_value_arbitr_map (l, key, value, depth)
    local index_before_key = #l+1
    l[#l+1] = '['
    format_value(key, math.huge, l)
    l[#l+1] = ']'
    l[#l+1] = { 'align', 'key', width_of_strings_in_l(l, index_before_key) }
    l[#l+1] = ' = '
    return format_value(value, depth, l)
end

local function format_key_and_value_sequence (l, key, value, depth)
    return format_value(value, depth, l)
end

local TABLE_TYPE_TO_PAIR_FORMAT = {
    [TABLE_TYPE.EMPTY]      = format_key_and_value_sequence,
    [TABLE_TYPE.SEQUENCE]   = format_key_and_value_sequence,
    [TABLE_TYPE.SET]        = format_key_and_value_arbitr_map,
    [TABLE_TYPE.MIXED]      = format_key_and_value_arbitr_map,
    [TABLE_TYPE.STRING_MAP] = format_key_and_value_string_map,
    [TABLE_TYPE.PURE_MAP]   = format_key_and_value_arbitr_map,
}

-- Formatting tables

local function format_table (t, depth, l)
    -- TODO: Add more nuanced formatting.

    -- Error Checking
    assert(type(t) == 'table')
    assert(type(depth) == 'number' and type(l) == 'table')

    -- Do stuff
    local table_info  =  l.info[t] or get_table_info(t)

    if l.options.recursion == 'marked' and table_info.marker then
        l[#l+1], l[#l+2], l[#l+3] = '<', table_info.marker, '>'
    end

    local already_visited  =  l.visited[t]
    l.visited[t]  =  true

    if table_info.type == TABLE_TYPE.EMPTY then
        -- Empty Map
        return l('{'..( l.options._table_addr_comment and (' --[['..table_info.address..']] ') or '')..'}')
    elseif depth >= l.options.max_depth or already_visited then
        -- Already visited or above max depth
        return l('{'..( l.options._table_addr_comment and (' --[['..table_info.address..']] ') or '')..'...}')
    end

    -- Get key-value pairs, and possibly fill holes.
    local key_value_pairs  = get_key_value_pairs_in_proper_order(t)
    if table_info.type == TABLE_TYPE.SEQUENCE and l.info[t].has_holes then
        fill_holes_in_key_value_pairs(key_value_pairs)
    end

    -- Find format function
    local pair_format_func = TABLE_TYPE_TO_PAIR_FORMAT[table_info.type]
    local start_of_table_i = #l + 1
    assert(pair_format_func)

    -- Begin formatting table.
    l[#l+1] = {'indent', '{'}
    if l.options._table_addr_comment then  l[#l+1], l[#l+2]  =  '--[['..table_info.address..']]', {'seperator'}  end
    for _, pair in ipairs(key_value_pairs) do
        pair_format_func(l, pair[1], pair[2], depth + 1)
        l[#l+1] = {'seperator', ','}
    end
    if l[#l][1] == 'seperator' then  l[#l] = nil  end
    l[#l+1] = {'unindent', '}'}

    --require 'fun' ()   -- DEBUG!
    --DEBUG = map(operator.identity, l) -- DEBUG!

    -- Decide for short or long table formatting.
    local table_width  =  width_of_strings_in_l(l, start_of_table_i)
    if table_width <= MAX_WIDTH_FOR_SINGLE_LINE_TABLE then
        -- Is short table: Ignore the "width of key"-shit
        l[start_of_table_i][3] = 'inline'
        ignore_alignment_info(l, start_of_table_i)
    else
        -- Is long table: Fix whitespace alignment
        fix_alignment(l, start_of_table_i)
    end
end

-- Formatting Strings

local function format_string (str, depth, l)
    -- TODO: Add option for escaping unicode characters.
    -- TODO: Improve cutstring argument.

    -- Error checking
    assert( type(str) == 'string' )
    assert(type(depth) == 'number' and type(l) == 'table')

    -- Do work

    local is_long_string         =  (str:len() >= NR_CHARS_IN_LONG_STRING)
    local newline_or_tab_index   =  str:find('[\n\t]')
    local single_quote_index     =  str:find('\'')
    local double_quote_index     =  str:find('\"')

    -- ...
    local chance_of_longform     =  is_long_string and ((newline_or_tab_index or math.huge) <= NR_CHARS_IN_LONG_STRING) or double_quote_index and single_quote_index
    local cut_string_index       =  l.options.cut_strings and (is_long_string or chance_of_longform)
                                                          and math.min(NR_CHARS_IN_LONG_STRING - 3, newline_or_tab_index or 1/0, double_quote_index or 1/0, single_quote_index or 1/0)

    local longform   = chance_of_longform and ((not cut_string_index) or cut_string_index < math.min(newline_or_tab_index or 1/0, double_quote_index or 1/0, single_quote_index or 1/0))

    local escape_newline_and_tab  =  not longform and newline_or_tab_index

    -- Determine string delimiters
    local left, right
    if longform then
        local level = smallest_secure_longform_string_level(str)
        left, right = '['..string.rep('=', level)..'[', ']'..string.rep('=', level)..']'
        if newline_or_tab_index then  str = '\n' .. str  end
    elseif not single_quote_index then
        left, right = '\'', '\''
    else
        left, right = '\"', '\"'
    end

    -- Cut string
    if cut_string_index then        str = str:sub(1, cut_string_index)  end
    str = escape_string(str)
    -- Escape newline and tab
    if escape_newline_and_tab then  str = str:gsub('\n', '\\n'):gsub('\t', '\\t')  end

    l[#l+1] = left
    l[#l+1] = str
    l[#l+1] = right
end

local function format_coroutine (value, depth, l)
    -- Error check
    assert(type(value) == 'thread')
    assert(type(depth) == 'number' and type(l) == 'table')
    -- Do stuff
    l[#l+1] = coroutine.status(value)
    l[#l+1] = ' coroutine: '
    l[#l+1] = tostring(value):sub(9)
end

local function format_primitive (value, depth, l)
    -- Error check
    assert(type(depth) == 'number' and type(l) == 'table')
    -- Do stuff
    l[#l+1] = tostring(value)
end

local TYPE_TO_FORMAT_FUNC = {
    ['nil']      =  format_primitive,
    ['boolean']  =  format_primitive,
    ['number']   =  format_number,
    ['string']   =  format_string,
    ['thread']   =  format_coroutine,
    ['table']    =  format_table,

    ['function'] =  format_function,   -- TODO: Improve a little
    ['userdata'] =  format_primitive,  -- TODO
    ['cdata']    =  format_primitive,  -- TODO & Luajit only
}

function format_value (value, depth, l)
    assert(type(depth) == 'number' and type(l) == 'table')
    local formatting = TYPE_TO_FORMAT_FUNC[type(value)]
    if formatting then
        formatting(value, depth, l, format_value)
    else
        error(ERROR_UNKNOWN_TYPE:format(type(value), tostring(value)), 2)
    end
end

--------------------------------------------------------------------------------
-- StringBuilder

local StringBuilder = {}
      StringBuilder.__index = StringBuilder
      StringBuilder.__call  = function (l, s) l[#l+1] = s end
setmetatable(StringBuilder, {
    __call = function() return setmetatable({}, StringBuilder) end
})

--------------------------------------------------------------------------------


local DEBUG_OPTIONS = { _all_function_info = true, _table_addr_comment = true }

local KNOWN_OPTIONS = {
    _table_addr_comment = { type = 'boolean', default = false },
    _all_function_info  = { type = 'boolean', default = false },
    cut_strings         = { type = 'boolean', default = false },
    include_closure     = { type = 'boolean', default = false },
    indent              = { type = 'string',  default = '    ' },
    max_depth           = { type = 'number',  default = math.huge },
    embed_loaded_funcs  = { type = 'boolean', default = false },
    more_function_info  = { type = 'boolean', default = false },
    recursion           = { type = 'string',  default = 'ignore', accepted = {['ignore'] = true, ['marked'] = true} },
    short_builtins      = { type = 'boolean', default = false },
}

local function ensure_that_all_options_are_known (options)
    assert(type(options) == 'table')
    -- Error check options
    for option_name, option_value in pairs(options) do
        if not KNOWN_OPTIONS[option_name] then
            error(('[pretty]: Unknown option: %s. Was given value %s (%s)'):format(option_name, option_value, type(option_value)), 2)
        elseif type(option_value) ~= KNOWN_OPTIONS[option_name].type then
            error(('[pretty]: Bad value given to option %s: %s (%s). Expected value of type %s'):format(option_name, option_value, type(option_value), KNOWN_OPTIONS[option_name].type), 2)
        elseif KNOWN_OPTIONS[option_name].accepted and not KNOWN_OPTIONS[option_name].accepted[option_value] then
            error(('[pretty]: Bad value given to option %s: %s (%s). Expected one of: %s'):format(option_name, option_value, type(option_value), table.concat(KNOWN_OPTIONS[option_name].accepted, ', ')), 2)
        end
    end
    -- Assign default values
    for option_name, option_info in pairs(KNOWN_OPTIONS) do
        if options[option_name] == nil then
            options[option_name] = option_info.default
        end
    end
    -- Returns options
    return options
end

local function pretty_format (value, options)
    -- Error checking
    local options  =  ensure_that_all_options_are_known(options or {})

    -- Setup StringBuilder
    local l = StringBuilder()
          l.visited = { next_mark = 1 }
          l.options = options
          l.info    = analyze_structure(value, options.max_depth)

    -- Format value.
    format_value(value, 0, l)

    -- If any alignment info still exists, ignore it
    fix_seperator_info(l, l.options.indent, l.options.max_depth)
    ignore_alignment_info(l)

    return table.concat(l, '')
end

return pretty_format