lexer.lua

local utils = require "__zk-lib__.lualib.luacheck.utils"

-- Lexer should support syntax of Lua 5.1, Lua 5.2, Lua 5.3 and LuaJIT(64bit and complex cdata literals).
local lexer = {}

local sbyte = string.byte
local schar = string.char
local sreverse = string.reverse
local tconcat = table.concat
local mfloor = math.floor

-- No point in inlining these, fetching a constant ~= fetching a local.
local BYTE_0, BYTE_9, BYTE_f, BYTE_F = sbyte("0"), sbyte("9"), sbyte("f"), sbyte("F")
local BYTE_x, BYTE_X, BYTE_i, BYTE_I = sbyte("x"), sbyte("X"), sbyte("i"), sbyte("I")
local BYTE_l, BYTE_L, BYTE_u, BYTE_U = sbyte("l"), sbyte("L"), sbyte("u"), sbyte("U")
local BYTE_e, BYTE_E, BYTE_p, BYTE_P = sbyte("e"), sbyte("E"), sbyte("p"), sbyte("P")
local BYTE_a, BYTE_z, BYTE_A, BYTE_Z = sbyte("a"), sbyte("z"), sbyte("A"), sbyte("Z")
local BYTE_DOT, BYTE_COLON = sbyte("."), sbyte(":")
local BYTE_OBRACK, BYTE_CBRACK = sbyte("["), sbyte("]")
local BYTE_OBRACE, BYTE_CBRACE = sbyte("{"), sbyte("}")
local BYTE_QUOTE, BYTE_DQUOTE = sbyte("'"), sbyte('"')
local BYTE_PLUS, BYTE_DASH, BYTE_LDASH = sbyte("+"), sbyte("-"), sbyte("_")
local BYTE_SLASH, BYTE_BSLASH = sbyte("/"), sbyte("\\")
local BYTE_EQ, BYTE_NE = sbyte("="), sbyte("~")
local BYTE_LT, BYTE_GT = sbyte("<"), sbyte(">")
local BYTE_LF, BYTE_CR = sbyte("\n"), sbyte("\r")
local BYTE_SPACE, BYTE_FF, BYTE_TAB, BYTE_VTAB = sbyte(" "), sbyte("\f"), sbyte("\t"), sbyte("\v")

local function to_hex(b)
   if BYTE_0 <= b and b <= BYTE_9 then
      return b-BYTE_0
   elseif BYTE_a <= b and b <= BYTE_f then
      return 10+b-BYTE_a
   elseif BYTE_A <= b and b <= BYTE_F then
      return 10+b-BYTE_A
   else
      return nil
   end
end

local function to_dec(b)
   if BYTE_0 <= b and b <= BYTE_9 then
      return b-BYTE_0
   else
      return nil
   end
end

local function to_utf(codepoint)
   if codepoint < 0x80 then  -- ASCII?
      return schar(codepoint)
   end

   local buf = {}
   local mfb = 0x3F

   repeat
      buf[#buf+1] = schar(codepoint % 0x40 + 0x80)
      codepoint = mfloor(codepoint / 0x40)
      mfb = mfloor(mfb / 2)
   until codepoint <= mfb

   buf[#buf+1] = schar(0xFE - mfb*2 + codepoint)
   return sreverse(tconcat(buf))
end

local function is_alpha(b)
   return (BYTE_a <= b and b <= BYTE_z) or
      (BYTE_A <= b and b <= BYTE_Z) or b == BYTE_LDASH
end

local function is_newline(b)
   return (b == BYTE_LF) or (b == BYTE_CR)
end

local function is_space(b)
   return (b == BYTE_SPACE) or (b == BYTE_FF) or
      (b == BYTE_TAB) or (b == BYTE_VTAB)
end

local keywords = utils.array_to_set({
   "and", "break", "do", "else", "elseif", "end", "false", "for", "function", "goto", "if", "in",
   "local", "nil", "not", "or", "repeat", "return", "then", "true", "until", "while"})

local simple_escapes = {
   [sbyte("a")] = sbyte("\a"),
   [sbyte("b")] = sbyte("\b"),
   [sbyte("f")] = sbyte("\f"),
   [sbyte("n")] = sbyte("\n"),
   [sbyte("r")] = sbyte("\r"),
   [sbyte("t")] = sbyte("\t"),
   [sbyte("v")] = sbyte("\v"),
   [BYTE_BSLASH] = BYTE_BSLASH,
   [BYTE_QUOTE] = BYTE_QUOTE,
   [BYTE_DQUOTE] = BYTE_DQUOTE
}

local function next_byte(state)
   local offset = state.offset + 1
   state.offset = offset
   return state.src:get_codepoint(offset)
end

-- Skipping helpers.
-- Take the current character, skip something, return next character.

local function skip_newline(state, newline)
   local first_newline_offset = state.offset
   local b = next_byte(state)

   if b ~= newline and is_newline(b) then
      b = next_byte(state)
   end

   local line = state.line
   local line_offsets = state.line_offsets
   state.line_lengths[line] = first_newline_offset - line_offsets[line]
   line = line + 1
   state.line = line
   line_offsets[line] = state.offset
   return b
end

local function skip_to_newline(state, b)
   while not is_newline(b) and b do
      b = next_byte(state)
   end

   return b
end

local function skip_space(state, b)
   while is_space(b) or is_newline(b) do
      if is_newline(b) then
         b = skip_newline(state, b)
      else
         b = next_byte(state)
      end
   end

   return b
end

-- Skips "[=*" or "]=*". Returns next character and number of "="s.
local function skip_long_bracket(state)
   local start = state.offset
   local b = next_byte(state)

   while b == BYTE_EQ do
      b = next_byte(state)
   end

   return b, state.offset-start-1
end

-- Token handlers.

-- Called after the opening "[=*" has been skipped.
-- Takes number of "=" in the opening bracket and token type(comment or string).
local function lex_long_string(state, opening_long_bracket, token)
   local b = next_byte(state)

   if is_newline(b) then
      b = skip_newline(state, b)
   end

   local lines = {}
   local line_start = state.offset

   while true do
      if is_newline(b) then
         -- Add the finished line.
         lines[#lines+1] = state.src:get_substring(line_start, state.offset-1)

         b = skip_newline(state, b)
         line_start = state.offset
      elseif b == BYTE_CBRACK then
         local long_bracket
         b, long_bracket = skip_long_bracket(state)

         if b == BYTE_CBRACK and long_bracket == opening_long_bracket then
            break
         end
      elseif b == nil then
         return nil, token == "string" and "unfinished long string" or "unfinished long comment"
      else
         b = next_byte(state)
      end
   end

   -- Add last line.
   lines[#lines+1] = state.src:get_substring(line_start, state.offset-opening_long_bracket-2)
   state.offset = state.offset + 1
   return token, tconcat(lines, "\n")
end

local function lex_short_string(state, quote)
   local b = next_byte(state)
   local chunks  -- Buffer is only required when there are escape sequences.
   local chunk_start = state.offset

   while b ~= quote do
      if b == BYTE_BSLASH then
         -- Escape sequence.

         if not chunks then
            -- This is the first escape sequence, init buffer.
            chunks = {}
         end

         -- Put previous chunk into buffer.
         if chunk_start ~= state.offset then
            chunks[#chunks+1] = state.src:get_substring(chunk_start, state.offset-1)
         end

         b = next_byte(state)

         -- The final string escape sequence evaluates to.
         local s

         local escape_byte = simple_escapes[b]

         if escape_byte then  -- Is it a simple escape sequence?
            b = next_byte(state)
            s = schar(escape_byte)
         elseif is_newline(b) then
            b = skip_newline(state, b)
            s = "\n"
         elseif b == BYTE_x then
            -- Hexadecimal escape.
            b = next_byte(state)  -- Skip "x".
            -- Exactly two hexadecimal digits.
            local c1, c2

            if b then
               c1 = to_hex(b)
            end

            if not c1 then
               return nil, "invalid hexadecimal escape sequence", -2
            end

            b = next_byte(state)

            if b then
               c2 = to_hex(b)
            end

            if not c2 then
               return nil, "invalid hexadecimal escape sequence", -3
            end

            b = next_byte(state)
            s = schar(c1*16 + c2)
         elseif b == BYTE_u then
            b = next_byte(state)  -- Skip "u".

            if b ~= BYTE_OBRACE then
               return nil, "invalid UTF-8 escape sequence", -2
            end

            b = next_byte(state)  -- Skip "{".

            local codepoint  -- There should be at least one digit.

            if b then
               codepoint = to_hex(b)
            end

            if not codepoint then
               return nil, "invalid UTF-8 escape sequence", -3
            end

            local hexdigits = 0

            while true do
               b = next_byte(state)
               local hex

               if b then
                  hex = to_hex(b)
               end

               if hex then
                  hexdigits = hexdigits + 1
                  codepoint = codepoint*16 + hex

                  if codepoint > 0x10FFFF then
                     -- UTF-8 value too large.
                     return nil, "invalid UTF-8 escape sequence", -hexdigits-3
                  end
               else
                  break
               end
            end

            if b ~= BYTE_CBRACE then
               return nil, "invalid UTF-8 escape sequence", -hexdigits-4
            end

            b = next_byte(state)  -- Skip "}".
            s = to_utf(codepoint)
         elseif b == BYTE_z then
            -- Zap following span of spaces.
            b = skip_space(state, next_byte(state))
         else
            -- Must be a decimal escape.
            local cb

            if b then
               cb = to_dec(b)
            end

            if not cb then
               return nil, "invalid escape sequence", -1
            end

            -- Up to three decimal digits.
            b = next_byte(state)

            if b then
               local c2 = to_dec(b)

               if c2 then
                  cb = 10*cb + c2
                  b = next_byte(state)

                  if b then
                     local c3 = to_dec(b)

                     if c3 then
                        cb = 10*cb + c3

                        if cb > 255 then
                           return nil, "invalid decimal escape sequence", -3
                        end

                        b = next_byte(state)
                     end
                  end
               end
            end

            s = schar(cb)
         end

         if s then
            chunks[#chunks+1] = s
         end

         -- Next chunk starts after escape sequence.
         chunk_start = state.offset
      elseif b == nil or is_newline(b) then
         return nil, "unfinished string"
      else
         b = next_byte(state)
      end
   end

   -- Offset now points at the closing quote.
   local string_value

   if chunks then
      -- Put last chunk into buffer.
      if chunk_start ~= state.offset then
         chunks[#chunks+1] = state.src:get_substring(chunk_start, state.offset-1)
      end

      string_value = tconcat(chunks)
   else
      -- There were no escape sequences.
      string_value = state.src:get_substring(chunk_start, state.offset-1)
   end

   -- Skip the closing quote.
   state.offset = state.offset + 1
   return "string", string_value
end

-- Payload for a number is simply a substring.
-- Luacheck is supposed to be forward-compatible with Lua 5.3 and LuaJIT syntax, so
--    parsing it into actual number may be problematic.
-- It is not needed currently anyway as Luacheck does not do static evaluation yet.
local function lex_number(state, b)
   local start = state.offset

   local exp_lower, exp_upper = BYTE_e, BYTE_E
   local is_digit = to_dec
   local has_digits = false
   local is_float = false

   if b == BYTE_0 then
      b = next_byte(state)

      if b == BYTE_x or b == BYTE_X then
         exp_lower, exp_upper = BYTE_p, BYTE_P
         is_digit = to_hex
         b = next_byte(state)
      else
         has_digits = true
      end
   end

   while b ~= nil and is_digit(b) do
      b = next_byte(state)
      has_digits = true
   end

   if b == BYTE_DOT then
      -- Fractional part.
      is_float = true
      b = next_byte(state)  -- Skip dot.

      while b ~= nil and is_digit(b) do
         b = next_byte(state)
         has_digits = true
      end
   end

   if b == exp_lower or b == exp_upper then
      -- Exponent part.
      is_float = true
      b = next_byte(state)

      -- Skip optional sign.
      if b == BYTE_PLUS or b == BYTE_DASH then
         b = next_byte(state)
      end

      -- Exponent consists of one or more decimal digits.
      if b == nil or not to_dec(b) then
         return nil, "malformed number"
      end

      repeat
         b = next_byte(state)
      until b == nil or not to_dec(b)
   end

   if not has_digits then
      return nil, "malformed number"
   end

   -- Is it cdata literal?
   if b == BYTE_i or b == BYTE_I then
      -- It is complex literal. Skip "i" or "I".
      state.offset = state.offset + 1
   else
      -- uint64_t and int64_t literals can not be fractional.
      if not is_float then
         if b == BYTE_u or b == BYTE_U then
            -- It may be uint64_t literal.
            local b1 = state.src:get_codepoint(state.offset+1)

            if b1 == BYTE_l or b1 == BYTE_L then
               local b2 = state.src:get_codepoint(state.offset+2)

               if b2 == BYTE_l or b2 == BYTE_L then
                  -- It is uint64_t literal.
                  state.offset = state.offset + 3
               end
            end
         elseif b == BYTE_l or b == BYTE_L then
            -- It may be uint64_t or int64_t literal.
            local b1 = state.src:get_codepoint(state.offset+1)

            if b1 == BYTE_l or b1 == BYTE_L then
               local b2 = state.src:get_codepoint(state.offset+2)

               if b2 == BYTE_u or b2 == BYTE_U then
                  -- It is uint64_t literal.
                  state.offset = state.offset + 3
               else
                  -- It is int64_t literal.
                  state.offset = state.offset + 2
               end
            end
         end
      end
   end

   return "number", state.src:get_substring(start, state.offset-1)
end

local function lex_ident(state)
   local start = state.offset
   local b = next_byte(state)

   while (b ~= nil) and (is_alpha(b) or to_dec(b)) do
      b = next_byte(state)
   end

   local ident = state.src:get_substring(start, state.offset-1)

   if keywords[ident] then
      return ident
   else
      return "name", ident
   end
end

local function lex_dash(state)
   local b = next_byte(state)

   -- Is it "-" or comment?
   if b ~= BYTE_DASH then
      return "-"
   end

   -- It is a comment.
   b = next_byte(state)
   local start = state.offset

   -- Is it a long comment?
   if b == BYTE_OBRACK then
      local long_bracket
      b, long_bracket = skip_long_bracket(state)

      if b == BYTE_OBRACK then
         return lex_long_string(state, long_bracket, "long_comment")
      end
   end

   -- Short comment.
   skip_to_newline(state, b)
   local comment_value = state.src:get_substring(start, state.offset - 1)
   return "short_comment", comment_value
end

local function lex_bracket(state)
   -- Is it "[" or long string?
   local b, long_bracket = skip_long_bracket(state)

   if b == BYTE_OBRACK then
      return lex_long_string(state, long_bracket, "string")
   elseif long_bracket == 0 then
      return "["
   else
      return nil, "invalid long string delimiter"
   end
end

local function lex_eq(state)
   local b = next_byte(state)

   if b == BYTE_EQ then
      state.offset = state.offset + 1
      return "=="
   else
      return "="
   end
end

local function lex_lt(state)
   local b = next_byte(state)

   if b == BYTE_EQ then
      state.offset = state.offset + 1
      return "<="
   elseif b == BYTE_LT then
      state.offset = state.offset + 1
      return "<<"
   else
      return "<"
   end
end

local function lex_gt(state)
   local b = next_byte(state)

   if b == BYTE_EQ then
      state.offset = state.offset + 1
      return ">="
   elseif b == BYTE_GT then
      state.offset = state.offset + 1
      return ">>"
   else
      return ">"
   end
end

local function lex_div(state)
   local b = next_byte(state)

   if b == BYTE_SLASH then
      state.offset = state.offset + 1
      return "//"
   else
      return "/"
   end
end

local function lex_ne(state)
   local b = next_byte(state)

   if b == BYTE_EQ then
      state.offset = state.offset + 1
      return "~="
   else
      return "~"
   end
end

local function lex_colon(state)
   local b = next_byte(state)

   if b == BYTE_COLON then
      state.offset = state.offset + 1
      return "::"
   else
      return ":"
   end
end

local function lex_dot(state)
   local b = next_byte(state)

   if b == BYTE_DOT then
      b = next_byte(state)

      if b == BYTE_DOT then
         state.offset = state.offset + 1
         return "...", "..."
      else
         return ".."
      end
   elseif b and to_dec(b) then
      -- Backtrack to dot.
      state.offset = state.offset - 2
      return lex_number(state, next_byte(state))
   else
      return "."
   end
end

local function lex_any(state, b)
   state.offset = state.offset + 1

   if b > 255 then
      b = 255
   end

   return schar(b)
end

-- Maps first bytes of tokens to functions that handle them.
-- Each handler takes the first byte as an argument.
-- Each handler stops at the character after the token and returns the token and,
--    optionally, a value associated with the token.
-- On error handler returns nil, error message and, optionally, start of reported location as negative offset.
local byte_handlers = {
   [BYTE_DOT] = lex_dot,
   [BYTE_COLON] = lex_colon,
   [BYTE_OBRACK] = lex_bracket,
   [BYTE_QUOTE] = lex_short_string,
   [BYTE_DQUOTE] = lex_short_string,
   [BYTE_DASH] = lex_dash,
   [BYTE_SLASH] = lex_div,
   [BYTE_EQ] = lex_eq,
   [BYTE_NE] = lex_ne,
   [BYTE_LT] = lex_lt,
   [BYTE_GT] = lex_gt,
   [BYTE_LDASH] = lex_ident
}

for b=BYTE_0, BYTE_9 do
   byte_handlers[b] = lex_number
end

for b=BYTE_a, BYTE_z do
   byte_handlers[b] = lex_ident
end

for b=BYTE_A, BYTE_Z do
   byte_handlers[b] = lex_ident
end

-- Creates and returns lexer state for source.
function lexer.new_state(src, line_offsets, line_lengths)
   local state = {
      src = src,
      line = 1,
      line_offsets = line_offsets or {},
      line_lengths = line_lengths or {},
      offset = 1
   }

   state.line_offsets[1] = 1

   if src:get_length() >= 2 and src:get_substring(1, 2) == "#!" then
      -- Skip shebang line.
      state.offset = 2
      skip_to_newline(state, next_byte(state))
   end

   return state
end

function lexer.get_quoted_substring_or_line(state, line, offset, end_offset)
   local line_length = state.line_lengths[line]

   if line_length then
      local line_end_offset = state.line_offsets[line] + line_length - 1

      if line_end_offset < end_offset then
         end_offset = line_end_offset
      end
   end

   return "'" .. state.src:get_printable_substring(offset, end_offset) .. "'"
end

-- Looks for next token starting from state.line, state.offset.
-- Returns next token, its value and its location (line, offset).
-- Sets state.line, state.offset to token end location + 1.
-- Fills state.line_offsets and state.line_lengths.
-- On error returns nil, error message, error location (line, offset), error end offset.
function lexer.next_token(state)
   local line_offsets = state.line_offsets
   local b = skip_space(state, state.src:get_codepoint(state.offset))

   -- Save location of token start.
   local token_line = state.line
   local line_offset = line_offsets[token_line]
   local token_offset = state.offset

   if not b then
      -- EOF token has length 1.
      state.offset = state.offset + 1
      state.line_lengths[token_line] = token_offset - line_offset
      return "eof", nil, token_line, token_offset
   end

   local token, token_value, relative_error_offset = (byte_handlers[b] or lex_any)(state, b)

   if relative_error_offset then
      -- Error relative to current offset.
      local error_offset = state.offset + relative_error_offset
      local error_end_offset = math.min(state.offset, state.src:get_length())
      local error_message = token_value .. " " .. lexer.get_quoted_substring_or_line(state,
         state.line, error_offset, error_end_offset)
      return nil, error_message, state.line, error_offset, error_end_offset
   end

   -- Single character errors fall through here.
   return token, token_value, token_line, token_offset, not token and token_offset
end

return lexer