assert-gooder/Parser.lua

require 'errors' 'shunt' . enable_strict_globals ()

---- Algorithm

local DEFAULT_ASSOC = 'left'
local DEFAULT_ARITY = 2

local function greater_than (a, b)
    assert(type(a) == 'number', a)
    assert(type(b) == 'number', b)

    return a > b
end

local function greater_than_or_equal (a, b)
    assert(type(a) == 'number', a)
    assert(type(b) == 'number', b)

    return a >= b
end

local function shunting_yard (tokens, lang, num_arity_to_token)
    -- Implementation of the shunting yard algorithm, for transforming
    -- infix notation math to postfix notation math.
    --
    -- Features extension for variadic first-class functions, which
    -- results in slightly complicated output.
    -- Whenever a function is called, the output will be formatted
    -- like: <function-identifier> <arguments...> <num-args>
    -- <call-operator>
    --
    -- For example: f(1, 2, 3)
    -- Output:      "f" 1 2 3 3 $
    --
    -- The exact behavior can be customized.
    --
    assert(type(tokens) == 'table' and #tokens > 0)
    assert(type(lang)   == 'function')

    local pop_i = 1
    local stack, output, call_stack  =  {}, {}, {}
    local prev_info = nil

    for _, val in ipairs(tokens) do
        local t_info  =  lang(val, prev_info)
        assert(type(t_info) == 'table' and t_info.val ~= nil)

        if t_info.type == 'imm' then
            output[#output+1] = t_info

        -- Binary operators
        elseif t_info.type == 'op' then
            -- Determine associativity
            t_info.associativity = t_info.associativity or DEFAULT_ASSOC
            t_info.arity = t_info.arity or DEFAULT_ARITY
            assert(type(t_info.precedence) == 'number')
            assert(t_info.associativity == 'left' or t_info.associativity == 'right')
            assert(t_info.arity         == 1      or t_info.arity         == 2)
            local test = t_info.associativity == 'left' and greater_than_or_equal or greater_than

            -- Pop operations with higher precedence from the stack
            if t_info.arity == 1 and t_info.associativity == 'right' then
                -- Do not pop operations from the stack
            elseif t_info.arity == 2 then
                while #stack > 0 and stack[#stack].type == 'op' and test(stack[#stack].precedence, t_info.precedence) do
                    output[#output+1], stack[#stack] = stack[#stack], nil
                end
            else
                assert(false)
            end

            -- Add this operation to the stack
            stack[#stack+1] = t_info

        -- Enclosures
        elseif t_info.type == 'bracket' then
            -- Determine opening and closing brackets
            assert(t_info.open ~= nil and t_info.close ~= nil and t_info.open ~= t_info.close)

            -- Are we opening or closing?
            if t_info.val == t_info.open then
                while #stack > 0 and stack[#stack].type == 'op' and stack[#stack].precedence >= t_info.precedence do
                    output[#output+1], stack[#stack] = stack[#stack], nil
                end
                if t_info.on_end_call_op and prev_info and prev_info.type ~= 'op' then
                    call_stack[#call_stack+1] = { arity = 1, bracket = t_info }
                    stack[#stack+1] = t_info.on_end_call_op
                end
                stack[#stack+1] = t_info

            elseif t_info.val == t_info.close then

                -- If closing, pop from stack into output
                while #stack > 0 and stack[#stack].val ~= t_info.open do
                    output[#output+1], stack[#stack] = stack[#stack], nil
                end

                -- Remove remaining bracket.
                -- Brackets are not needed in postfix notation.
                assert(#stack > 0 and stack[#stack].val == t_info.open)
                stack[#stack] = nil

                -- If top of stack is a function, we should pop it
                if #stack > 0 and stack[#stack].call_op then
                    assert(#call_stack > 0)
                    if prev_info.type == 'bracket' and prev_info.open == prev_info.val and t_info.open == prev_info.val then
                        -- No arguments for this function
                        call_stack[#call_stack].arity = 0
                    end

                    -- Pop from stacks and construct output
                    output[#output+1], call_stack[#call_stack]  =  num_arity_to_token(call_stack[#call_stack].arity), nil
                    output[#output+1], stack[#stack] = stack[#stack], nil
                end
            else
                assert(false, require'pretty'(t_info))
            end

        elseif t_info.type == 'arg-sep' then

            -- Ensure that we are in a call
            assert(#call_stack > 0 and call_stack[#call_stack].bracket.open == t_info.open)

            -- Pop from stack
            while #stack > 0 and stack[#stack].val ~= t_info.open do
                output[#output+1], stack[#stack] = stack[#stack], nil
            end

            -- Ensure open is on the stack
            assert(#stack > 0 and stack[#stack].val == t_info.open)

            -- Increment arity in call-stack
            call_stack[#call_stack].arity = call_stack[#call_stack].arity + 1

        else
            assert(false, 'Unknown type: '..tostring(t_info.type))
        end

        prev_info = t_info

    end

    -- Pop remaining from stack onto output
    while #stack > 0 do
        output[#output+1], stack[#stack] = stack[#stack], nil
    end

    -- Find value for all tokens on the stack
    for i = 1, #output do
        assert(output[i].val and output[i].type ~= 'bracket', require'pretty'(output[i]))
        output[i] = output[i].val
    end

    -- Post asserts
    --assert(type(output) == 'table' and #output <= #tokens)
    assert(#stack == 0)
    --
    return output
end

return shunting_yard