pretty/pretty.lua

-- pretty.pretty
-- Main module of the `pretty` library.

-- TODO: Maybe move table formatting into its own file?


--[=[ Thoughts on displaying tables in an intuitive way.

Lua's table datastructure is likely to be the most consise data structure ever
invented. (If not, please send me a link!) Lists, maps, objects, classes,
proxies, etc. This obviously brings about it some difficulty when attempting to
represent these tables. What do we want to highlight, and what do we choose to
avoid?

One notable issue is whether to show every key that a table answers to, or to
just display those it contains. That is, do we think about `__index` and what it
returns, or do we ignore `__index`? For cases where `__index` is a function,
we cannot say anything about the keys that the table answers to. If `__index` is
a table, we have a better idea, but then the output would be cluttered.

1.  Native representation: Lua's native representation includes the type and
    address of the table. It allows for distinguishing between unique tables,
    but won't tell about the contents.
2.  It continues: By representing tables as the pseudo-parsable `{...}`, it's
    clear we are talking about a table. We disregard the ability to
    distinguish between tables.
3.  Special case for 2. If the table is empty, we could represent it as `{}`.
    But what if the table has a metatable with `__index` defined? We could
    continue to represent it as `{}`, but `{...}` would be more "honest".
4.  Single-line: TODO
5.  Multi-line: TODO
6.  Format into columns: (like ls) TODO
7.  Tabular: TODO
8.  Special cases: (Array-tree, Table-Tree, Linked-List) TODO

--]=]

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

-- 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, 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'
    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)
    -- Compares two strings alphanumerically.

    assert(type(a) == 'string')
    assert(type(b) == 'string')

    local a_padded  =  a:gsub("%.?%d+", padnum)..("\0%3d"):format(#b)
    local b_padded  =  b:gsub("%.?%d+", padnum)..("\0%3d"):format(#a)
    local A_padded, B_padded  =  a_padded:upper(), b_padded:upper()
    if A_padded == B_padded then  return  a_padded < b_padded  end
    return A_padded < B_padded
end

local function compare_key_value_pair (a, b)
    -- Function for comparing two key-value pairs, given as `{ key, value }`.
    -- Pretty complex due to our high standards for sorting.

    assert(type(a) == 'table')
    assert(type(b) == 'table')

    -- 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_pair)

    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_pair)
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')

    -- 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 attempt_to_align_into_columns (l, start_i, stop_i, nr_items_pr_row)
    assert(type(l) == 'table')
    assert(type(start_i) == 'number')
    assert(type(stop_i) == 'number')
    assert(type(nr_items_pr_row) == 'number')

	local column = {}
    ---
    local start_of_item_i, item_nr  =  nil, 0
	for i = start_i, stop_i do
        if type(l[i]) == 'table' and (l[i][1] == 'indent' or l[i][1] == 'seperator' or l[i][1] == 'unindent') then
            if start_of_item_i then
                local width_of_item = width_of_strings_in_l(l, start_of_item_i, i-1)
                local column_i = (item_nr-1)%nr_items_pr_row+1
                column[column_i]  =  math.max(column[column_i] or 0, width_of_item)
            end
            start_of_item_i, item_nr  =  i + 1, item_nr + 1
        end
	end
    ---
    local width  =  nr_items_pr_row * 2 - 1  -- FIXME: Magic numbers: 2 = #', ', 1 = #' '
    for i = 1, #column do  width = width + column[i]  end
    --
	return width, column
end

local function align_into_columns (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')

    -- Find columns
    local columns = nil
    for nr_items_pr_row = 10, 1, -1 do  -- TODO: Do this more intelligently.
        local column_width
        column_width, columns  =  attempt_to_align_into_columns(l, start_i, stop_i, nr_items_pr_row)
        if column_width <= MAX_WIDTH_FOR_SINGLE_LINE_TABLE then  break  end
    end

    -- Change alignment of columns
    local start_of_item_i, item_nr  =  nil, 0
    for i = start_i, stop_i do
        if type(l[i]) == 'table' and l[i][1] == 'align' then
            local column_i = (item_nr-1)%#columns+1
            l[i][2] = l[i][2] .. '_column_'..column_i
        elseif (l[i][1] == 'indent' or l[i][1] == 'seperator' or l[i][1] == 'unindent') then
            start_of_item_i, item_nr  =  i + 1, item_nr + 1
        end
    end

    -- Fix newly changed alignment
    fix_alignment(l, start_i, stop_i)

    -- Quick-exit on only a single column
    if #columns == 1 then  return  end

    -- Fit into columns.
    local start_of_item_i, item_nr  =  nil, 0
	for i = start_i, stop_i do
        if type(l[i]) ~= 'table' then
            -- Do nothing
        elseif (l[i][1] == 'indent' or l[i][1] == 'seperator' or l[i][1] == 'unindent') then
            if start_of_item_i and l[i][1] == 'seperator' then
                local column_i = (item_nr-1)%#columns+1
                if column_i ~= #columns then
                    local width_of_item = width_of_strings_in_l(l, start_of_item_i, i-1)
                    l[i] = l[i][2] .. ' ' .. (' '):rep(columns[column_i]-width_of_item)
                end
            end
            start_of_item_i, item_nr  =  i + 1, item_nr + 1
        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
    if not l.info[t] then analyze_structure(t, l.options.max_depth-depth, l.info)  end
    local table_info  =  l.info[t]
    assert(table_info)

    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 empty, visited or above max-depth, give a small represetation: `{...}`
    if table_info.type == TABLE_TYPE.EMPTY or depth >= l.options.max_depth or already_visited then
        l '{'
        if l.options._table_addr_comment       then  l[#l+1] = ' --[[' .. table_info.address .. ']] '  end
        if table_info.type ~= TABLE_TYPE.EMPTY then l[#l+1] =  '...'  end
        return l '}'
    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 pair formatting 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', '}'}

    -- 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)
    elseif table_info.is_leaf_node then
        -- Is leaf node: Can format into columns.
        -- NOTE: Currently we only allow leaf-nodes to format into columns, due
        -- to issues with table alignment.
        align_into_columns(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.
    -- TODO: Move this to it's own file.

    -- 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)
    -- Formats a coroutine. Unfortunantly we cannot gather a lot of information
    -- about coroutines.

    -- 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_OPTION_USED = {  }

local KNOWN_OPTIONS = {
    _table_addr_comment = { type = 'boolean', default = false, debug = 'debug' },

    cut_strings         = { type = 'boolean', default = false },
    indent              = { type = 'string',  default = '    ' },
    max_depth           = { type = 'number',  default = math.huge },
    short_builtins      = { type = 'boolean', default = false }, -- TODO: Outphase this. Rather automatically use the short versions in places where it would be strange to find the function, like keys, etc.
    recursion           = { type = 'string',  default = 'ignore', accepted = {['ignore'] = true, ['marked'] = true} },
}

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)
        elseif KNOWN_OPTIONS[option_name].debug and not DEBUG_OPTION_USED[option_name] then
            DEBUG_OPTION_USED[option_name] = true
            print(('[pretty]: Using %s option "%s".\n    Please note that this option may change at any time. It is not stable,\n    not tested, and may indeed break or be removed without warning.'):format(KNOWN_OPTIONS[option_name].debug, option_name))
        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