lexer.lua

local ffi = require('ffi')

local band = bit.band
local strsub, strbyte, strchar = string.sub, string.byte, string.char

local ASCII_0, ASCII_9 = 48, 57
local ASCII_a, ASCII_f, ASCII_z = 97, 102, 122
local ASCII_A, ASCII_Z = 65, 90

local END_OF_STREAM = -1

local ReservedKeyword = {['and'] = 1, ['break'] = 2, ['do'] = 3, ['else'] = 4, ['elseif'] = 5, ['end'] = 6, ['false'] = 7, ['for'] = 8, ['function'] = 9, ['goto'] = 10, ['if'] = 11, ['in'] = 12, ['local'] = 13, ['nil'] = 14, ['not'] = 15, ['or'] = 16, ['repeat'] = 17, ['return'] = 18, ['then'] = 19, ['true'] = 20, ['until'] = 21, ['while'] = 22 }

local uint64, int64 = ffi.typeof('uint64_t'), ffi.typeof('int64_t')
local complex = ffi.typeof('complex')

local TokenSymbol = { TK_ge = '>=', TK_le = '<=' , TK_concat = '..', TK_eq = '==', TK_ne = '~=', TK_eof = '<eof>' }

local function token2str(tok)
    if string.match(tok, "^TK_") then
        return TokenSymbol[tok] or string.sub(tok, 4)
    else
        return tok
    end
end

local function error_lex(chunkname, tok, line, em, ...)
    local emfmt = string.format(em, ...)
    local msg = string.format("%s:%d: %s", chunkname, line, emfmt)
    if tok then
        msg = string.format("%s near '%s'", msg, tok)
    end
    error("LLT-ERROR" .. msg, 0)
end

local function lex_error(ls, token, em, ...)
    local tok
    if token == 'TK_name' or token == 'TK_string' or token == 'TK_number' then
        tok = ls.save_buf
    elseif token then
        tok = token2str(token)
    end
    error_lex(ls.chunkname, tok, ls.linenumber, em, ...)
end

local function char_isident(c)
    if type(c) == 'string' then
        local b = strbyte(c)
        if b >= ASCII_0 and b <= ASCII_9 then
            return true
        elseif b >= ASCII_a and b <= ASCII_z then
            return true
        elseif b >= ASCII_A and b <= ASCII_Z then
            return true
        else
            return (c == '_')
        end
    end
    return false
end

local function char_isdigit(c)
    if type(c) == 'string' then
        local b = strbyte(c)
        return b >= ASCII_0 and b <= ASCII_9
    end
    return false
end

local function char_isspace(c)
    local b = strbyte(c)
    return b >= 9 and b <= 13 or b == 32
end

local function byte(ls, n)
    local k = ls.p + n
    return strsub(ls.data, k, k)
end

local function skip(ls, n)
    ls.n = ls.n - n
    ls.p = ls.p + n
end

local function pop(ls)
    local k = ls.p
    local c = strsub(ls.data, k, k)
    ls.p = k + 1
    ls.n = ls.n - 1
    return c
end

local function fillbuf(ls)
    local data = ls:read_func()
    if not data then
        return END_OF_STREAM
    end
    ls.data, ls.n, ls.p = data, #data, 1
    return pop(ls)
end

local function nextchar(ls)
    local c = ls.n > 0 and pop(ls) or fillbuf(ls)
    ls.current = c
    return c
end

local function curr_is_newline(ls)
    local c = ls.current
    return (c == '\n' or c == '\r')
end

local function resetbuf(ls)
    ls.save_buf = ''
end

local function resetbuf_tospace(ls)
    ls.space_buf = ls.space_buf .. ls.save_buf
    ls.save_buf = ''
end

local function spaceadd(ls, str)
    ls.space_buf = ls.space_buf .. str
end

local function save(ls, c)
    ls.save_buf = ls.save_buf .. c
end

local function savespace_and_next(ls)
    ls.space_buf = ls.space_buf .. ls.current
    nextchar(ls)
end

local function save_and_next(ls)
    ls.save_buf = ls.save_buf .. ls.current
    nextchar(ls)
end

local function get_string(ls, init_skip, end_skip)
    return strsub(ls.save_buf, init_skip + 1, - (end_skip + 1))
end

local function get_space_string(ls)
    local s = ls.space_buf
    ls.space_buf = ''
    return s
end

local function inclinenumber(ls)
    local old = ls.current
    savespace_and_next(ls) -- skip `\n' or `\r'
    if curr_is_newline(ls) and ls.current ~= old then
        savespace_and_next(ls) -- skip `\n\r' or `\r\n'
    end
    ls.linenumber = ls.linenumber + 1
end

local function skip_sep(ls)
    local count = 0
    local s = ls.current
    assert(s == '[' or s == ']')
    save_and_next(ls)
    while ls.current == '=' do
        save_and_next(ls)
        count = count + 1
    end
    return ls.current == s and count or (-count - 1)
end

local function build_64int(str)
    local u = str[#str - 2]
    local x = (u == 117 and uint64(0) or int64(0))
    local i = 1
    while str[i] >= ASCII_0 and str[i] <= ASCII_9 do
        x = 10 * x + (str[i] - ASCII_0)
        i = i + 1
    end
    return x
end

-- Only lower case letters are accepted.
local function byte_to_hexdigit(b)
    if b >= ASCII_0 and b <= ASCII_9 then
        return b - ASCII_0
    elseif b >= ASCII_a and b <= ASCII_f then
        return 10 + (b - ASCII_a)
    else
        return -1
    end
end

local function build_64hex(str)
    local u = str[#str - 2]
    local x = (u == 117 and uint64(0) or int64(0))
    local i = 3
    while str[i] do
        local n = byte_to_hexdigit(str[i])
        if n < 0 then break end
        x = 16 * x + n
        i = i + 1
    end
    return x
end

local function strnumdump(str)
    local t = {}
    for i = 1, #str do
        local c = strsub(str, i, i)
        if char_isident(c) then
            t[i] = strbyte(c)
        else
            return nil
        end
    end
    return t
end

local function lex_number(ls)
    local lower = string.lower
    local xp = 'e'
    local c = ls.current
    if c == '0' then
        save_and_next(ls)
        local xc = ls.current
        if xc == 'x' or xc == 'X' then xp = 'p' end
    end
    while char_isident(ls.current) or ls.current == '.' or
        ((ls.current == '-' or ls.current == '+') and lower(c) == xp) do
        c = lower(ls.current)
        save(ls, c)
        nextchar(ls)
    end
    local str = ls.save_buf
    local x
    if strsub(str, -1, -1) == 'i' then
        local img = tonumber(strsub(str, 1, -2))
        if img then x = complex(0, img) end
    elseif strsub(str, -2, -1) == 'll' then
        local t = strnumdump(str)
        if t then
            x = xp == 'e' and build_64int(t) or build_64hex(t)
        end
    else
        x = tonumber(str)
    end
    if x then
        return x
    else
        lex_error(ls, 'TK_number', "malformed number")
    end
end

local function read_long_string(ls, sep, ret_value)
    save_and_next(ls) -- skip 2nd `['
    if curr_is_newline(ls) then -- string starts with a newline?
        inclinenumber(ls) -- skip it
    end
    while true do
        local c = ls.current
        if c == END_OF_STREAM then
            lex_error(ls, 'TK_eof', ret_value and "unfinished long string" or "unfinished long comment")
        elseif c == ']' then
            if skip_sep(ls) == sep then
                save_and_next(ls) -- skip 2nd `['
                break
            end
        elseif c == '\n' or c == '\r' then
            save(ls, '\n')
            inclinenumber(ls)
            if not ret_value then
                resetbuf(ls) -- avoid wasting space
            end
        else
            if ret_value then save_and_next(ls)
            else nextchar(ls) end
        end
    end
    if ret_value then
        return get_string(ls, 2 + sep, 2 + sep)
    end
end

local Escapes = {
    a = '\a', b = '\b', f = '\f', n = '\n', r = '\r', t = '\t',
    v = '\v',
}

local function hex_char(c)
    if string.match(c, '^%x') then
        local b = band(strbyte(c), 15)
        if not char_isdigit(c) then b = b + 9 end
        return b
    end
end

local function read_escape_char(ls)
    local c = nextchar(ls) -- Skip the '\\'.
    local esc = Escapes[c]
    if esc then
        save(ls, esc)
        nextchar(ls)
    elseif c == 'x' then -- Hexadecimal escape '\xXX'.
        local ch1 = hex_char(nextchar(ls))
        local hc
        if ch1 then
            local ch2 = hex_char(nextchar(ls))
            if ch2 then
                hc = strchar(ch1 * 16 + ch2)
            end
        end
        if not hc then
            lex_error(ls, 'TK_string', "invalid escape sequence")
        end
        save(ls, hc)
        nextchar(ls)
    elseif c == 'z' then -- Skip whitespace.
        nextchar(ls)
        while char_isspace(ls.current) do
            if curr_is_newline(ls) then inclinenumber(ls) else nextchar(ls) end
        end
    elseif c == '\n' or c == '\r' then
        save(ls, '\n')
        inclinenumber(ls)
    elseif c == '\\' or c == '\"' or c == '\'' then
        save(ls, c)
        nextchar(ls)
    elseif c == END_OF_STREAM then
    else
        if not char_isdigit(c) then
            lex_error(ls, 'TK_string', "invalid escape sequence")
        end
        local bc = band(strbyte(c), 15) -- Decimal escape '\ddd'.
        if char_isdigit(nextchar(ls)) then
            bc = bc * 10 + band(strbyte(ls.current), 15)
            if char_isdigit(nextchar(ls)) then
                bc = bc * 10 + band(strbyte(ls.current), 15)
                if bc > 255 then
                    lex_error(ls, 'TK_string', "invalid escape sequence")
                end
                nextchar(ls)
            end
        end
        save(ls, strchar(bc))
    end
end

local function read_string(ls, delim)
    save_and_next(ls)
    while ls.current ~= delim do
        local c = ls.current
        if c == END_OF_STREAM then
            lex_error(ls, 'TK_eof', "unfinished string")
        elseif c == '\n' or c == '\r' then
            lex_error(ls, 'TK_string', "unfinished string")
        elseif c == '\\' then
            read_escape_char(ls)
        else
            save_and_next(ls)
        end
    end
    save_and_next(ls) -- skip delimiter
    return get_string(ls, 1, 1)
end

local function skip_line(ls)
    while not curr_is_newline(ls) and ls.current ~= END_OF_STREAM do
        savespace_and_next(ls)
    end
end

local function llex(ls)
    resetbuf(ls)
    while true do
        local current = ls.current
        if char_isident(current) then
            if char_isdigit(current) then -- Numeric literal.
                return 'TK_number', lex_number(ls)
            end
            repeat
                save_and_next(ls)
            until not char_isident(ls.current)
            local s = get_string(ls, 0, 0)
            local reserved = ReservedKeyword[s]
            if reserved then
                return 'TK_' .. s
            else
                return 'TK_name', s
            end
        end
        if current == '\n' or current == '\r' then
            inclinenumber(ls)
        elseif current == ' ' or current == '\t' or current == '\b' or current == '\f' then
            savespace_and_next(ls)
            -- nextchar(ls)
        elseif current == '-' then
            nextchar(ls)
            if ls.current ~= '-' then return '-' end
            -- else is a comment
            nextchar(ls)
            spaceadd(ls, '--')
            if ls.current == '[' then
                local sep = skip_sep(ls)
                resetbuf_tospace(ls) -- `skip_sep' may dirty the buffer
                if sep >= 0 then
                    read_long_string(ls, sep, false) -- long comment
                    resetbuf_tospace(ls)
                else
                    skip_line(ls)
                end
            else
                skip_line(ls)
            end
        elseif current == '[' then
            local sep = skip_sep(ls)
            if sep >= 0 then
                local str = read_long_string(ls, sep, true)
                return 'TK_string', str
            elseif sep == -1 then
                return '['
            else
                lex_error(ls, 'TK_string', "delimiter error")
            end
        elseif current == '=' then
            nextchar(ls)
            if ls.current ~= '=' then return '=' else nextchar(ls); return 'TK_eq' end
        elseif current == '<' then
            nextchar(ls)
            if ls.current ~= '=' then return '<' else nextchar(ls); return 'TK_le' end
        elseif current == '>' then
            nextchar(ls)
            if ls.current ~= '=' then return '>' else nextchar(ls); return 'TK_ge' end
        elseif current == '~' then
            nextchar(ls)
            if ls.current ~= '=' then return '~' else nextchar(ls); return 'TK_ne' end
        elseif current == '*' then
            nextchar(ls)
            if ls.current ~= '*' then return '*' else nextchar(ls); return '**' end
        elseif current == '^' then
            nextchar(ls)
            if ls.current ~= '^' then return '^' else nextchar(ls); return '^^' end
        elseif current == ':' then
            nextchar(ls)
            if ls.current ~= ':' then return ':' else nextchar(ls); return 'TK_label' end
        elseif current == '"' or current == "'" then
            local str = read_string(ls, current)
            return 'TK_string', str
        elseif current == '.' then
            save_and_next(ls)
            if ls.current == '.' then
                nextchar(ls)
                if ls.current == '.' then
                    nextchar(ls)
                    return 'TK_dots' -- ...
                end
                return 'TK_concat' -- ..
            elseif not char_isdigit(ls.current) then
                return '.'
            else
                return 'TK_number', lex_number(ls)
            end
        elseif current == END_OF_STREAM then
            return 'TK_eof'
        else
            nextchar(ls)
            return current -- Single-char tokens (+ - / ...).
        end
    end
end

local Lexer = {
    token2str = token2str,
    error = lex_error,
}

function Lexer.next(ls)
    ls.lastline = ls.linenumber
    if ls.tklookahead == 'TK_eof' then -- No lookahead token?
        ls.token, ls.tokenval = llex(ls) -- Get nextchar token.
        ls.space = get_space_string(ls)
    else
        ls.token, ls.tokenval = ls.tklookahead, ls.tklookaheadval
        ls.space = ls.spaceahead
        ls.tklookahead = 'TK_eof'
    end
end

function Lexer.lookahead(ls)
    assert(ls.tklookahead == 'TK_eof')
    ls.tklookahead, ls.tklookaheadval = llex(ls)
    ls.spaceahead = get_space_string(ls)
    return ls.tklookahead
end

local LexerClass = { __index = Lexer }

local function lex_setup(read_func, chunkname)
    local header = false
    local ls = {
        n = 0,
        tklookahead = 'TK_eof', -- No look-ahead token.
        linenumber = 1,
        lastline = 1,
        read_func = read_func,
        chunkname = chunkname,
        space_buf = ''
    }
    nextchar(ls)
    if ls.current == '\xef' and ls.n >= 2 and
        byte(ls, 0) == '\xbb' and byte(ls, 1) == '\xbf' then -- Skip UTF-8 BOM (if buffered).
        ls.n = ls.n - 2
        ls.p = ls.p + 2
        nextchar(ls)
        header = true
    end
    if ls.current == '#' then
        repeat
            nextchar(ls)
            if ls.current == END_OF_STREAM then return ls end
        until curr_is_newline(ls)
        inclinenumber(ls)
        header = true
    end
    return setmetatable(ls, LexerClass)
end

return lex_setup