tlc.lua

--------------------------------------------------------------------------------
-- TINY LUA COMPILER (TLC)
--------------------------------------------------------------------------------
-- URL:     https://github.com/bytexenon/tiny-lua-compiler
-- License: MIT (free to use and modify, attribution appreciated)
-- Version: 1.0.1
--------------------------------------------------------------------------------
-- DESCRIPTION:
--   A minimal, educational compiler for Lua 5.1, written in pure Lua 5.1.
--   It allows for dynamic compilation and execution of Lua code within a Lua
--   environment without needing to rely on the `load` or `loadstring`
--   functions.
--   It features a high-performance lexer, a complete parser, proper
--   operator precedence, lexical scoping, a bytecode generator, and a
--   register-based virtual machine.
--
--   Pipeline Stages:
--     1. Tokenizer:        Source Text -> Token Stream
--     2. Parser:           Tokens      -> Abstract Syntax Tree (AST)
--     3. Code Generator:   AST         -> Function Prototype (proto)
--     4. Bytecode Emitter: Prototype   -> Binary Chunk (bytecode string)
--     5. Virtual Machine:  Prototype   -> Execution
--
-- USAGE:
--   local tlc = require("tlc")
--   tlc.run("print('Hello from TLC!')")
--
-- API SUMMARY:
--   High-Level
--     tlc.tokenize(code) ....................... Returns Token List
--     tlc.parse(code) .......................... Returns AST
--     tlc.compileToProto(code) ................. Returns Prototype
--     tlc.compile(code) ........................ Returns Bytecode String
--     tlc.run(code, env?, ...?) ................ Executes Code
--
--   Low-Level (Explicit Stages)
--     tlc.parseTokens(tokens) .................. Returns AST
--     tlc.generate(ast) ........................ Returns Prototype
--     tlc.emit(proto) .......................... Returns Bytecode String
--     tlc.execute(proto, env?, ...?) ........... Executes Prototype
--
--   Manual Pipeline (All Stages)
--     tlc.Tokenizer.new(code):tokenize()            -> tokens
--     tlc.Parser.new(tokens):parse()                -> ast
--     tlc.CodeGenerator.new(ast):generate()         -> proto
--     tlc.BytecodeEmitter.emit(proto)               -> bytecode string
--     tlc.VirtualMachine.execute(proto, env?, ...?) -> executes proto
--
-- NOTE:
--   * Internal component classes should be accessed via high-level API.
--   * TLC can run on all Lua versions (5.1 to 5.5), but generated bytecode from
--     BytecodeEmitter is only compatible with Lua 5.1.
--------------------------------------------------------------------------------
-- stylua: ignore start

-- In Lua, iterating over a list to check for existence is slow. By converting
-- the list `{"a", "b"}` into `{a=true, b=true}`, we leverage Lua's internal
-- hash map architecture. This makes checking `if keywords[word]` faster than
-- iterating through the list each time.
local function createLookupTable(list)
  local lookup = {}
  for _, value in ipairs(list) do
    lookup[value] = true
  end
  return lookup
end

-- A helper function for creating a trie (prefix tree) from a list of strings.
--
-- This structure solves the "Longest Prefix Match" problem. When the tokenizer
-- encounters a character like `>`, it needs to decide if it's a standalone `>`,
-- or part of a longer operator like `>=`.
--
-- Instead of complex lookahead logic, we traverse this tree. If we can travel
-- deeper (e.g., from `>` to `=`), we continue. If we stop, we know we've
-- matched the longest valid operator.
local function createTrie(ops)
  local trie = {}
  for _, op in ipairs(ops) do
    local node = trie
    for char in op:gmatch(".") do
      node[char] = node[char] or {}
      node = node[char]
    end

    -- Mark this node as a valid ending point for an operator.
    node.value = op
  end
  return trie
end

-- Running `string.match(char, "%d")` on every character of the source code
-- is computationally expensive due to pattern matching overhead.
--
-- Since a byte can only have 256 possible values, we can pre-calculate
-- the result of the pattern match for every possible character (1-255).
-- This transforms a regex operation into a simple array lookup.
local function createPatternLookup(pattern)
  local lookup = {}
  for code = 1, 255 do
    local char = string.char(code)
    if string.match(char, pattern) then
      lookup[char] = true
    end
  end
  return lookup
end

-- ============================================================================
--                                    (/^▽^)/
--                                 THE TOKENIZER!
-- ============================================================================
--
-- The first stage of compilation. The tokenizer (also called a lexer or
-- scanner) reads raw source text character by character and produces a flat
-- stream of labeled tokens - the smallest meaningful units in the language.
--
-- Example:
--   Input:  "local x = 42 + y"
--   Output: [Keyword:local] [Identifier:x] [Equals] [Number:42]
--           [Operator:+] [Identifier:y]
--
-- Separating tokenization from parsing keeps both stages simple. The tokenizer
-- only needs to recognize lexical patterns ("is this a number?", "is this a
-- keyword?"). The parser only needs to understand structure ("is this a valid
-- if-statement?"). Neither needs to worry about the other's concerns.

-- In Lua, when you call string.sub with an index beyond the end of the string,
-- it returns an empty string. Storing it as a named constant makes the intent explicit.
local END_OF_FILE = ""

local TOKENIZER_CONFIG = {
  DIGIT = createPatternLookup("%d"), -- 0-9
  WHITESPACE = createPatternLookup("%s"), -- Whitespace (space, newline, etc.)
  IDENTIFIER_START = createPatternLookup("[%a_]"), -- a-z, A-Z, underscore

  -- The main lookup table for single-character tokens that can never be part
  -- of a longer operator. Splitting them from UNSAFE_PUNCTUATION speeds up
  -- common tokens like parentheses and commas: we emit their token immediately
  -- without having to attempt an operator trie match first.
  SAFE_PUNCTUATION = {
    [":"] = "Colon",        [";"] = "Semicolon",
    ["("] = "LeftParen",    [")"] = "RightParen",
    ["{"] = "LeftBrace",    ["}"] = "RightBrace",
    ["]"] = "RightBracket", [","] = "Comma"

    -- These three are moved to UNSAFE_PUNCTUATION because they require
    -- lookahead to determine their full meaning. We cannot immediately emit
    -- a token upon seeing them the way we can for "(" or ",".
    --
    -- ["["] = "LeftBracket" - Can start long strings
    -- ["."] = "Dot"         - Can start ".." or "..."
    -- ["="] = "Equals"      - Can start "=="
  },

  -- Single-character tokens that may begin a longer operator sequence.
  -- These are checked only after the operator trie has had a chance to match
  -- a longer form: for example, `[` could start a long string `[[`, and `=`
  -- could start `==`. We keep them here as a fallback for when the trie finds
  -- no multi-character match.
  UNSAFE_PUNCTUATION = {
    ["["] = "LeftBracket",
    ["."] = "Dot",
    ["="] = "Equals"
  },

  -- This trie allows us to efficiently match operators of varying lengths
  -- without hardcoding complex lookahead logic. When we encounter a character
  -- that could start an operator, we traverse the trie to find the longest valid
  -- match.
  OPERATOR_TRIE = createTrie({
    "^",  "*",  "/",  "%",  "+", "-", "<", ">", "#",
    "<=", ">=", "==", "~=", ".."
  }),

  -- Maps characters following a backslash '\' to their actual byte
  -- values. For example, 'n' maps to the newline character '\n'.
  ESCAPES = {
    ["a"]  = "\a", ["b"]  = "\b", ["f"]  = "\f",
    ["n"]  = "\n", ["r"]  = "\r", ["t"]  = "\t",
    ["v"]  = "\v", ["\\"] = "\\", ["\""] = "\"",
    ["\'"] = "\'"
  },

  -- Words that look like identifiers but carry special grammatical meaning in Lua.
  -- The tokenizer first consumes any word as a generic identifier, then performs
  -- an O(1) lookup against this table. If the word is found here, its token type
  -- is changed from "Identifier" to "Keyword".
  KEYWORDS = createLookupTable({
    "and",      "break", "do",    "else",
    "elseif",   "end",   "false", "for",
    "function", "if",    "in",    "local",
    "nil",      "not",   "or",    "repeat",
    "return",   "then",  "true",  "until",
    "while"
  })
}

local Tokenizer = {}
Tokenizer.__index = Tokenizer

function Tokenizer.new(code)
  assert(
    type(code) == "string",
    "Tokenizer.new requires a string as input code. Got: " .. type(code)
  )

  -- Create the instance table and set its metatable to `Tokenizer`.
  -- This grants the instance access to all methods defined in the Tokenizer
  -- class (like `consume`, `getNextToken`, `tokenize`, etc.) via prototypal
  -- inheritance, avoiding the need to copy methods for each instance.
  return setmetatable({
    code       = code,
    curChar    = code:sub(1, 1),
    curCharPos = 1
  }, Tokenizer)
end

--// Character Navigation //--

-- Looks ahead by 'n' characters in the character stream without advancing
-- the current position.
function Tokenizer:lookAhead(n)
  local targetCharPos   = self.curCharPos + n
  local lookedAheadChar = self.code:sub(targetCharPos, targetCharPos)
  return lookedAheadChar
end

-- Consumes (advances past) 'n' characters in the character stream.
function Tokenizer:consume(n)
  if self.curChar == END_OF_FILE then
    error("Internal Tokenizer Error: Attempted to consume past end of stream.")
  end

  local nextCharPos = self.curCharPos + n
  local nextChar    = self.code:sub(nextCharPos, nextCharPos)

  -- Update the tokenizer's state.
  self.curCharPos = nextCharPos
  self.curChar    = nextChar
end

-- When string.find() is used to bulk-consume characters, it returns an end
-- position but doesn't update the tokenizer state. Call this function to
-- synchronize self.curCharPos and self.curChar with the new position.
--
-- NOTE: self.curChar will be stale after string.find() and before calling
-- this function. Don't access self.curChar in that window!
function Tokenizer:setCurrentPosition(pos)
  self.curCharPos = pos
  self.curChar    = self.code:sub(pos, pos)
end

--// Multi-Character Checkers //--

-- Checks if the current character and the next character form the
-- prefix of a hexadecimal number (0x or 0X).
function Tokenizer:isHexadecimalNumberPrefix()
  if self.curChar == "0" then
    local nextChar = self:lookAhead(1)
    return nextChar == "x" or nextChar == "X"
  end

  return false
end

--// Utils //--

-- Calculates the "depth" of a long string or comment delimiter.
-- This counts how many '=' characters follow the initial '['.
--
-- For example:
--   [==[  -> depth of 2
--   [=[   -> depth of 1
--   [[    -> depth of 0
function Tokenizer:calculateDelimiterDepth()
  -- Match a sequence of '=' characters starting from the current position.
  -- "^" ensures we match from the current position, and "=*" matches zero or
  -- more '=' characters.
  local startMatchPos, endMatchPos = string.find(self.code, "^=*", self.curCharPos)
  return endMatchPos - startMatchPos + 1
end

--// Consumers //--
function Tokenizer:skipWhitespace()
  -- Match optional whitespace characters after the initial whitespace character.
  local _, endPos = string.find(self.code, "^%s*", self.curCharPos + 1)
  self:setCurrentPosition(endPos + 1)
end

function Tokenizer:consumeIdentifier()
  -- The initial character has already been verified to be a
  -- valid identifier start. Match any subsequent characters
  -- that are valid followers (letters, digits, underscores).
  local startPos, endPos = string.find(self.code, "^[0-9a-zA-Z_]*", self.curCharPos)
  self:setCurrentPosition(endPos + 1)
  return self.code:sub(startPos, endPos)
end

-- Must be called after verifying the hex prefix (0x or 0X) is present.
-- Returns the string representation including the prefix (e.g., "0xFF").
function Tokenizer:consumeHexadecimalNumber()
  local startPos = self.curCharPos

  -- Find hex digits after the "0x" or "0X" prefix (skip first 2 chars).
  -- At least one hex digit is required after the prefix.
  local _, endPos = string.find(self.code, "^[0-9a-fA-F]+", self.curCharPos + 2)
  if not endPos then
    error(string.format(
      "malformed number near '%s'",
      self.code:sub(startPos, self.curCharPos - 1)
    ))
  end

  local nextChar = self.code:sub(endPos + 1, endPos + 1)

  -- Is it a hexadecimal literal with an exponent part (e.g., 0x2p4)?
  if nextChar == "p" or nextChar == "P" then
    local _, exponentEndPos = string.find(self.code, "^[0-9]+", endPos + 2)
    if not exponentEndPos then
      error(string.format(
        "malformed number near '%s'",
        self.code:sub(startPos, self.curCharPos - 1)
      ))
    end
    endPos = exponentEndPos
  end

  self:setCurrentPosition(endPos + 1)
  return self.code:sub(startPos, self.curCharPos - 1)
end

-- Recognizes all Lua 5.1 number formats:
--   - Decimal integers:      42, 123
--   - Floating point:        3.14, .5, 42. (trailing dot is allowed)
--   - Hexadecimal:           0xFF, 0X1A3, 0x2b, 0x2p4 (delegates to helper)
--   - Scientific notation:   1e10, 3.14E-2, 5e+3
--   - Mixed formats:         1.e3 (1000), .5e-1 (0.05)
function Tokenizer:consumeNumber()
  local startPos = self.curCharPos

  -- Hexadecimal Case: Check for "0x" or "0X" prefix.
  -- Pattern: 0[xX][0-9a-fA-F]+[[pP][0-9]+]? (handled in helper)
  if self:isHexadecimalNumberPrefix() then
    return self:consumeHexadecimalNumber()
  end

  -- Decimal Case: Parse the integer part first.
  -- Match optional leading digits (e.g., "42" in "42.5" or "0" in "0.5").
  local _, endPos = string.find(self.code, "^[0-9]*", startPos)

  -- Fractional Part: Check for decimal point followed by optional digits.
  local nextChar = self.code:sub(endPos + 1, endPos + 1)
  if nextChar == "." then
    _, endPos = string.find(self.code, "^[0-9]*", endPos + 2)
  end

  -- Exponent Part: Check for scientific notation (e.g., 1e10, 3.5E-2).
  -- The exponent must have at least one digit after the optional sign.
  nextChar = self.code:sub(endPos + 1, endPos + 1)
  if nextChar == "e" or nextChar == "E" then
    _, endPos = string.find(self.code, "^[+-]?[0-9]+", endPos + 2)

    if not endPos then
      error(string.format(
        "malformed number near '%s'",
        self.code:sub(startPos, self.curCharPos - 1)
      ))
    end
  end

  -- Done parsing the number, update the tokenizer's
  -- position to the end of the number.
  self:setCurrentPosition(endPos + 1)

  -- Lua forbids identifiers immediately after numbers (e.g., "123abc").
  -- This prevents ambiguity: is "123abc" a number or an identifier?
  -- The official Lua lexer treats this as a malformed number, not an
  -- identifier. We replicate this behavior for compatibility.
  if TOKENIZER_CONFIG.IDENTIFIER_START[self.curChar] then
    error(string.format(
      "malformed number near '%s'",
      self.code:sub(startPos, self.curCharPos - 1)
    ))
  end

  return self.code:sub(startPos, self.curCharPos - 1)
end

-- Called after encountering a backslash (`\`) in a simple string.
-- Returns the resolved character (e.g., `\n` -> newline, `\65` -> `A`).
function Tokenizer:consumeEscapeSequence()
  self:consume(1) -- Consume the backslash.
  local curChar = self.curChar

  -- Named escape? (e.g., \n -> newline, \t -> tab)
  local convertedChar = TOKENIZER_CONFIG.ESCAPES[curChar]
  if convertedChar then
    self:consume(1)
    return convertedChar

  -- Numeric escape? (e.g., \65 -> `A`, \123 -> `{`)
  elseif TOKENIZER_CONFIG.DIGIT[curChar] then
    -- Match up to three digits (0-9), but at least one is required.
    local pattern = "^[0-9][0-9]?[0-9]?"
    local startPos, endPos = string.find(self.code, pattern, self.curCharPos)
    local numberString = self.code:sub(startPos, endPos)
    local number = tonumber(numberString)
    if number > 255 then
      error(string.format("escape sequence too large near '\\%s'", numberString))
    end

    self:setCurrentPosition(endPos + 1)
    return string.char(number)
  end

  -- Unknown escape: Lua's lexer is lenient and passes
  -- the character through verbatim (e.g., "\e" -> "e").
  -- We replicate this quirk for compatibility.
  self:consume(1)
  return curChar
end

function Tokenizer:consumeSimpleString()
  local quoteChar   = self.curChar
  local stringTable = {}
  local stringIndex = 0
  self:consume(1) -- Consume the opening quote.

  while self.curChar ~= quoteChar do
    if self.curChar == END_OF_FILE then
      error("Unclosed string")
    elseif self.curChar == "\n" then
      error("Unclosed string (newline found)")

    -- Is it an escape sequence?
    elseif self.curChar == "\\" then
      local consumedEscape = self:consumeEscapeSequence()
      stringIndex = stringIndex + 1
      stringTable[stringIndex] = consumedEscape
    else
      stringIndex = stringIndex + 1
      stringTable[stringIndex] = self.curChar
      self:consume(1)
    end
  end
  self:consume(1) -- Consume the closing quote.

  return table.concat(stringTable)
end

-- Consumes a long string literal using [[...]] or [=[...]=] syntax.
-- Automatically skips a leading newline if present (Lua 5.1 behavior).
function Tokenizer:consumeLongString()
  self:consume(1) -- Consume the first '['.
  local depth = self:calculateDelimiterDepth()
  self:consume(depth) -- Consume the '=' characters.

  if self.curChar ~= "[" then
    error("Invalid long string delimiter")
  end
  self:consume(1)
  local firstStringChar = self.curChar

  -- Long string starts with a newline?
  if firstStringChar == "\n" then
    -- If the first character inside the long string is a newline,
    -- we skip it entirely. This is an edge case in the official lexer source.
    -- Source: https://www.lua.org/source/5.1/llex.c.html#read_long_string

    -- Skip it.
    self:consume(1)
  end

  local stringStartPos   = self.curCharPos
  local closingDelimiter = "%]" .. string.rep("=", depth) .. "%]"
  local startPos, endPos = string.find(self.code, closingDelimiter, self.curCharPos)
  if not endPos then
    error("Unclosed long string")
  end

  self:setCurrentPosition(endPos + 1)
  return self.code:sub(stringStartPos, startPos - 1)
end

-- Uses the operator trie for longest prefix matching.
function Tokenizer:consumeOperator()
  local node = TOKENIZER_CONFIG.OPERATOR_TRIE
  local operator

  -- Walk the trie character by character.
  local startPos   = self.curCharPos
  local newCharPos = startPos
  while true do
    -- Optimization: Use `sub` instead of `lookAhead` to
    -- avoid extra function call overhead in this tight loop.
    local character = string.sub(self.code, newCharPos, newCharPos)
    node = node[character]

    -- If the path doesn't exist in the trie,
    -- it means that we've gone as far as we can.
    if not node then break end

    -- If this node represents a valid operator, remember it.
    -- We keep going to see if there's a longer match (greedy matching).
    operator   = node.value
    newCharPos = newCharPos + 1
  end

  if not operator then return end
  self:consume(newCharPos - startPos)

  return operator
end

function Tokenizer:skipShortComment()
  -- Match anything until the next newline or end of input.
  local _, endPos = string.find(self.code, "[^\n]*", self.curCharPos)
  self:setCurrentPosition(endPos + 1)
end

-- Consumes a multi-line comment using --[[...]] or --[=[...]=] syntax.
function Tokenizer:skipLongComment()
  self:consume(1) -- Consume the first '['.
  local depth = self:calculateDelimiterDepth()
  self:consume(depth) -- Consume the '=' characters.

  -- Not a complete long-comment delimiter - fall through to the short-comment path.
  if self.curChar ~= "[" then
    self:skipShortComment()
    return
  end
  self:consume(1)

  -- Consume the closing delimiter using Lua pattern matching.
  -- (We're putting the modulo signs around the `]` to escape it in the pattern).
  local closingDelimiter = "%]" .. string.rep("=", depth) .. "%]"
  local _, endPos = string.find(self.code, closingDelimiter, self.curCharPos)
  if not endPos then
    error("Unclosed long comment")
  end

  self:setCurrentPosition(endPos + 1)
end

function Tokenizer:skipComment()
  self:consume(2)
  if self.curChar == "[" then
    return self:skipLongComment()
  end
  return self:skipShortComment()
end

--// Token Consumer Handler //--

-- Consumes and returns tokenKind, tokenValue, tokenRaw for the next
-- token in the stream. If it returns a nil "tokenKind", it means that
-- no token was produced (e.g., whitespace or comment), so the caller should
-- call this method again to get the next token.
--
-- Note that "tokenValue" and "tokenRaw" may also be nil for tokens that don't
-- have associated values (e.g., punctuation).
--
-- For optimization, the checks in this function are ordered by expected
-- frequency in typical Lua source code. Whitespace and identifiers come
-- first since they appear most often. Safe punctuation and strings follow.
-- Numbers, comments, dot-prefixed numbers, varargs, and long strings are
-- progressively less common. Symbolic operators are matched via the trie.
-- Unsafe punctuation is handled last as a fallback.
function Tokenizer:getNextToken()
  local curChar = self.curChar

  if TOKENIZER_CONFIG.WHITESPACE[curChar] then
    self:skipWhitespace()
    curChar = self.curChar

    -- If the stream ends with whitespace, check here to avoid returning
    -- a stale token to the caller.
    if curChar == END_OF_FILE then
      return nil
    end
  end

  if TOKENIZER_CONFIG.IDENTIFIER_START[curChar] then
    local identifier = self:consumeIdentifier()

    -- Is it a keyword? (E.g., "while", "local", "function", etc.)
    if TOKENIZER_CONFIG.KEYWORDS[identifier] then
      return "Keyword", identifier
    end

    return "Identifier", identifier
  end

  -- Processes "safe" punctuation first (single-character tokens that
  -- cannot be part of multi-character tokens).
  local safeTokenKind = TOKENIZER_CONFIG.SAFE_PUNCTUATION[curChar]
  if safeTokenKind then
    self:consume(1)
    return safeTokenKind, nil

  elseif curChar == '"' or curChar == "'" then
    return "String", self:consumeSimpleString()

  elseif TOKENIZER_CONFIG.DIGIT[curChar] then
    local numberString = self:consumeNumber()
    local numberValue = tonumber(numberString)
    if not numberValue then
      error(string.format(
        "malformed number near '%s'",
        numberString
      ))
    end

    return "Number", numberValue, numberString

  elseif curChar == "-" and self:lookAhead(1) == "-" then
    self:skipComment()
    return nil

  -- Handle numbers starting with a decimal point and vararg ("...").
  elseif curChar == "." then
    local nextChar = self:lookAhead(1)

    -- Repeat the number tokenization logic for the case where a number starts
    -- with a decimal point (e.g., ".5", ".5e-1", "1.e3").
    if TOKENIZER_CONFIG.DIGIT[nextChar] then
      local numberString = self:consumeNumber()
      local numberValue = tonumber(numberString)
      if not numberValue then
        error(string.format(
          "malformed number near '%s'",
          numberString
        ))
      end

      return "Number", numberValue, numberString

    elseif nextChar == "." and self:lookAhead(2) == "." then
      self:consume(3)
      return "Vararg", nil
    end

  -- Handle long string literals (with [[...]] or [=[...]=], etc.).
  elseif curChar == "[" then
    local nextChar = self:lookAhead(1)
    if nextChar == "[" or nextChar == "=" then
      return "String", self:consumeLongString()
    end
  end

  -- Symbolic operators like "+", "==", "<=", "..".
  local operator = self:consumeOperator()
  if operator then
    return "Operator", operator
  end

  -- Handle "unsafe" punctuation (single-character tokens that
  -- could be part of multi-character tokens).
  local unsafeTokenKind = TOKENIZER_CONFIG.UNSAFE_PUNCTUATION[curChar]
  if unsafeTokenKind then
    self:consume(1)
    return unsafeTokenKind, nil
  end

  -- If we reach this point, it means the current character doesn't match
  -- any known token kind. This is an error condition, as the character
  -- is not valid in the Lua language syntax.
  error(string.format(
    "Unexpected character '%s' at position %d",
    tostring(curChar),
    self.curCharPos
  ))
end

--// Tokenizer Main Method //--
function Tokenizer:tokenize()
  local tokens, tokenIndex = {}, 0

  while self.curChar ~= END_OF_FILE do
    local tokenKind, tokenValue, tokenRaw = self:getNextToken()

    if tokenKind then
      tokenIndex = tokenIndex + 1
      tokens[tokenIndex] = {
        kind  = tokenKind,
        value = tokenValue,
        raw   = tokenRaw
      }
    end
  end

  -- Append a sentinel EOF token at the end of the stream. The parser can then
  -- assume the current token is always non-nil, removing the need for
  -- `if token then` guards on every token check.
  table.insert(tokens, {
    kind  = "EOF",
    value = nil,
    raw   = nil
  })

  return tokens
end

-- ============================================================================
--                         (ﾉ◕ヮ◕)ﾉ*:･ﾟ THE PARSER!
-- ============================================================================
--
-- The parser is the second stage of compilation. It transforms a flat stream
-- of tokens into an Abstract Syntax Tree (AST) - a nested structure that makes
-- the grammatical meaning of the source explicit.
--
-- Where the tokenizer asks "what symbols are present?", the parser asks "how
-- are those symbols related?" A token stream is flat; the AST makes nesting
-- and operator precedence visible through its shape.
--
-- Two complementary techniques are used. Recursive descent handles the top-down
-- structure of statements: each construct (`if`, `while`, `for`, `function`,
-- `local`, and so on) has its own parsing function. Precedence climbing handles
-- expressions, correctly binding `*` tighter than `+` and keeping `^` and `..`
-- right-associative - all without writing a separate grammar rule for each
-- precedence level.
--
-- Example:
--   Input source:
--     local x = 1 + 2 * 3
--
--   Tokens:
--     [
--       { kind = "Keyword",    value = "local" },
--       { kind = "Identifier", value = "x" },
--       { kind = "Equals" },
--       { kind = "Number",     value = 1 },
--       { kind = "Operator",   value = "+" },
--       { kind = "Number",     value = 2 },
--       { kind = "Operator",   value = "*" },
--       { kind = "Number",     value = 3 },
--       { kind = "EOF" }
--     ]
--
--   AST:
--     Program
--       LocalDeclarationStatement
--         variables: ["x"]
--         initializers:
--           BinaryOperator("+")
--             left:  NumericLiteral(1)
--             right: BinaryOperator("*")
--               left:  NumericLiteral(2)
--               right: NumericLiteral(3)
--
-- Notice the tree shape: `2 * 3` is nested under `+`, reflecting that
-- multiplication binds more tightly and must be evaluated first.

local PARSER_CONFIG = {
  -- Operator Precedence & Associativity Table
  -- -----------------------------------------
  -- This table drives the "Precedence Climbing" algorithm in `parseBinaryExpression`.
  -- Each entry is a pair { left_binding_power, right_binding_power } that encodes
  -- both the operator's precedence and its associativity.
  --
  -- Higher numbers mean tighter binding - that is why `*` is evaluated before
  -- `+` in `1 + 2 * 3`. Left-associative operators (e.g., `+`, `-`) carry
  -- equal left and right powers, so `1 - 2 - 3` parses as `(1 - 2) - 3`.
  -- Right-associative operators (e.g., `^`, `..`) carry a slightly higher right
  -- power, so `a ^ b ^ c` parses as `a ^ (b ^ c)`.
  PRECEDENCE = {
    ["+"]   = {6, 6},  ["-"]   = {6, 6}, -- Addition/Subtraction
    ["*"]   = {7, 7},  ["/"]   = {7, 7}, -- Multiplication/Division
    ["%"]   = {7, 7},                    -- Modulo
    ["^"]   = {10, 9},                   -- Exponentiation (Right-associative!)
    [".."]  = {5, 4},                    -- String Concatenation (Right-associative!)
    ["=="]  = {3, 3},  ["~="]  = {3, 3}, -- Equality / Inequality
    ["<"]   = {3, 3},  [">"]   = {3, 3}, -- Less Than / Greater Than
    ["<="]  = {3, 3},  [">="]  = {3, 3}, -- Less Than or Equal / Greater Than or Equal
    ["and"] = {2, 2},                    -- Logical AND (Short-circuiting)
    ["or"]  = {1, 1}                     -- Logical OR (Short-circuiting)
  },

  -- Precedence for all unary operators (like `-x` or `not x`).
  -- They bind tighter than most binary operators but looser than exponentiation.
  -- (e.g., `-2^3` parses as `-(2^3)`, not `(-2)^3`).
  UNARY_PRECEDENCE = 8,

  -- L-Values (Locator Values).
  -- These are the specific node types that are allowed to appear on the
  -- LEFT side of an assignment (e.g., `x = 1` or `t.k = 1`).
  -- You can't assign to a number (`5 = 1` is invalid), so "NumericLiteral" isn't here.
  LVALUE_NODES = createLookupTable({ "Identifier", "IndexExpression" }),

  -- Keywords that signal the end of a block. When the parser sees one of
  -- these, it stops parsing statements in the current block and returns
  -- control to the caller that owns that keyword.
  TERMINATION_KEYWORDS = createLookupTable({ "end", "else", "elseif", "until" }),

  UNARY_OPERATORS  = createLookupTable({ "-", "#", "not" }),
  BINARY_OPERATORS = createLookupTable({
    "+",  "-",   "*",  "/",
    "%",  "^",   "..", "==",
    "~=", "<",   ">",  "<=",
    ">=", "and", "or"
  })
}

local Parser = {}
Parser.__index = Parser

function Parser.new(tokens)
  assert(
    type(tokens) == "table",
    "Parser.new requires a table of tokens. Got: " .. type(tokens)
  )

  local self = setmetatable({}, Parser)
  self.tokens = tokens
  self.currentToken = tokens[1]
  self.currentTokenIndex = 1

  -- A special flag which gets set to `true` when we encounter either
  -- a `break` statement, or a `return` statement. This is used to prevent
  -- parsing code that syntactically may be correct, but is semantically invalid.
  self.shouldTerminateBlock = false

  -- A lookup table that maps statement-starting keywords to their corresponding
  -- parsing functions. This allows the main statement parsing loop to quickly
  -- dispatch to the correct parser based on the current token.
  self.statementHandlers = {
    ["do"]       = self.parseDoStatement,
    ["break"]    = self.parseBreakStatement,
    ["for"]      = self.parseForStatement,
    ["function"] = self.parseFunctionDeclaration,
    ["if"]       = self.parseIfStatement,
    ["local"]    = self.parseLocalStatement,
    ["repeat"]   = self.parseRepeatStatement,
    ["return"]   = self.parseReturnStatement,
    ["while"]    = self.parseWhileStatement
  }

  return self
end

--// Token Navigation //--

-- NOTE: This is called only in `:parseTableConstructor`, the one place in the
-- parser that needs to look more than one token ahead.
function Parser:peekNext()
  return self.tokens[self.currentTokenIndex + 1]
end

function Parser:advance(n)
  local newTokenIndex = self.currentTokenIndex + n
  self.currentTokenIndex = newTokenIndex
  self.currentToken = self.tokens[newTokenIndex]
end

function Parser:getTokenDescription(token)
  if not token.value then return tostring(token.kind) end
  return string.format("%s [%s]", tostring(token.kind), tostring(token.value))
end

function Parser:expect(tokenKind, tokenValue)
  local token = self.currentToken
  if token.kind == tokenKind and token.value == tokenValue then
    self:advance(1)
    return
  end

  error(string.format(
    "Expected %s [%s] token, got: %s",
    tostring(tokenKind),
    tostring(tokenValue),
    self:getTokenDescription(token)
  ))
end

-- Conditionally consume the current token by kind.
--
-- This is the parser equivalent of asking "is the next token a comma?" and,
-- if so, advancing past it.
function Parser:advanceIfMatch(tokenKind)
  if self.currentToken.kind == tokenKind then
    self:advance(1)
    return true
  end

  return false
end

--// Token Checkers //--
-- Lightweight checks to see if a token matches a specific kind and/or value
-- without throwing an error if it doesn't match.

function Parser:isTokenKind(kind)
  return self.currentToken.kind == kind
end

function Parser:isKeyword(keyword)
  local token = self.currentToken
  return token.kind == "Keyword" and token.value == keyword
end

function Parser:isUnaryOperator()
  local token = self.currentToken
  return (token.kind == "Operator" or token.kind == "Keyword")
        and PARSER_CONFIG.UNARY_OPERATORS[token.value]
end

function Parser:isBinaryOperator()
  local token = self.currentToken
  return (token.kind == "Operator" or token.kind == "Keyword")
        and PARSER_CONFIG.BINARY_OPERATORS[token.value]
end

--// Token Expectation //--
-- These functions verify that the current token matches an expected kind or
-- keyword and then consume it. If the check fails, they throw an error.

function Parser:expectKind(expectedKind)
  if not self:isTokenKind(expectedKind) then
    error(string.format(
      "Expected a %s, but found %s",
      tostring(expectedKind),
      tostring(self.currentToken.kind)
    ))
  end

  self:advance(1)
end

function Parser:expectKeyword(keyword)
  if not self:isKeyword(keyword) then
    error(string.format(
      "Expected keyword '%s', but found %s",
      keyword,
      self:getTokenDescription(self.currentToken)
    ))
  end

  self:advance(1)
end

--// AST Node Checkers //--
-- Functions to check properties of already-parsed AST nodes.

-- Valid LValues are single variables or table access expressions.
function Parser:isValidAssignmentLvalue(node)
  return node and PARSER_CONFIG.LVALUE_NODES[node.kind]
end

--// Parsers //--
-- These functions are responsible for parsing specific syntactic constructs
-- in the token stream and building the corresponding AST nodes.

-- `(<parameters>) <body> end`
function Parser:parseFunctionExpression()
  local parameters, isVararg = self:parseParameterList()
  local body = self:parseCodeBlock()
  self:expectKeyword("end")
  return { kind = "FunctionExpression",
    body       = body,
    parameters = parameters,
    isVararg   = isVararg
  }
end

-- `if/elseif <condition> then <body>`
function Parser:parseIfClause(keyword)
  self:expectKeyword(keyword) -- "if" / "elseif"
  local condition = self:expectExpression()
  self:expectKeyword("then")
  local body = self:parseCodeBlock()
  return { kind = "IfClause",
    condition = condition,
    body      = body
  }
end

-- `do <body> end`
--
-- Used as a utility parser for every statement that opens a `do ... end` block
-- (e.g., `while`, `for`, and standalone `do` statements).
function Parser:parseDoEndLoopBlock()
  self:expectKeyword("do")
  local body = self:parseCodeBlock()
  self:expectKeyword("end")
  return body
end

function Parser:expectIdentifier()
  local currentToken = self.currentToken
  self:expectKind("Identifier")
  return currentToken.value
end

-- `<identifier>, <identifier>, ...`
function Parser:parseIdentifierList()
  local identifiers = { self:expectIdentifier() }
  while self:advanceIfMatch("Comma") do
    table.insert(identifiers, self:expectIdentifier())
  end

  return identifiers
end

-- `(<identifier>, <identifier>, ..., <vararg>)`
function Parser:parseParameterList()
  self:expect("LeftParen")
  local parameters = {}
  local isVararg   = false

  if self:advanceIfMatch("RightParen") then
    return parameters, isVararg
  end

  -- Keep consuming parameters until we hit `)` or the vararg marker.
  -- Unlike some languages, Lua does not allow trailing commas here.
  while true do
    if self:isTokenKind("Identifier") then
      table.insert(parameters, self.currentToken.value)
      self:advance(1)

    elseif self:isTokenKind("Vararg") then
      isVararg = true
      self:advance(1)
      -- Lua 5.1 grammar: once "..." is seen, no more parameters are allowed.
      -- This matches the reference parser. We enforce it with a break rather
      -- than an error because the ")" check below will catch malformed syntax.
      break

    else
      error(string.format(
        "Expected parameter name or '...', but got: %s",
        self:getTokenDescription(self.currentToken)
      ))
    end

    if self.currentToken.kind == "RightParen" then
      break
    elseif not self:advanceIfMatch("Comma") then
      error(string.format(
        "Expected ',' or ')' after parameter, but got: %s",
        self:getTokenDescription(self.currentToken)
      ))
    end
  end
  self:expect("RightParen")

  return parameters, isVararg
end

-- `.<identifier>` | `[<expression>]`
--
-- Parse one indexing step after a prefix expression.
--
-- Lua has two surface syntaxes here:
--   `tbl.name`   -> key is the string literal "name"
--   `tbl[expr]`  -> key is the result of `expr`
-- Both lower to the same AST node shape so later compiler stages only need to
-- handle one concept: "index this base expression by this key expression".
function Parser:parseIndexSuffix(tableExpression)
  local indexExpression

  -- `table.index`
  if self:advanceIfMatch("Dot") then
    indexExpression = {
      kind  = "StringLiteral",
      value = self:expectIdentifier()
    }

  -- `table[expression]`
  elseif self:advanceIfMatch("LeftBracket") then
    indexExpression = self:expectExpression()
    self:expectKind("RightBracket")

  else
    error(string.format(
      "Expected '.' or '[' for index expression, but got: %s",
      self:getTokenDescription(self.currentToken)
    ))
  end

  return { kind = "IndexExpression",
    base  = tableExpression,
    index = indexExpression
  }
end

-- `{ <table fields> }`
--
-- `<table fields>` can be:
--   - `[<expression>] = <expression>` (explicit key)
--   - `<identifier> = <expression>`
--       (syntactic sugar for `["<identifier>"] = <expression>`)
--   - `<expression>`
--       (implicit numeric key, automatically increments key)
--
-- table fields can be separated by commas or semicolons, and
-- the final field can optionally have a trailing separator.
function Parser:parseTableConstructor()
  self:expect("LeftBrace")

  local implicitKeyCounter = 0
  local tableElements      = {}

  -- Loop through tokens until we find the closing "}".
  while not self:isTokenKind("RightBrace") do
    local key, value
    local isImplicit = false

    -- Explicit key in brackets, e.g., `[1+2] = "value"`.
    -- [<expression>] = <expression>
    if self:advanceIfMatch("LeftBracket") then
      key = self:expectExpression()
      self:expectKind("RightBracket")
      self:expectKind("Equals")
      value = self:expectExpression()

    -- Identifier key, e.g., `name = "value"`.
    -- This is syntactic sugar for `["name"] = "value"`.
    -- <identifier> = <expression>
    --
    -- This is the only part of the parser that requires lookahead.
    -- We could potentially use left factoring to avoid it, but it would
    -- complicate the parser for little gain.
    elseif self:isTokenKind("Identifier") and
           self:peekNext().kind == "Equals" then
      key = { kind = "StringLiteral", value = self.currentToken.value }
      self:expectKind("Identifier")
      self:expectKind("Equals")
      value = self:expectExpression()

    -- Implicit numeric key, e.g. `"value1", MY_VAR, 42`
    -- This is syntactic sugar for `[1] = "value1", [2] = MY_VAR, [3] = 42`.
    -- <expression>
    else
      isImplicit = true
      implicitKeyCounter = implicitKeyCounter + 1
      key = { kind = "NumericLiteral", value = implicitKeyCounter }
      value = self:expectExpression()
    end

    table.insert(tableElements, {
      kind       = "TableElement",
      key        = key,
      value      = value,
      isImplicit = isImplicit
    })

    if
      -- Is there a separator for another element?
      not self:advanceIfMatch("Comma")
      and not self:advanceIfMatch("Semicolon")
      -- Or is the next token the end of the table?
      and not self:isTokenKind("RightBrace")
    then
      error(string.format(
        "Expected ',' or ';' after table element, but got: %s",
        self:getTokenDescription(self.currentToken)
      ))
    end
  end
  self:expect("RightBrace")
  return { kind = "TableConstructor", elements = tableElements }
end

-- `( <expressions> )` | `<string>` | `<table_constructor>`
--
-- Parses arguments after a callee expression, forming a function call node.
--
-- It handles the following Lua call syntaxes:
--   f(a, b)      -- Explicit call with parentheses
--   f "string"   -- Implicit call with a single string argument
--   f {table}    -- Implicit call with a single table argument
--
-- All three end up as the same AST node kind. The only difference is how the
-- argument list is parsed.
function Parser:parseCallSuffix(calleeExpression, isMethodCall)
  local curToken     = self.currentToken
  local curTokenKind = curToken.kind
  local arguments

  -- Explicit call with parentheses: `f(a, b)`.
  if curTokenKind == "LeftParen" then
    self:advance(1)
    arguments = self:parseExpressionList()
    self:expect("RightParen")

  -- Call with a single string argument: `print "hello" `
  elseif curTokenKind == "String" then
    arguments = { {
      kind  = "StringLiteral",
      value = curToken.value
    } }
    self:advance(1)

  -- Call with a single table constructor argument: `f {1, 2, 3} `.
  elseif curTokenKind == "LeftBrace" then
    arguments = { self:parseTableConstructor() }

  else
    error(string.format(
      "Expected function call arguments (parentheses, string, or table), but got: %s",
      self:getTokenDescription(curToken)
    ))
  end

  return { kind = "FunctionCall",
    callee       = calleeExpression,
    arguments    = arguments,
    isMethodCall = isMethodCall and true
  }
end

-- `:<identifier><call_suffix>`
--
-- Parses a method call expression (e.g., `object:method(args)`), and
-- returns a function call AST node.
function Parser:parseMethodCallSuffix(calleeExpression)
  self:expect("Colon")

  -- tb:method(args) is syntactic sugar for tb["method"](tb, args), so we
  -- construct an `IndexExpression` node for the method access, and then
  -- pass that as the base for a normal function call.
  local callBaseNode = {
    kind = "IndexExpression",
    base = calleeExpression,
    index = { kind = "StringLiteral", value = self:expectIdentifier() }
  }
  return self:parseCallSuffix(callBaseNode, true)
end

--// Precedence Climbing Expression Parser //--

-- The "Precedence Climbing" algorithm is a powerful technique for parsing
-- expressions with different precedence levels and associativity, without
-- needing backtracking or a complicated grammar. Below is the core of the
-- algorithm, which handles all of Lua's expression syntax correctly.
--
-- Lua 5.1 expression term hierarchy:
--
--   Call   ::= '(' <arguments> ')'
--            | <string>
--            | <table_constructor>
--   Index  ::= '.' <identifier>
--            | '[' <expression> ']'
--   Method ::= ':' <identifier> <call>
--
--   Binary  ::= <unary> {bin_op <unary>}
--   Unary   ::= un_op <unary> | <postfix>
--   Postfix ::= <primary> { <suffix> }
--   Suffix  ::= <call> | <index> | <method>
--   Primary ::= <number>
--             | <string>
--             | <identifier>
--             | <vararg>
--             | '(' <expression> ')'
--             | <anonymous_function>
--             | <table_constructor>
--
-- Note that in this notation, `{ term }` means a term can be repeated zero
-- or more times.

-- Parses a primary expression (e.g., Number, String, Variable, etc.).
-- This is the first term that actually gets parsed in an expression.
function Parser:parsePrimaryExpression()
  local currentToken = self.currentToken
  local tokenKind    = currentToken.kind

  if tokenKind == "Number" then
    self:advance(1)
    return { kind = "NumericLiteral",
      value = currentToken.value,
      raw   = currentToken.raw
    }
  elseif tokenKind == "String" then
    self:advance(1)
    return { kind = "StringLiteral",
      value = currentToken.value,
      raw   = currentToken.raw
    }
  elseif tokenKind == "Keyword" then
    local keyword = currentToken.value
    if keyword == "nil" then
      self:advance(1)
      return { kind = "NilLiteral" }
    elseif keyword == "true" or keyword == "false" then
      self:advance(1)
      return { kind = "BooleanLiteral", value = (keyword == "true") }

    -- Anonymous function, e.g., `function(arg1) ... end`.
    elseif keyword == "function" then
      self:advance(1)
      return self:parseFunctionExpression()
    end

  -- Variable, e.g., `MY_VAR`.
  elseif tokenKind == "Identifier" then
    return { kind = "Identifier", value = self:expectIdentifier() }

  -- Parenthesized expression, e.g., `(a + b)`.
  elseif tokenKind == "LeftParen" then
    self:advance(1)
    local expression = self:expectExpression()
    self:expectKind("RightParen")

    -- The "ParenthesizedExpression" node is NOT just for operator precedence.
    -- In Lua, parentheses also force a multi-return expression to adjust to
    -- a single value. e.g., `a, b = f()` is different from `a, b = (f())`.
    -- Therefore, we must keep this wrapper node in the AST.
    return { kind = "ParenthesizedExpression", expression = expression }

  -- Variadic argument (vararg) - `...`
  elseif tokenKind == "Vararg" then
    self:advance(1)
    return { kind = "VarargExpression" }

  -- Table constructor, e.g., `{ 42, "string", ["my"] = variable }`.
  elseif tokenKind == "LeftBrace" then
    return self:parseTableConstructor()
  end

  return nil
end

-- Parses a prefix expression with possible suffixes (function calls, indexing).
-- Handles chaining of calls and accesses (e.g., `obj:method().field["key"]`).
function Parser:parsePrefixExpression()
  -- <prefix> ::= <primary> { <suffix> }
  --
  -- Reads as: A prefix expression starts with a primary expression, followed by
  -- optional suffixes. The suffixes can be function calls, table accesses, or
  -- method calls, and they can be chained together indefinitely.

  local primaryExpression = self:parsePrimaryExpression()
  if not primaryExpression then return end

  -- Can't have a suffix if the primary expression is not an identifier.
  -- This check is the exact reason why you can't index strings/numbers/etc. in Lua
  -- without them being wrapped in parentheses first. E.g., `variable.index` is valid,
  -- but `"hello".index` is not - you'd need `("hello").index` instead.
  --
  -- Examples:
  --   variable.index    valid   (Identifier has a suffix)
  --   (func()).key      valid   (ParenthesizedExpression has a suffix)
  --   "hello".length    invalid (StringLiteral cannot have a suffix here)
  if
    primaryExpression.kind ~= "Identifier"
    and primaryExpression.kind ~= "ParenthesizedExpression"
  then
    return primaryExpression
  end

  while true do
    local currentToken     = self.currentToken
    local currentTokenKind = currentToken.kind

    -- Function call.
    if currentTokenKind == "LeftParen" then
      -- NOTE: Lua's grammar can be ambiguous when a function call is
      -- followed by a parenthesized expression. A semicolon is often
      -- needed to prevent a parse error. Our parser does not enforce this.
      primaryExpression = self:parseCallSuffix(primaryExpression, false)

    -- Table access.
    elseif
      currentTokenKind == "Dot"
      or currentTokenKind == "LeftBracket"
    then
      primaryExpression = self:parseIndexSuffix(primaryExpression)

    -- Method call.
    elseif currentTokenKind == "Colon" then
      primaryExpression = self:parseMethodCallSuffix(primaryExpression)

    -- Implicit function call syntax (string or table): `f "str"` / `f {}`.
    elseif
      currentTokenKind == "String"
      or currentTokenKind == "LeftBrace"
    then
      primaryExpression = self:parseCallSuffix(primaryExpression, false)
    else
      break
    end
  end

  return primaryExpression
end

function Parser:parseUnaryOperator()
  -- <unary> ::= <unary_operator> <unary> | <prefix>
  --
  -- Reads as: A unary expression is either a unary operator followed by another
  -- unary expression (e.g., `-x`, `not true`), OR a prefix expression on its own.

  if not self:isUnaryOperator() then
    return self:parsePrefixExpression()
  end

  local operator = self.currentToken.value
  self:advance(1)
  local operand = self:parseBinaryExpression(PARSER_CONFIG.UNARY_PRECEDENCE)
  if not operand then error("Unexpected end") end

  return { kind = "UnaryOperator",
    operator = operator,
    operand  = operand
  }
end

function Parser:parseBinaryExpression(previousPrecedence)
  -- <binary> ::= <unary> <binary_operator> <binary> | <unary>
  --
  -- Reads as: A binary expression is either a unary expression followed by
  -- a binary operator and another binary expression, OR just a unary expression
  -- on its own. This is a recursive definition, and the `previousPrecedence`
  -- parameter is what allows us to handle operator precedence correctly.

  local expression = self:parseUnaryOperator() -- <unary>
  if not expression then return end

  -- <binary_operator> <binary>
  while self:isBinaryOperator() do
    local operator   = self.currentToken.value
    local precedence = PARSER_CONFIG.PRECEDENCE[operator]
    if precedence[1] <= previousPrecedence then
      break
    end
    self:advance(1)

    local right = self:parseBinaryExpression(precedence[2])
    if not right then error("Unexpected end") end

    expression = { kind = "BinaryOperator",
      operator = operator,
      left     = expression,
      right    = right
    }
  end

  return expression
end

function Parser:parseExpression()
  return self:parseBinaryExpression(0)
end

-- Parse one expression and require that something valid was found.
--
-- This is useful in grammar positions where an empty list is not allowed, such
-- as the condition of an `if` or the right-hand side of an assignment.
function Parser:expectExpression()
  local expression = self:parseExpression()
  if not expression then
    error(
      "Expected an expression, but found: "
      .. self:getTokenDescription(self.currentToken)
    )
  end

  return expression
end

-- Parse a comma-separated expression list.
--
-- Returning an empty list is intentional: some grammar positions allow it,
-- such as `return` with no values or an empty call argument list `f()`.
function Parser:parseExpressionList()
  local firstExpression = self:parseExpression()
  if not firstExpression then return {} end

  local expressionList = { firstExpression }
  while self:advanceIfMatch("Comma") do
    table.insert(expressionList, self:expectExpression())
  end

  return expressionList
end

function Parser:expectOneOrMoreExpressions()
  local expressions = self:parseExpressionList()
  if #expressions == 0 then
    error(
      "Expected at least one expression, but found: "
      .. self:getTokenDescription(self.currentToken)
    )
  end

  return expressions
end

--// Statement Parsers //--

-- Parses a local variable or function declaration.
--
-- Handles both 'local var = value' and 'local function name() end' forms.
function Parser:parseLocalStatement()
  self:expectKeyword("local")

  if self:isKeyword("function") then
    self:expectKeyword("function")

    -- Unlike non-local function declarations (which desugar to assignments),
    -- local function declarations require special handling to support recursion.
    -- Consider: `local function foo() return foo() end`
    -- If we naively translated this to `local foo = function() return foo() end`,
    -- the inner `foo` would reference the global `foo` instead of the local one.
    -- The correct desugaring is: `local foo; foo = function() return foo() end`,
    -- but since our parser returns a single AST node per statement, we represent
    -- local function declarations as a dedicated node type.
    return { kind = "LocalFunctionDeclaration",
      name = self:expectIdentifier(),
      body = self:parseFunctionExpression()
    }
  end

  local variables = self:parseIdentifierList()
  local initializers = {}
  if self:advanceIfMatch("Equals") then
    initializers = self:expectOneOrMoreExpressions()
  end

  return { kind = "LocalDeclarationStatement",
    variables    = variables,
    initializers = initializers
  }
end

function Parser:parseWhileStatement()
  self:expectKeyword("while")

  return { kind = "WhileStatement",
    condition = self:expectExpression(),
    body      = self:parseDoEndLoopBlock()
  }
end

function Parser:parseRepeatStatement()
  self:expectKeyword("repeat")
  local body = self:parseCodeBlock()
  self:expectKeyword("until")

  return { kind = "RepeatStatement",
    body      = body,
    condition = self:expectExpression()
  }
end

function Parser:parseDoStatement()
  return { kind = "DoStatement", body = self:parseDoEndLoopBlock() }
end

function Parser:parseReturnStatement()
  self:expectKeyword("return")
  local expressions = self:parseExpressionList()
  self.shouldTerminateBlock = true
  return { kind = "ReturnStatement", expressions = expressions }
end

function Parser:parseBreakStatement()
  self:expectKeyword("break")
  self.shouldTerminateBlock = true
  return { kind = "BreakStatement" }
end

function Parser:parseIfStatement()
  -- Parse the first `if` clause, then any number of `elseif` clauses.
  -- The optional `else` body is stored separately because it has no condition.
  local clauses = { self:parseIfClause("if") }
  while self:isKeyword("elseif") do
    table.insert(clauses, self:parseIfClause("elseif"))
  end
  local elseClause
  if self:isKeyword("else") then
    self:expectKeyword("else")
    elseClause = self:parseCodeBlock()
  end
  self:expectKeyword("end")
  return { kind = "IfStatement",
    clauses    = clauses,
    elseClause = elseClause
  }
end

function Parser:parseForStatement()
  self:expectKeyword("for")
  local variableName = self:expectIdentifier()

  -- At this point, we must distinguish which form of `for` loop we're parsing.
  -- A generic loop looks like `for k, v in pairs(t) do ... end`, while a numeric
  -- loop looks like `for i = 1, 10 do ... end`. The presence of a comma or the
  -- `in` keyword after the first variable name signals the generic form.
  local isGenericFor = self:isTokenKind("Comma") or self:isKeyword("in")
  if isGenericFor then
    local iteratorVariables = { variableName }
    while self:advanceIfMatch("Comma") do
      table.insert(iteratorVariables, self:expectIdentifier())
    end
    self:expectKeyword("in")
    return { kind = "ForGenericStatement",
      iterators   = iteratorVariables,
      expressions = self:expectOneOrMoreExpressions(),
      body        = self:parseDoEndLoopBlock()
    }
  else
    self:expect("Equals")
    local expressions = self:expectOneOrMoreExpressions()
    if #expressions > 3 then
      error(
        "Numeric 'for' loop allows at most a start, limit, "
        .. "and optional step expression."
      )
    elseif #expressions < 2 then
      error(
        "Numeric 'for' loop requires at least a start and limit expression."
      )
    end

    return { kind = "ForNumericStatement",
      variable = variableName,
      start    = expressions[1],
      limit    = expressions[2],
      step     = expressions[3],
      body     = self:parseDoEndLoopBlock()
    }
  end
end

-- Parses the following forms:
--   function func(params) ... end
--   function tb.method(params) ... end
--   function tb:method(params) ... end
--   function tb.subtb:method(params) ... end
--
-- Desugars into an assignment statement for consistency.
function Parser:parseFunctionDeclaration()
  self:expectKeyword("function")

  -- Parse the base function name (e.g., `myFunc` in `myFunc.new`).
  local base = { kind = "Identifier", value = self:expectIdentifier() }

  -- Parse the `.field`/`:method` suffix chains, if any.
  local isMethodDeclaration = false
  while true do
    local isDot   = self:isTokenKind("Dot")
    local isColon = self:isTokenKind("Colon")
    if not (isDot or isColon) then
      break
    end
    self:advance(1) -- Consume the dot or colon.

    base = {
      kind  = "IndexExpression",
      base  = base,
      index = { kind = "StringLiteral",
        value = self:expectIdentifier()
      }
    }

    if isColon then
      isMethodDeclaration = true
      -- A method (with `:`) is always the last part of the function name.
      break
    end
  end

  local funcExpr = self:parseFunctionExpression()
  if isMethodDeclaration then
    table.insert(funcExpr.parameters, 1, "self")
  end

  -- Non-local function declarations desugar to assignment statements.
  -- Behaviorally, `function f() end` is equivalent to `f = function() end`.
  -- This uniform representation simplifies code generation by eliminating
  -- the need for a separate FunctionDeclaration node handler.
  return { kind = "AssignmentStatement",
    lvalues     = { base },
    expressions = { funcExpr }
  }
end

-- Parse the left-hand side list of an assignment after the first lvalue has
-- already been read by the caller.
function Parser:parseAssignment(firstLvalue)
  local lvalues = { firstLvalue }
  while self:advanceIfMatch("Comma") do
    local nextLValue = self:parsePrefixExpression()
    if not self:isValidAssignmentLvalue(nextLValue) then
      error(
        "Invalid assignment target: expected identifier or table index, got "
        .. tostring(nextLValue and nextLValue.kind or "nil")
      )
    end
    table.insert(lvalues, nextLValue)
  end
  self:expectKind("Equals")
  return { kind = "AssignmentStatement",
    lvalues     = lvalues,
    expressions = self:expectOneOrMoreExpressions()
  }
end

-- Parse the ambiguous "prefix-expression statement" form.
--
-- In Lua, a statement that starts with an identifier can be either a function call
-- or an assignment, and we can't know which one until we see what comes after it.
--
-- For example:
--   `f(x)`  -> call statement
--   `f = 1` -> assignment
--
-- We solve this ambiguity by first parsing the prefix expression - which can be
-- either a simple variable or a more complex expression with indexing and method
-- calls - and then checking if it forms a valid function call or assignment target.
-- (Lua's grammar is technically LL(1) but this particular form requires
-- unbounded lookahead if parsed naively - the prefix-expression approach
-- is the standard solution.)
function Parser:parseCallOrAssignment()
  local expression = self:parsePrefixExpression()

  -- Function call statement (e.g., `f(x)`, `obj:method()`, `print "hello"`).
  if expression and expression.kind == "FunctionCall" then
    -- It's already parsed as a FunctionCall node, so we just need to wrap it in
    -- a CallStatement node for the statement-level AST.
    return { kind = "CallStatement", expression = expression }

  -- Assignment statement (e.g., `a = 1`, `tbl.key = "value"`).
  elseif self:isValidAssignmentLvalue(expression) then
    return self:parseAssignment(expression)
  end

  error(string.format(
    "Invalid statement: expected an identifier or function call, but got '%s'",
    tostring(expression and expression.kind)
  ))
end

--// CODE BLOCK PARSERS //--

-- Parses the next statement in the current block and returns its AST node.
-- Complete syntax of Lua 5.1: https://www.lua.org/manual/5.1/manual.html#8
function Parser:parseNextStatement()
  local node

  local currentToken = self.currentToken
  local tokenKind    = currentToken.kind
  if tokenKind == "Keyword" then
    local keyword = currentToken.value

    -- Stop at block terminators and let the parent parser handle them.
    -- `end` / `else` / `elseif` / `until`
    if PARSER_CONFIG.TERMINATION_KEYWORDS[keyword] then
      return nil
    end

    local handler = self.statementHandlers[keyword]
    if not handler then
      error("Unsupported keyword used as a statement starter: " .. keyword)
    end
    node = handler(self) -- e.g., self:parseWhileStatement().

  elseif tokenKind == "EOF" then
    return nil
  else
    node = self:parseCallOrAssignment()
  end

  -- Statements can optionally be separated by a semicolon.
  -- In Lua 5.1, a semicolon can't appear after another semicolon.
  self:advanceIfMatch("Semicolon")

  return node
end

-- Parse a sequence of statements until we hit a block terminator such as
-- `end`, `elseif`, `else`, `until`, or EOF.
--
-- `return` and `break` are tracked with `shouldTerminateBlock` so we stop
-- parsing statements after them, matching Lua's rule that nothing meaningful
-- can appear later in the same block.
function Parser:parseCodeBlock()
  local node = { kind = "Block", statements = {} }
  local nodeList  = node.statements
  local nodeCount = 0

  while true do
    local statementNode = self:parseNextStatement()

    -- Have we reached a termination keyword or end of file (EOF)?
    if not statementNode then
      break
    end

    -- Optimization: Avoid using table.insert in a loop, which can be slow.
    nodeCount = nodeCount + 1
    nodeList[nodeCount] = statementNode

    -- `return` / `break` terminates the current block.
    if self.shouldTerminateBlock then
      self.shouldTerminateBlock = false
      break
    end
  end

  return node
end

--// MAIN //--
function Parser:parse()
  local body = self:parseCodeBlock()
  self:expect("EOF") -- Make sure we consumed all tokens.

  return { kind = "Program", body = body }
end

-- ============================================================================
--                             (づ｡◕‿‿◕｡)づ
--                         THE CODE GENERATOR!!!
-- ============================================================================
--
-- The third stage: transform AST nodes into a Lua 5.1 prototype - the compiled
-- form of a function that the VM knows how to execute. A prototype holds the
-- instruction list, the constant pool, upvalue descriptors, and any nested
-- function prototypes produced by inner functions.
--
-- If the AST says "what the program means", the prototype says "how to run it".
--
-- Lua's VM is register-based. Every expression result, local variable, and
-- temporary value lives in a numbered slot (R(0), R(1), ...). The code
-- generator allocates and frees these slots as it emits instructions. These
-- are not the registers of a physical CPU - they are just positions in a
-- per-function array that the VM manages at runtime.
--
-- Each function in the source compiles into its own prototype. Nested functions
-- produce nested prototypes, collected as children of the enclosing function's
-- prototype.
--
-- Scopes drive how variables are accessed. The generator maintains a linked
-- list of scopes, one per block. When it encounters an identifier, it walks
-- outward: if the variable is found in the current function it is a local,
-- if it is found in an enclosing function it is an upvalue, and if it is not
-- found at all it is a global.
--
-- When an inner function reads a variable from an outer scope, that variable
-- becomes an upvalue. The generator tracks which variables are captured and
-- emits the MOVE/GETUPVAL pseudo-instructions that wire up the closure at
-- runtime.
--
-- Example (continuing from the parser's `local x = 1 + 2 * 3`):
--
--   Proto:
--     constants: [1, 2, 3]            -- K(0)=1, K(1)=2, K(2)=3
--     upvalues:  0
--     numParams: 0
--     isVararg:  true                 -- the main chunk always accepts varargs
--     maxStack:  2                    -- R(0) and R(1) are both needed
--     code:
--       [1] MUL    R(1), K(1), K(2)   ; R(1) = 2 * 3
--       [2] ADD    R(0), K(0), R(1)   ; R(0) = 1 + R(1)
--       [3] RETURN 0,   1             ; default return (no values)
--
-- Notice that 2 and 3 are used directly as constant RK operands - no register
-- is spent loading them. Only the intermediate product needs a temporary slot
-- (R(1)), which is freed as soon as ADD consumes it. The local `x` maps to
-- R(0), which is exactly where the final result already lives.

-- Maximum amount of table elements for a single SETLIST instruction.
-- If more elements are needed, multiple SETLIST instructions are emitted.
local SETLIST_MAX = 50

local CODE_GENERATOR_CONFIG = {
  -- Stack Limits.
  -- Lua 5.1 allows up to 255 registers per function. We reserve a few for
  -- internal safety, capping user registers slightly lower.
  MIN_STACK_SIZE = 2,
  MAX_REGISTERS  = 250,

  -- Node kinds that can return multiple values (multret).
  -- These require special handling when they appear as the last element
  -- of a list (e.g., `return f()`).
  MULTRET_NODES = createLookupTable({"FunctionCall", "VarargExpression"}),

  -- Short-circuit operators. Unlike arithmetic, these don't just
  -- compute a value -- they conditionally skip evaluation of the
  -- right-hand operand, which requires jump instructions.
  LOGICAL_OPERATOR_LOOKUP = createLookupTable({"and", "or"}),

  -- Maps Lua unary operator tokens to their VM opname names.
  UNARY_OPERATOR_LOOKUP = {
    ["-"]   = "UNM",
    ["#"]   = "LEN",
    ["not"] = "NOT"
  },

  -- Maps Lua arithmetic operator tokens to their VM opname names.
  -- All six use the same instruction format:
  --   OP [A, B, C]  R(A) := RK(B) op RK(C)
  -- where B and C can be either registers or constant indices (RK operands).
  ARITHMETIC_OPERATOR_LOOKUP = {
    ["+"] = "ADD", ["-"] = "SUB",
    ["*"] = "MUL", ["/"] = "DIV",
    ["%"] = "MOD", ["^"] = "POW"
  },

  -- Maps Lua comparison operators to their VM instruction and the `A` operand
  -- value used to select the comparison sense (true/false branch).
  --
  -- Lua's VM only has EQ, LT, and LE instructions. The "opposite" operators
  -- (>, >=, ~=) are implemented by:
  --   - ~=: Uses EQ with A=0 (inverts the equality test)
  --   - >, >=: Uses LT/LE with operands B and C swapped
  --
  -- For > and >=, the operand ORDER is also flipped: `a > b` is compiled as
  -- `b < a` by swapping which value goes into the B field and which into C.
  -- This is why the table maps both ">" and "<" to "LT" with the same A=1:
  -- the difference is entirely in the swap, which is handled in BinaryOperator.
  --
  -- Format: [operator] = {instruction_name, A_operand}
  COMPARISON_INSTRUCTION_LOOKUP = {
    ["=="] = {"EQ", 1}, ["~="] = {"EQ", 0},
    ["<"]  = {"LT", 1}, [">"]  = {"LT", 1},
    ["<="] = {"LE", 1}, [">="] = {"LE", 1}
  }
}

--// Utility Functions //--

-- Packs an 8-bit table array operand into the "exponential" format used
-- by the Lua VM for the NEWTABLE instruction. This allows us to encode
-- larger table sizes than would fit in a single byte, at the cost of
-- some precision for larger sizes (e.g., 17 and 18 both encode to 18).
--
-- Reference: https://www.lua.org/source/5.1/lobject.c.html#luaO_int2fb
local function encodeTableSize(size)
  -- This is not the real limit; it was chosen arbitrarily to prevent
  -- potential overflows in instruction operands.
  if size >= 65536 then
    error("Table size too large to encode: " .. size)
  end

  local exponent = 0
  while size >= 16 do
    size = math.floor((size + 1) / 2)
    exponent = exponent + 1
  end

  if size < 8 then
    return size
  end
  return ((exponent + 1) * 8) + (size - 8)
end

local CodeGenerator = {}
CodeGenerator.__index = CodeGenerator

function CodeGenerator.new(ast)
  assert(
    type(ast) == "table",
    "Expected 'ast' argument to be a table, got: " .. type(ast)
  )
  assert(
    ast.kind == "Program",
    "Expected 'ast' to be a 'Program' node, got: " .. ast.kind
  )

  return setmetatable({
    ast   = ast,
    proto = nil,

    currentScope      = nil,
    pendingBreakJumps = nil
  }, CodeGenerator)
end

--// Scope Management //--

-- By the time we leave a function scope, every temporary register should have
-- been freed. That means the only live registers left should be the ones owned
-- by declared locals in that function scope.
function CodeGenerator:validateRegisterUsage()
  local currentScope = self.currentScope
  local currentRegisters = currentScope.allocatedRegisters
  local variableCount = 0
  for _ in pairs(currentScope.locals) do
    variableCount = variableCount + 1
  end

  if currentRegisters ~= variableCount then
    error(string.format(
      "Register allocation mismatch: allocated %d registers but have %d "
      .. "variables in scope. This is a compiler bug.",
      currentRegisters,
      variableCount
    ))
  end
end

function CodeGenerator:enterScope(isFunctionScope)
  local parentScope = self.currentScope
  local allocatedRegisters = 0
  if parentScope and not isFunctionScope then
    allocatedRegisters = parentScope.allocatedRegisters
  end

  self.currentScope = {
    parentScope        = parentScope,
    locals             = {},
    isFunction         = isFunctionScope,
    needClose          = false,
    allocatedRegisters = allocatedRegisters
  }
end

-- Exits the current scope, returning to the parent scope.
--
-- If any variable in this scope was captured as an upvalue
-- (scope.needClose=true), emits OP_CLOSE to transition those
-- upvalues from "open" (pointing into the register window) to
-- "closed" (owning their own copy of the value). Without this,
-- the upvalue would point to a register that gets reused, causing
-- stale reads.
function CodeGenerator:leaveScope()
  local currentScope = self.currentScope
  if not currentScope then
    error("CodeGenerator: No scope to leave!")
  end

  if currentScope.isFunction then
    self:validateRegisterUsage()
  end

  self.currentScope = currentScope.parentScope

  -- Function scopes don't need OP_CLOSE here because the VM handles
  -- closing all upvalues automatically when a function returns.
  local needClose = currentScope.needClose and not currentScope.isFunction
  if needClose then
    -- If we're inside a loop, mark the break jumps as needing
    -- OP_CLOSE too, since break exits the scope early.
    if self.pendingBreakJumps then
      self.pendingBreakJumps.needClose = true
    end

    -- CLOSE expects the first register that belongs to the scope being left.
    -- Once we restore the parent scope, its allocated register count is
    -- exactly that boundary.
    -- OP_CLOSE [A]    close all variables in the stack up to (>=) R(A)
    self:emitInstruction("CLOSE", self.currentScope.allocatedRegisters, 0, 0)
  end
end

--// Variable Management //--
-- Variables are mapped to registers in their declaring scope.
-- Looking up a variable means searching scopes from innermost
-- to outermost, stopping at function boundaries (since a different
-- function means the variable is an upvalue, not a local).

-- Binds a local variable name to a register in the current scope.
-- If no register is specified, allocates the next available one.
function CodeGenerator:declareLocalVariable(varName, register)
  register = register or self:allocateRegisters(1)
  self.currentScope.locals[varName] = register
  return register
end

function CodeGenerator:declareLocalVariables(varNames)
  for _, varName in ipairs(varNames) do
    self:declareLocalVariable(varName)
  end
end

function CodeGenerator:getLocalRegister(varName)
  local scope = self.currentScope
  while scope do
    local register = scope.locals[varName]
    if register then
      return register
    elseif scope.isFunction then
      break
    end
    scope = scope.parentScope
  end

  -- This error should never happen. If it does, it's most likely
  -- due to a bug in the `CodeGenerator:getVariableType()` function.
  error(string.format(
    "Internal error: Assumed local %s is not found. This should never happen.",
    varName
  ))
end

function CodeGenerator:getVariableType(variableName, markScopeForClose)
  local scope, isUpvalue = self.currentScope, false

  while scope do
    if scope.locals[variableName] then
      if markScopeForClose then
        scope.needClose = true
      end
      return (isUpvalue and "Upvalue") or "Local"
    elseif scope.isFunction then
      -- Crossing a function boundary changes the meaning of the next match:
      -- from this point on, any found local belongs to an outer function and
      -- must therefore be captured as an upvalue.
      isUpvalue = true
    end
    scope = scope.parentScope
  end

  return "Global"
end

--// Prototype Management //--

-- A prototype is the "compiled template" of a function. It holds
-- everything the VM needs to create a closure and execute the function:
--   - code:      List of instructions
--   - constants: Pool of literal values (numbers, strings)
--   - upvalues:  Names of captured variables from enclosing scopes
--   - protos:    Nested function prototypes (children)
function CodeGenerator:createPrototype(properties)
  return {
    code      = {},
    constants = {},
    protos    = {},

    -- Tracks the maximum register index used by this function, which is used
    -- in the Lua VM to preallocate the most optimal amount of registers for
    -- the function. We could potentially always just set it to MAX_STACK_SIZE,
    -- but dynamically tracking the actual maximum allows us to generate more
    -- efficient bytecode that uses less memory and can run faster on the VM.
    maxStackSize = CODE_GENERATOR_CONFIG.MIN_STACK_SIZE,

    numUpvals = 0, -- Will be set during finalization to #proto.upvalues.
    numParams = properties.numParams or 0,
    isVararg  = properties.isVararg  or false,

    -- Internal lookup tables to avoid O(n) searches when referencing constants
    -- and upvalues by name. These will get removed during finalization, since
    -- the VM only needs the lists, not the lookups.
    constantLookup = {}, -- format: {[constant] = index in proto.constants}
    upvalueLookup  = {}, -- format: {[varName]  = index in proto.upvalues}
    upvalues       = {}  -- format: {varName1, varName2, ...}
  }
end

--// Register/Constant/Upvalue Management //--
-- These functions manage the three kinds of operands that instructions
-- reference: registers (local slots), constants (literal values), and
-- upvalues (captured variables from enclosing functions).

-- Allocates `count` consecutive registers from the current scope's
-- register window. Returns the index of the LAST allocated register
-- (e.g., if allocatedRegisters was 3 and count is 2, this returns 4 and
-- updates allocatedRegisters to 5).
function CodeGenerator:allocateRegisters(count)
  if count <= 0 then return end -- Just ignore non-positive counts.

  local currentScope      = self.currentScope
  local previousStackSize = currentScope.allocatedRegisters
  local updatedStackSize  = previousStackSize + count
  if updatedStackSize > CODE_GENERATOR_CONFIG.MAX_REGISTERS then
    error(string.format(
      "CodeGenerator: Register overflow! exceeded maximum of %d registers.",
      CODE_GENERATOR_CONFIG.MAX_REGISTERS
    ))
  end

  currentScope.allocatedRegisters = updatedStackSize

  -- Read the note in `:createPrototype()` about maxStackSize.
  if updatedStackSize >= self.proto.maxStackSize then
    self.proto.maxStackSize = updatedStackSize
  end

  return updatedStackSize - 1
end

function CodeGenerator:freeRegisters(count)
  if not count or count <= 0 then return end

  local currentScope = self.currentScope
  local newStackSize = currentScope.allocatedRegisters - count
  if newStackSize < 0 then
    error("CodeGenerator: Stack underflow! cannot free below minimum of 0 registers.")
  end

  currentScope.allocatedRegisters = newStackSize
end

-- Frees an RK operand if it's a register, ignores if it's a constant.
--
-- Throughout the code generator, we use a convention where:
--   >= 0  means a register index (needs to be freed after use)
--   <  0  means a constant index (lives in the constant pool, no cleanup)
--
-- This helper avoids the need to check "is this a register or constant?"
-- at every call site where we use compileConstantOrExpression.
function CodeGenerator:freeIfRegister(rkOperand)
  -- Is this a constant? We can't free it.
  if rkOperand < 0 then return end

  -- Register freeing must happen in strict LIFO (last-in-first-out) order to
  -- avoid accidentally freeing registers that are still in use. Only the most
  -- recently allocated register can be freed at any given time.
  if rkOperand ~= self.currentScope.allocatedRegisters - 1 then
    error(string.format(
      "CodeGenerator: Attempting to free register R(%d) but only the "
      .. "most recently allocated register R(%d) can be freed. "
      .. "This is a compiler bug, not a user error.",
      rkOperand,
      self.currentScope.allocatedRegisters - 1
    ))
  end

  self:freeRegisters(1)
end

-- Finds or creates a constant pool entry for the given value.
-- Returns the negative index of the constant in the proto.constants list.
--
-- Note:
--  Constant indices are returned as negative numbers (e.g., -1 for the first
--  constant) to distinguish them from register indices, which are non-negative.
--  This mirrors how Lua 5.1's VM encodes RK (Register or Konstant) operands in
--  bytecode: registers use values 0-255, while constants are encoded as 256
--  plus the absolute constant index. By using negative indices internally,
--  TLC simplifies operand handling without losing compatibility.
function CodeGenerator:internConstant(value)
  local constantLookup = self.proto.constantLookup
  local constants      = self.proto.constants
  local constantIndex  = constantLookup[value]

  if constantIndex then
    return constantIndex
  end

  table.insert(constants, value)
  constantIndex = -#constants -- Constant indices are negative-based.
  constantLookup[value] = constantIndex
  return constantIndex
end

-- Finds or creates an upvalue entry for the given variable name.
-- Returns the 0-based upvalue index used by OP_GETUPVAL / OP_SETUPVAL.
--
-- Upvalue indices are separate from register indices. Each function
-- prototype has its own upvalue list, and the indices refer to slots
-- in that list. The CLOSURE pseudo-instructions (see generateClosure)
-- populate these slots at runtime.
function CodeGenerator:internUpvalue(varName)
  local upvalueLookup = self.proto.upvalueLookup
  local upvalues      = self.proto.upvalues
  local upvalueIndex  = upvalueLookup[varName]

  if upvalueIndex then
    return upvalueIndex
  end

  table.insert(upvalues, varName)

  -- Lua upvalue indices are 0-based, table indices are 1-based.
  -- Therefore, we subtract 1 from the length.
  upvalueIndex = #upvalues - 1
  upvalueLookup[varName] = upvalueIndex
  return upvalueIndex
end

-- Tries to encode a node as a constant index (RK optimization).
--
-- Many Lua instructions accept "RK" operands that can be either a
-- register or a constant pool index. If the node is a simple literal
-- (number or string) that fits in the 9-bit RK field, we return its
-- constant index directly, avoiding a register allocation entirely.
--
-- If it can't be encoded as a constant (complex expression, or the
-- constant pool is too large), falls back to evaluating the expression
-- into a register. The caller must call freeIfRegister on the result.
function CodeGenerator:compileConstantOrExpression(node, register)
  local nodeKind = node.kind

  if nodeKind == "NumericLiteral" or nodeKind == "StringLiteral" then
    local constantIndex = self:internConstant(node.value)

    -- Before returning the constant register we must check if it can fit in an
    -- 8-bit unsigned operand (0-255) first, if not, we put it in a register instead.
    if -constantIndex < 255 then
      -- Return the constant index (which is negative) directly as an RK operand.
      return constantIndex
    end
  end

  return self:compileExpressionNode(node, register)
end

--// Instruction Emission //--
function CodeGenerator:emitInstruction(opname, a, b, c)
  if not a then
    error("CodeGenerator: Missing 'a' operand for instruction " .. opname)
  elseif not b then
    error("CodeGenerator: Missing 'b' operand for instruction " .. opname)
  end

  -- TLC stores instructions in a normalized table form first, then the
  -- BytecodeEmitter later packs them into their final binary encoding.
  --
  -- `b` is always present in this normalized form. For iABx/iAsBx
  -- instructions it simply carries the wider Bx/sBx payload.
  local instruction = {
    opname = opname,
    a      = a,
    b      = b,
    c      = c or 0
  }
  table.insert(self.proto.code, instruction)
  return instruction
end

function CodeGenerator:emitJump(offset)
  -- OP_JMP [A, sBx]    pc+=sBx
  self:emitInstruction("JMP", 0, offset or 1, 0)
  return #self.proto.code
end

function CodeGenerator:emitJumpBack(relativeDest)
  self:emitJump(relativeDest - (#self.proto.code + 1))
end

function CodeGenerator:patchJump(fromPC, toPC)
  local instruction = self.proto.code[fromPC]
  if not instruction then
    error(
      "CodeGenerator: Invalid 'fromPC' for jump patching: "
      .. tostring(fromPC)
    )
  end

  -- NOTE: `JMP` uses the `sBx` operand for the jump offset, but in our table
  -- it is represented as the `b` field for simplicity.
  instruction.b = toPC - (fromPC + 1)
end

function CodeGenerator:patchJumpToHere(fromPC)
  self:patchJump(fromPC, #self.proto.code + 1)
end

function CodeGenerator:patchJumpsToHere(jumpPCs)
  if not jumpPCs then return end
  local currentPC = #self.proto.code + 1

  for _, pc in ipairs(jumpPCs) do
    self:patchJump(pc, currentPC)
  end
end

--// Loop Helpers //--

-- When the code generator encounters a `break` inside a loop, it
-- emits a forward jump to an unknown location (the loop's exit).
-- This function collects all such break jumps, then patches them
-- to point to the instruction immediately after the loop body.
--
-- If any scope inside the loop captured upvalues (needClose), an
-- OP_CLOSE is emitted at the break target to clean them up.
function CodeGenerator:breakable(callback)
  local previousPendingJumps = self.pendingBreakJumps
  self.pendingBreakJumps = {}

  callback()

  -- Patch breaks before emitting CLOSE so they land exactly
  -- on the cleanup instruction when one is needed.
  self:patchJumpsToHere(self.pendingBreakJumps)
  if #self.pendingBreakJumps > 0 and self.pendingBreakJumps.needClose then
    -- OP_CLOSE [A]    close all variables in the stack up to (>=) R(A)
    self:emitInstruction("CLOSE", self.currentScope.allocatedRegisters, 0, 0)
  end

  self.pendingBreakJumps = previousPendingJumps
end

--// Utility Methods //--

-- In Lua, only the LAST expression in a list can expand to multiple
-- values. Earlier expressions are always adjusted to one value.
-- These helpers detect and handle that distinction.
--
-- Examples:
--   return f(), g()     -- f() adjusted to 1 value, g() expands fully
--   local a, b = f()    -- f() expands to fill both a and b
--   local a, b = f(), 1 -- f() adjusted to 1 value, b gets 1

-- Checks if a node can produce multiple return values.
-- Only function calls and vararg expressions (...) can do this.
function CodeGenerator:isMultiReturnNode(node)
  return node and CODE_GENERATOR_CONFIG.MULTRET_NODES[node.kind]
end

function CodeGenerator:isMultiReturnList(expressions)
  if #expressions == 0 then return false end

  local lastExpression = expressions[#expressions]
  return self:isMultiReturnNode(lastExpression)
end

-- Used to detect tail call optimization opportunities, e.g.: `return f()`.
function CodeGenerator:isSingleFunctionCall(expressions)
  if #expressions ~= 1 then return false end

  local expression = expressions[1]
  return expression.kind == "FunctionCall"
end

-- Recursively assigns values to lvalues.
--
-- Assignments like `a, a.x = 1, 2` require evaluating `a` for the table index
-- before `a` is overwritten with `1`. This recursion ensures that base and
-- index expressions for all assignments are evaluated before the actual
-- writes occur.
function CodeGenerator:assignValuesToRegisters(nodes, index, copyFromRegister)
  if index > #nodes then return end
  local node     = nodes[index]
  local nodeKind = node.kind

  if nodeKind == "Identifier" then
    -- First give the other index assignments a chance to read
    -- the base and index before it might change.
    self:assignValuesToRegisters(nodes, index + 1, copyFromRegister + 1)

    local variableName = node.value
    local variableType = self:getVariableType(variableName)
    if variableType == "Local" then
      -- OP_MOVE [A, B]    R(A) := R(B)
      self:emitInstruction(
        "MOVE",
        self:getLocalRegister(variableName),
        copyFromRegister
      )
    elseif variableType == "Upvalue" then
      -- OP_SETUPVAL [A, B]    UpValue[B] := R(A)
      self:emitInstruction(
        "SETUPVAL",
        copyFromRegister,
        self:internUpvalue(variableName)
      )
    elseif variableType == "Global" then
      -- OP_SETGLOBAL [A, Bx]    Gbl[Kst(Bx)] := R(A)
      self:emitInstruction(
        "SETGLOBAL",
        copyFromRegister,
        self:internConstant(variableName)
      )
    end

    return
  elseif nodeKind == "IndexExpression" then
    local baseNode  = node.base
    local indexNode = node.index

    local baseRegister  = self:compileConstantOrExpression(baseNode)
    local indexRegister = self:compileConstantOrExpression(indexNode)

    -- This index assignment did read its base and index.
    -- Now give other index assignments a chance before doing the assignment.
    self:assignValuesToRegisters(nodes, index + 1, copyFromRegister + 1)

    -- OP_SETTABLE [A, B, C]    R(A)[RK(B)] := RK(C)
    self:emitInstruction("SETTABLE", baseRegister, indexRegister, copyFromRegister)
    self:freeIfRegister(indexRegister)
    self:freeIfRegister(baseRegister)

    return
  end

  error(
    "CodeGenerator: Unsupported lvalue kind in assignValuesToRegisters: "
    .. nodeKind
  )
end

-- Processes one SETLIST "page" of implicit table elements.
--
-- Lua's SETLIST instruction uses the C operand to encode which "page"
-- of 50 elements to write (C=1 means indices 1-50, C=2 means 51-100,
-- etc.). Large tables need multiple SETLIST instructions, one per page.
--
-- If the very last element is a multret (e.g., `{f()}`), we pass B=0
-- to SETLIST, which means "take all values up to the stack top."
function CodeGenerator:compileTablePage(
  startFromRegister,
  implicitElems,
  page,
  totalPages,
  isLastElemImplicit
)
  local startIndex   = (page - 1) * SETLIST_MAX + 1
  local endIndex     = math.min(page * SETLIST_MAX, #implicitElems)
  local elementCount = endIndex - startIndex + 1
  local isLastPage   = (page == totalPages)
  local lastElement  = implicitElems[#implicitElems]

  -- Determine if the last element of this batch is the absolute last element
  -- of the table constructor AND if it should be treated as a multret.
  local isMultret = isLastPage
                  and isLastElemImplicit
                  and self:isMultiReturnNode(lastElement.value)

  for elemIdx = startIndex, endIndex do
    local isLastInBatch = (elemIdx == endIndex)

    -- Only the absolute last element gets -1 (multret), others get 1.
    local returnCount  = (isLastInBatch and isMultret and -1) or 1
    local elemValue    = implicitElems[elemIdx].value
    local elemRegister = self:allocateRegisters(1)

    self:compileExpressionNode(elemValue, elemRegister, returnCount)
  end

  -- B = 0 means "take all values from stack top" (used for multret).
  -- Otherwise, B is the number of elements to set.
  local setListCount = (isMultret and 0) or elementCount

  -- OP_SETLIST [A, B, C]    R(A)[(C-1)*FPF+i] := R(A+i), 1 <= i <= B
  -- FPF stands for "fields per flush" and is 50 in Lua 5.1.
  self:emitInstruction("SETLIST", startFromRegister, setListCount, page)
  self:freeRegisters(elementCount)
end

-- Splits table constructor elements into two lists:
--   - implicitElems: positional entries like {1, 2, 3} (use SETLIST)
--   - explicitElems: keyed entries like {a=1, [2]=3} (use SETTABLE)
--
-- They are processed differently because SETLIST can batch-assign
-- consecutive integer keys efficiently, while SETTABLE handles
-- one key-value pair at a time.
function CodeGenerator:splitTableElements(elements)
  local implicitElems = {}
  local explicitElems = {}
  for _, element in ipairs(elements) do
    if element.isImplicit then
      table.insert(implicitElems, element)
    else
      table.insert(explicitElems, element)
    end
  end

  return implicitElems, explicitElems
end

--// Expression Handlers //--
-- Each method below handles one AST expression node kind.
-- Convention: receives a `register` to store the result in,
-- and returns that register. The caller is responsible for
-- allocating the register beforehand (compileExpressionNode
-- does this automatically if no register is provided).

---@diagnostic disable-next-line: unused-local
function CodeGenerator:NilLiteral(node, register)
  -- OP_LOADNIL [A, B]    R(A) := ... := R(B) := nil
  self:emitInstruction("LOADNIL", register, register, 0)
end

function CodeGenerator:BooleanLiteral(node, register)
  local value = (node.value and 1) or 0
  self:emitInstruction("LOADBOOL", register, value)
end

function CodeGenerator:NumericLiteral(node, register)
  self:emitInstruction("LOADK", register, self:internConstant(node.value))
end

function CodeGenerator:StringLiteral(node, register)
  self:emitInstruction("LOADK", register, self:internConstant(node.value))
end

function CodeGenerator:ParenthesizedExpression(node, register)
  return self:compileExpressionNode(node.expression, register)
end

-- Binary operators are split into four categories, each with
-- different code generation strategies:
--   1. Arithmetic (+, -, *, /, %, ^)     -> Single RK instruction
--   2. Short-circuit (and, or)           -> TEST + conditional jump
--   3. Comparison (==, ~=, <, >, <=, >=) -> Compare + LOADBOOL pair
--   4. Concatenation (..)                -> Flatten chain + CONCAT
function CodeGenerator:BinaryOperator(node, register)
  local operator = node.operator
  local left     = node.left
  local right    = node.right

  -- Arithmetic operators (+, -, /, *, %, ^)
  -- These map directly to a single VM instruction with RK operands.
  if CODE_GENERATOR_CONFIG.ARITHMETIC_OPERATOR_LOOKUP[operator] then
    local opname = CODE_GENERATOR_CONFIG.ARITHMETIC_OPERATOR_LOOKUP[operator]

    local leftOperand  = self:compileConstantOrExpression(left, register)
    local rightOperand = self:compileConstantOrExpression(right)
    self:emitInstruction(opname, register, leftOperand, rightOperand)
    self:freeIfRegister(rightOperand)

  -- Short-circuit operators (and, or)
  --
  -- These skip evaluating the right operand if the left operand's value
  -- determines the result.
  --
  -- Implementation:
  --   1. Evaluate left operand into `register`.
  --   2. TEST if short-circuit should trigger; skip next JMP if so.
  --   3. JMP over right operand evaluation.
  --   4. Evaluate right operand into `register` (overwrites left).
  elseif CODE_GENERATOR_CONFIG.LOGICAL_OPERATOR_LOOKUP[operator] then
    self:compileExpressionNode(left, register)
    local booleanValue = (operator == "and" and 0) or 1
    -- OP_TEST [A, C]    if not (R(A) <=> C) then pc++
    self:emitInstruction("TEST", register, 0, booleanValue)
    local jumpPC = self:emitJump()
    self:compileExpressionNode(right, register)
    self:patchJumpToHere(jumpPC)

  -- Comparison operators (==, ~=, <, >, <=, >=)
  --
  -- Lua's comparison instructions (EQ, LT, LE) skip the next instruction
  -- if the condition is false. We use a JMP and two LOADBOOLs to
  -- materialize a boolean result:
  --
  --   EQ/LT/LE A, B, C     ; if (RK(B) op RK(C)) ~= A then skip
  --   JMP +1               ; Condition was false, jump to "false" LOADBOOL
  --   LOADBOOL R, 0, 1     ; R := false, skip next instruction
  --   LOADBOOL R, 1, 0     ; R := true
  elseif CODE_GENERATOR_CONFIG.COMPARISON_INSTRUCTION_LOOKUP[operator] then
    local leftReg      = self:compileConstantOrExpression(left, register)
    local rightReg     = self:compileConstantOrExpression(right)
    local nodeOperator = CODE_GENERATOR_CONFIG.COMPARISON_INSTRUCTION_LOOKUP[operator]
    local instruction, opposite = nodeOperator[1], nodeOperator[2]

    -- The Lua VM doesn't have separate instructions for `>` and `>=`.
    -- Instead, it flips the operands and uses the opposite comparison sense
    -- (e.g., `a > b` becomes `b < a`).
    local flip = (operator == ">" or operator == ">=")
    local b, c = leftReg, rightReg
    if flip then b, c = c, b end

    -- OP_EQ / OP_LT / OP_LE [A, B, C]    if ((RK(B) op RK(C)) ~= A) then pc++
    self:emitInstruction(instruction, opposite, b, c)
    self:emitJump(1) -- Skip next instruction if comparison is false.

    -- OP_LOADBOOL [A, B, C]    R(A) := (Bool)B; if (C) pc++
    self:emitInstruction("LOADBOOL", register, 0, 1) -- R = false (skip next)
    self:emitInstruction("LOADBOOL", register, 1, 0) -- R = true
    self:freeIfRegister(rightReg)

  -- String concatenation (..)
  --
  -- Optimization: the parser builds `a .. b .. c` as a right-leaning
  -- tree: `(a .. (b .. c))`. Naively, each `..` would emit its own
  -- CONCAT instruction. Instead, we flatten the chain by walking
  -- right-hand `..` nodes, loading all operands into consecutive
  -- registers, then emitting a single CONCAT over the full range.
  --
  -- Example: "Hello, " .. "world" .. "!"
  --   LOADK  0, -1     ; "Hello, "
  --   LOADK  1, -2     ; "world"
  --   LOADK  2, -3     ; "!"
  --   CONCAT 3, 0, 2   ; R(3) = R(0) .. R(1) .. R(2)
  elseif operator == ".." then
    local bottomRegister = self:compileExpressionNode(left, register)

    -- Walk down the right-leaning chain, loading each left operand.
    local curRight = right
    while curRight.operator == ".." do
      self:compileExpressionNode(curRight.left)
      curRight = curRight.right
    end

    -- Load the final (rightmost) operand.
    local topRegister = self:compileExpressionNode(curRight)

    -- OP_CONCAT [A, B, C]    R(A) := R(B).. ... ..R(C)
    self:emitInstruction("CONCAT", register, bottomRegister, topRegister)
    self:freeRegisters(topRegister - bottomRegister)
  else
    error("CodeGenerator: Unsupported binary operator: " .. operator)
  end
end

function CodeGenerator:UnaryOperator(node, register)
  local opname = CODE_GENERATOR_CONFIG.UNARY_OPERATOR_LOOKUP[node.operator]
  if not opname then
    error("CodeGenerator: Unsupported unary operator: " .. node.operator)
  end

  self:compileExpressionNode(node.operand, register)

  -- OP_UNM / OP_NOT / OP_LEN [A, B]    R(A) := op R(B)
  self:emitInstruction(opname, register, register)
end

-- Tables are built in two phases. Explicit key-value pairs like
-- `{key="val", [expr]=val}` are each emitted with their own SETTABLE
-- instruction. Consecutive positional elements like `{val1, val2, val3}`
-- are batched into pages of up to 50 entries and emitted with SETLIST.
--
-- The last implicit element can be a multret (e.g., {f()}) which
-- expands to fill the remaining table slots.
function CodeGenerator:TableConstructor(node, register)
  local elements    = node.elements
  local lastElement = node.elements[#node.elements]
  local isLastElementImplicit = lastElement and lastElement.isImplicit
  local implicitElems, explicitElems = self:splitTableElements(elements)

  -- OP_NEWTABLE [A, B, C]    R(A) := {} (size = B,C)
  self:emitInstruction(
    "NEWTABLE",
    register,
    encodeTableSize(#explicitElems),
    encodeTableSize(#implicitElems)
  )

  for _, elem in ipairs(explicitElems) do
    local keyRegister   = self:compileExpressionNode(elem.key)
    local valueRegister = self:compileExpressionNode(elem.value)

    -- OP_SETTABLE [A, B, C]    R(A)[RK(B)] := RK(C)
    self:emitInstruction("SETTABLE", register, keyRegister, valueRegister)
    self:freeRegisters(2)
  end

  local totalPages = math.ceil(#implicitElems / SETLIST_MAX)

  -- Table page limit in SETLIST instructions.
  -- The C operand is 9 bits, limiting it to 511 (1-511 valid range).
  -- This caps table constructors at 511 * SETLIST_MAX elements (25550 with
  -- SETLIST_MAX=50). Official Lua 5.1 handles overflows by setting C=0 and
  -- encoding the page count in the next pseudo-instruction. TLC errors out
  -- instead for simplicity.
  if totalPages >= 512 then
    error("CodeGenerator: Table page overflow! Maximum of 511 pages allowed.")
  end

  for page = 1, totalPages do
    self:compileTablePage(
      register,
      implicitElems,
      page,
      totalPages,
      isLastElementImplicit
    )
  end
end

function CodeGenerator:Identifier(node, register)
  local varName = node.value
  local varType = self:getVariableType(varName)

  if varType == "Local" then
    -- OP_MOVE [A, B]    R(A) := R(B)
    self:emitInstruction("MOVE", register, self:getLocalRegister(varName))
  elseif varType == "Global" then
    -- OP_GETGLOBAL [A, Bx]    R(A) := Gbl[Kst(Bx)]
    self:emitInstruction("GETGLOBAL", register, self:internConstant(varName))
  elseif varType == "Upvalue" then
    -- OP_GETUPVAL [A, B]    R(A) := UpValue[B]
    self:emitInstruction("GETUPVAL", register, self:internUpvalue(varName))
  else
    error(string.format(
      "CodeGenerator: Unknown variable type for '%s': %s. This should never happen.",
      varName,
      varType
    ))
  end
end

-- Register layout for a call `f(a, b)`:
--   R(A)   = function to call
--   R(A+1) = first argument (or `self` for method calls)
--   R(A+2) = second argument
--   ...
--   R(A+C-2) ... R(A) = return values after the call
--
-- The B operand encodes argument count (0 = all results).
-- The C operand encodes expected return count (0 = all results).
function CodeGenerator:FunctionCall(node, register, resultRegisters)
  local callee       = node.callee
  local arguments    = node.arguments
  local isMethodCall = node.isMethodCall

  -- Method calls (e.g., `obj:method()`) use the SELF instruction,
  -- which does two things in one step:
  --   1. Copies the table (obj) into R(A+1) as the implicit `self` arg
  --   2. Loads the method function (obj["method"]) into R(A)
  if isMethodCall then
    self:compileExpressionNode(callee.base, register)
    self:allocateRegisters(1) -- Used for implicit `self` argument.
    local calleeIndexRegister = self:compileConstantOrExpression(callee.index)

    -- OP_SELF [A, B, C]    R(A+1) := R(B); R(A) := R(B)[RK(C)]
    -- Load the table into R(A+1) and the method into R(A) in one instruction.
    self:emitInstruction("SELF", register, register, calleeIndexRegister)
    self:freeIfRegister(calleeIndexRegister)
  else
    self:compileExpressionNode(callee, register)
  end

  local argumentRegisterCount, numArgs = self:compileExpressionList(arguments)
  local isArgsMultret = self:isMultiReturnList(arguments)
  if isMethodCall then
    -- Make sure the `CALL` instruction captures the "self" (table) param.
    argumentRegisterCount = argumentRegisterCount + 1
    if not isArgsMultret then
      -- Check if it's not multret, as multret already includes all args.
      -- Increasing it by one again in `argumentCount` would turn multret into a
      -- fixed count (1), which will make it capture zero args instead of all.
      numArgs = numArgs + 1
    end
  end

  -- Convert `numArgs` from our representation (where -1 = MULTRET, >= 0 = fixed count)
  local argumentCount    = numArgs + 1
  local returnValueCount = (resultRegisters and resultRegisters + 1) or 2
  local resultCount      = (register + returnValueCount) - 2

  -- OP_CALL [A, B, C]    R(A), ... ,R(A+C-2) := R(A)(R(A+1), ... ,R(A+B-1))
  self:emitInstruction("CALL", register, argumentCount, returnValueCount)
  self:freeRegisters(argumentRegisterCount)
  self:allocateRegisters(returnValueCount - 2)

  return resultCount
end

function CodeGenerator:FunctionExpression(node, register)
  self:compileFunction(node, register)
end

function CodeGenerator:IndexExpression(node, register)
  local baseRegister  = self:compileExpressionNode(node.base, register)
  local indexRegister = self:compileConstantOrExpression(node.index)

  -- OP_GETTABLE [A, B, C]    R(A) := R(B)[RK(C)]
  self:emitInstruction("GETTABLE", baseRegister, baseRegister, indexRegister)
  self:freeIfRegister(indexRegister)
end

---@diagnostic disable-next-line: unused-local
function CodeGenerator:VarargExpression(node, register, resultRegisters)
  -- VARARG's B operand encodes the number of requested values as (count + 1),
  -- where B=0 (count=-1) means "take all varargs". For example:
  --   `local a, b, c = ...`  requests 3 values -> B=4
  --   `return ...`           requests all      -> B=0 (resultRegisters=-1)
  -- OP_VARARG [A, B]    R(A), R(A+1), ..., R(A+B-1) = vararg
  self:emitInstruction("VARARG", register, resultRegisters + 1, 0)

  -- the first register was already allocated by the caller via
  -- `:compileExpressionNode`, so we only need to allocate the
  -- remaining (count - 1) registers here.
  self:allocateRegisters(resultRegisters - 1)
end

--// Statement Handlers //--

-- Instruction Layout:
--   <evaluate callee>              ; Load function into register
--   <evaluate arguments>           ; Process argument list
--   CALL R(A), B, 1                ; Call with 0 return values (B=arg count, C=1)
--
function CodeGenerator:CallStatement(node)
  self:FunctionCall(node.expression, self:allocateRegisters(1), 0)

  -- We shouldn't have allocated registers from statements.
  -- (Unless it's a statement that declares local variables.)
  self:freeRegisters(1)
end

-- Instruction Layout:
--   <enter scope>                  ; Create new variable scope
--   <body statements>              ; Execute block body
--   [CLOSE R(A)]                   ; Close upvalues if needed
--   <exit scope>                   ; Restore parent scope
--
function CodeGenerator:DoStatement(node)
  self:compileChunk(node.body)
end

-- Instruction Layout:
--   JMP [loop_end]                 ; Jump to end of enclosing loop
--                                  ; (patched later when loop ends)
--
function CodeGenerator:BreakStatement()
  if not self.pendingBreakJumps then
    error("CodeGenerator: no loop to break")
  end

  table.insert(self.pendingBreakJumps, self:emitJump())
end

-- Instruction Layout:
--   local a, b, c = expr1, expr2
--   <evaluate expr1>               ; Result in R(0)
--   <evaluate expr2>               ; Result in R(1)
--   [LOADNIL R(2), R(2)]           ; Pad remaining slots with nil if needed
--   ; Variables a, b, c now map to registers R(0), R(1), R(2)
--
function CodeGenerator:LocalDeclarationStatement(node)
  local baseVariableRegister = self.currentScope.allocatedRegisters - 1

  self:compileExpressionList(node.initializers, #node.variables)
  for index, variableName in ipairs(node.variables) do
    self:declareLocalVariable(variableName, baseVariableRegister + index)
  end
end

-- Instruction Layout:
--   local function foo() ... end  ; Declare local 'foo' first (enables recursion)
--   CLOSURE R(A), [proto_index]   ; Create closure from prototype
--   [MOVE/GETUPVAL ...]           ; Capture upvalues (pseudo-instructions)
--   ; R(A) now contains the function, mapped to 'foo'
--
function CodeGenerator:LocalFunctionDeclaration(node)
  self:compileFunction(node.body, self:declareLocalVariable(node.name))
end

-- Instruction Layout:
--   a, b.c, d[e] = expr1, expr2, expr3
--   <evaluate expr1>               ; Result in R(0)
--   <evaluate expr2>               ; Result in R(1)
--   <evaluate expr3>               ; Result in R(2)
--   MOVE/SETGLOBAL R(dest), R(0)   ; Assign to 'a'
--   SETTABLE R(base), K(c), R(1)   ; Assign to 'b.c'
--   SETTABLE R(base), R(idx), R(2) ; Assign to 'd[e]'
--
function CodeGenerator:AssignmentStatement(node)
  local lvalues     = node.lvalues
  local expressions = node.expressions

  local variableBaseRegister = self.currentScope.allocatedRegisters
  local lvalueRegisterCount  = self:compileExpressionList(expressions, #lvalues)

  self:assignValuesToRegisters(lvalues, 1, variableBaseRegister)
  self:freeRegisters(lvalueRegisterCount)
end

-- Instruction Layout:
--   if condition1 then
--     <evaluate condition1>        ; Result in R(A)
--     TEST R(A), 0                 ; Test if false
--     JMP [next_clause]            ; Skip body if condition false
--     <clause1 body>               ; Execute if true
--     JMP [end]                    ; Skip remaining clauses
--
--   elseif condition2 then         ; [next_clause] offset target
--     <evaluate condition2>        ; Result in R(B)
--     TEST R(B), 0
--     JMP [else_clause]            ; Skip to else if false
--     <clause2 body>
--     JMP [end]
--
--   else                           ; [else_clause] offset target
--     <else body>                  ; No condition to test
--
--   end                            ; [end] offset - all exit jumps land here
--
function CodeGenerator:IfStatement(node)
  local clauses    = node.clauses
  local elseClause = node.elseClause

  local jumpToEndPCs = {}
  local lastClause   = clauses[#clauses]
  for _, clause in ipairs(clauses) do
    local condition = clause.condition
    local body      = clause.body

    local conditionRegister = self:compileExpressionNode(condition)

    -- OP_TEST [A, C]    if not (R(A) <=> C) then pc++
    -- If the condition is false, jump to the next instruction.
    -- (which is always a jump to the next clause)
    self:emitInstruction("TEST", conditionRegister, 0, 0)
    self:freeRegisters(1)

    -- This is the jump that occurs if the condition is false.
    -- We will patch it later to jump to the next clause.
    local conditionIsFalseJumpPC = self:emitJump()
    self:compileChunk(body)

    local isLastClause    = (clause == lastClause)
    local shouldJumpToEnd = not isLastClause or elseClause
    if shouldJumpToEnd then
      local jumpToEndPC = self:emitJump()
      table.insert(jumpToEndPCs, jumpToEndPC)
    end

    -- Patch the condition false jump to here.
    self:patchJumpToHere(conditionIsFalseJumpPC)
  end

  -- NOTE: The previous clause's condition already jumps
  -- to here, so we don't need to patch anything.
  if elseClause then
    self:compileChunk(elseClause)
  end

  -- Patch all jumps at the end of clauses to here.
  self:patchJumpsToHere(jumpToEndPCs)
end

-- Instruction Layout:
--   for i, v in pairs(t) do ... end
--   <evaluate iterator expressions>  ; R(A)=iterator, R(A+1)=state, R(A+2)=control
--   JMP [loop_test]                  ; Jump to iterator call
--   [loop_start]:                    ; Loop body entry point
--     <body statements>              ; User code, can access i, v
--   [loop_test]:                     ; Iterator call happens here
--     TFORLOOP R(A), #iterators      ; Call R(A)(R(A+1), R(A+2))
--                                    ; Sets R(A+3)...R(A+2+C) with results
--                                    ; If R(A+3) != nil: R(A+2)=R(A+3), continue
--     JMP [loop_start]               ; Jump back if more values
--   [loop_end]:                      ; Exit when iterator returns nil
--
function CodeGenerator:ForGenericStatement(node)
  local iterators   = node.iterators
  local expressions = node.expressions
  local body        = node.body

  self:enterScope()
  local baseRegister = self.currentScope.allocatedRegisters
  self:compileExpressionList(expressions, 3)
  self:declareLocalVariables(iterators)

  local startJmpInstruction = self:emitJump()
  local loopStartPC = #self.proto.code
  self:breakable(function()
    self:compileChunk(body)
    self:patchJumpToHere(startJmpInstruction)

    -- OP_TFORLOOP [A, C]    R(A+3), ... ,R(A+2+C) := R(A)(R(A+1), R(A+2))
    --                       if R(A+3) ~= nil then R(A+2)=R(A+3) else pc++
    -- Skips next instruction if there are no more values.
    self:emitInstruction("TFORLOOP", baseRegister, 0, #iterators)

    -- Emit jump back to the start of the loop if there are more values.
    self:emitJumpBack(loopStartPC)
  end)
  self:leaveScope()
end

-- Instruction Layout:
--   for i = start, limit, step do ... end
--   <evaluate start>               ; Result in R(A) - internal control var
--   <evaluate limit>               ; Result in R(A+1)
--   <evaluate step or load 1>      ; Result in R(A+2)
--   FORPREP R(A), [loop_start]     ; R(A) -= R(A+2), jump to body
--   [loop_start]:                  ; Loop body entry
--     ; User variable 'i' is in R(A+3), copied from R(A) each iteration
--     <body statements>
--   [loop_end]:
--     FORLOOP R(A), [loop_start]   ; R(A) += R(A+2), test limit, loop if valid
--   [exit]:                        ; Continue after loop
--
-- Register layout: R(A)=control, R(A+1)=limit, R(A+2)=step, R(A+3)=user var
--
function CodeGenerator:ForNumericStatement(node)
  local varName   = node.variable
  local startExpr = node.start
  local limitExpr = node.limit
  local stepExpr  = node.step
  local body      = node.body

  self:enterScope()
  local controlRegister = self:compileExpressionNode(startExpr)
  self:compileExpressionNode(limitExpr)
  local stepRegister = self:allocateRegisters(1)

  if stepExpr then
    self:compileExpressionNode(stepExpr, stepRegister)
  else
    -- OP_LOADK [A, Bx]    R(A) := Kst(Bx)
    -- Default step is 1.
    self:emitInstruction("LOADK", stepRegister, self:internConstant(1))
  end

  -- OP_FORPREP [A, sBx]    R(A)-=R(A+2) pc+=sBx
  -- Prepare the internal loop variable by subtracting the step from the start.
  -- "sBx" ("b" in our implementation) is the signed offset to jump to the loop body.
  local forPrepInstruction = self:emitInstruction("FORPREP", controlRegister, 0)
  local loopStartPC = #self.proto.code

  -- Declare the loop variable in R(A + 3), not R(A).
  -- We must use the copied "for" variable value instead of the actual one,
  -- this is because Lua 5.1 semantics state that the loop variable is constant
  -- within the loop body. If the user modifies it, the loop should not be affected.
  -- See: https://github.com/bytexenon/Tiny-Lua-Compiler/issues/11
  --
  -- Current register layout:
  --   R(A)   - internal loop variable (control variable)
  --   R(A+1) - limit
  --   R(A+2) - step
  --   R(A+3) - user loop variable "varName" (copied from R(A) every iteration)
  self:declareLocalVariable(varName, self:allocateRegisters(1))
  self:breakable(function()
    self:compileChunk(body)

    local loopEndPC = #self.proto.code
    local relativeLoopStartPC = loopEndPC - loopStartPC
    forPrepInstruction.b = relativeLoopStartPC

    -- OP_FORLOOP [A, sBx]   R(A)+=R(A+2)
    --                       if R(A) <?= R(A+1) then { pc+=sBx R(A+3)=R(A) }
    --                       (<?= is < or > depending on step sign)
    --
    -- Increment the internal loop variable R(A) by the step R(A+2), then jump
    -- back to the start of the loop if the limit R(A+1) has not been reached,
    -- copying the new value into the user loop variable R(A+3).
    self:emitInstruction("FORLOOP", controlRegister, loopStartPC - loopEndPC - 1)
  end)
  self:leaveScope()
end

-- Instruction Layout:
--   while condition do ... end
--   [loop_start]:                  ; Test condition at loop start
--     <evaluate condition>         ; Result in R(A)
--     TEST R(A), 0                 ; Test if false
--     JMP [loop_end]               ; Exit loop if condition false
--     <body statements>            ; Execute loop body
--     JMP [loop_start]             ; Jump back to re-test condition
--   [loop_end]:                    ; Exit point
--
function CodeGenerator:WhileStatement(node)
  local condition = node.condition
  local body      = node.body

  local startPC = #self.proto.code
  local conditionRegister = self:compileExpressionNode(condition)

  -- OP_TEST [A, C]    if not (R(A) <=> C) then pc++
  self:emitInstruction("TEST", conditionRegister, 0, 0)
  self:freeRegisters(1)

  local endJumpPC = self:emitJump()
  self:breakable(function()
    self:compileChunk(body)
    self:emitJumpBack(startPC)
  end)
  self:patchJumpToHere(endJumpPC)
end

-- Instruction Layout:
--   repeat ... until condition
--   [loop_start]:                  ; Body executes at least once
--     <body statements>            ; Locals declared here visible in condition
--     <evaluate condition>         ; Result in R(A)
--     [CLOSE R(B)]                 ; Close upvalues if needed
--     TEST R(A), 0                 ; Test if condition false
--     JMP [loop_start]             ; Repeat if condition false
--   [loop_end]:                    ; Exit when condition true
--
-- Note: Unlike while loops, body and condition share the same scope.
--
function CodeGenerator:RepeatStatement(node)
  local body        = node.body
  local condition   = node.condition
  local loopStartPC = #self.proto.code

  -- NOTE: Repeat statements' body and condition are in the same scope.
  self:enterScope()
  self:breakable(function()
    self:compileStatementList(body.statements)
    local conditionRegister = self:compileExpressionNode(condition)

    if self.currentScope.needClose then
      local parentScope = self.currentScope.parentScope
      local allocatedRegisters = parentScope.allocatedRegisters

      -- OP_CLOSE [A]    close all variables in the stack up to (>=) R(A)
      self:emitInstruction("CLOSE", allocatedRegisters, 0, 0)
      self.currentScope.needClose = false
    end

    -- OP_TEST [A, C]    if not (R(A) <=> C) then pc++
    self:emitInstruction("TEST", conditionRegister, 0, 0)
    self:emitJumpBack(loopStartPC)
    self:freeRegisters(1)
  end)
  self:leaveScope()
end

-- Instruction Layout (multiple return values):
--   return expr1, expr2, func()
--   <evaluate expr1>               ; Result in R(A)
--   <evaluate expr2>               ; Result in R(A+1)
--   <evaluate func()>              ; Results in R(A+2)...stack top
--   RETURN R(A), B                 ; Return R(A)...R(A+B-2)
--                                  ; B=0 means return all values to stack top
--
-- Instruction Layout (tail call optimization):
--   return func()                  ; Single function call
--   <evaluate func with args>
--   TAILCALL R(A), B               ; Call and return in one step
--                                  ; Reuses current function's stack frame
--
function CodeGenerator:ReturnStatement(node)
  local expressions      = node.expressions
  local resultRegister   = self.currentScope.allocatedRegisters
  local lastExpression   = expressions[#expressions]
  local isSingleFuncCall = self:isSingleFunctionCall(expressions)
  local numResults, exprRegisterCount

  -- Optimization:
  -- If it's a single function call, we can optimize it to a TAILCALL.
  -- Tail calls are more efficient as they reuse the current function's stack frame.
  if isSingleFuncCall then
    self:FunctionCall(
      lastExpression,
      self:allocateRegisters(1), -- This register will be manually freed later.
      -1 -- Capture all return values from the call (MULTRET).
    )

    local lastInstruction = self.proto.code[#self.proto.code]
    if not lastInstruction or lastInstruction.opname ~= "CALL" then
      error("CodeGenerator: Expected CALL instruction for tail call optimization.")
    end

    -- OP_TAILCALL [A, B, C]    return R(A)(R(A+1), ... ,R(A+B-1))
    lastInstruction.opname = "TAILCALL" -- Change CALL to TAILCALL
    numResults = -1 -- -1 = MULTRET (all results)

    self:freeRegisters(1) -- Free the temp register used for the call.
  else
    exprRegisterCount, numResults = self:compileExpressionList(expressions)
    self:freeRegisters(exprRegisterCount)
  end

  -- OP_RETURN [A, B]    return R(A), ... ,R(A+B-2)
  self:emitInstruction("RETURN", resultRegister, numResults + 1, 0)
end

--// Main Dispatch Methods //--
function CodeGenerator:compileExpressionNode(node, register, resultRegisters)
  if not node then error("Node not found.") end
  register = register or self:allocateRegisters(1)

  local func = self[node.kind]
  if not func then
    error("CodeGenerator: Expression node kind not supported: " .. node.kind)
  end

  return register, func(self, node, register, resultRegisters or 1)
end

function CodeGenerator:compileExpressionList(expressionList, expectedRegisters)
  local exprCount = #expressionList

  -- Fast path: No expressions to process.
  if exprCount == 0 and not expectedRegisters then return 0, 0 end

  local lastExpression = expressionList[exprCount]
  local isLastMultret = self:isMultiReturnNode(lastExpression)

  -- Process all expressions that produce exactly one value (all but the last
  -- expression if it can return multiple values).
  local singleValueCount = (isLastMultret and (exprCount - 1)) or exprCount
  for expressionIndex = 1, singleValueCount do
    self:compileExpressionNode(
      expressionList[expressionIndex],
      self:allocateRegisters(1),
      1
    )

    -- Is it not needed but it got processed anyway?
    -- Discard the result by freeing the allocated register. We still had to
    -- EVALUATE the expression (for side effects - e.g., `f(), g()` where we
    -- only want g()'s result must still call f()), but we don't need the value.
    if expectedRegisters and expressionIndex > expectedRegisters then
      self:freeRegisters(1)
    end
  end

  -- Handle the last expression if it can return multiple values.
  if isLastMultret then
    -- Calculate how many remaining slots need to be filled.
    -- Use -1 to signal "take all results" when no limit is specified.
    local remaining = (expectedRegisters and (expectedRegisters - singleValueCount)) or -1
    self:compileExpressionNode(lastExpression, self:allocateRegisters(1), remaining)

  -- Pad with nils if we have fewer expressions than expected slots.
  -- This implements Lua's semantics: `local a, b, c = 1` sets b and c to nil.
  elseif expectedRegisters and expectedRegisters > exprCount then
    local needed   = expectedRegisters - exprCount
    local startReg = self.currentScope.allocatedRegisters
    local endReg   = self:allocateRegisters(needed)

    -- OP_LOADNIL [A, B]    R(A) := ... := R(B) := nil
    self:emitInstruction("LOADNIL", startReg, endReg, 0)
  end

  local allocatedCount = expectedRegisters or exprCount
  local resultCount    = (isLastMultret and -1) or allocatedCount

  return allocatedCount, resultCount
end

-- This function cannot be inlined into `compileChunk` because `repeat ... until`
-- allows the condition to see variables declared in the loop body, which means
-- the body and the condition must share the same scope. Handling that requires
-- explicit scope enter/leave calls in the RepeatStatement handler anyway.
function CodeGenerator:compileStatementList(statementList)
  for _, node in ipairs(statementList) do
    local nodeHandler = self[node.kind]
    if not nodeHandler then
      error("CodeGenerator: Statement node kind not supported: " .. node.kind)
    end

    nodeHandler(self, node)
  end
end

-- Compiles a block of statements (chunk) by entering a new scope, compiling the
-- statements, and then leaving the scope. This is used for the main chunk as
-- well as for blocks in control structures (e.g., if bodies, loop bodies).
function CodeGenerator:compileChunk(block)
  self:enterScope()
  self:compileStatementList(block.statements)
  self:leaveScope()
end

--// Function Compilation //--
-- These methods handle compiling function definitions into prototypes
-- and emitting CLOSURE instructions to instantiate them at runtime.

-- Compiles the body of a function: declares parameters as locals,
-- generates the body statements, and emits a default RETURN.
function CodeGenerator:compileFunctionBody(node)
  local previousBreakJumpPCs = self.pendingBreakJumps
  self.pendingBreakJumps = nil
  self:enterScope(true) -- Enter function scope.
  for _, paramName in ipairs(node.parameters) do
    self:declareLocalVariable(paramName)
  end

  self:compileChunk(node.body)

  -- Generate default return statement.
  -- Every function must always end with a return instruction at the end in order
  -- to be valid. If we do have an explicit return at the end, we will insert
  -- the default return regardless.
  -- OP_RETURN [A, B]    return R(A), ... ,R(A+B-2)
  self:emitInstruction("RETURN", 0, 1, 0)
  self:leaveScope()
  self.pendingBreakJumps = previousBreakJumpPCs
end

-- Generates a CLOSURE instruction to instantiate a function from its
-- prototype, followed by pseudo-instructions that capture upvalues.
--
-- A closure is a runtime object that pairs a prototype (compiled code)
-- with the actual upvalue bindings from the defining scope. This is
-- how Lua implements first-class functions with lexical scoping.
--
-- The pseudo-instructions after CLOSURE tell the VM how to find each
-- upvalue:
--   MOVE     -> upvalue is a local in the immediately enclosing function
--   GETUPVAL -> upvalue is itself an upvalue of the enclosing function
function CodeGenerator:generateClosure(proto, outputRegister)
  local parentProto = self.proto
  local protoIndex  = #parentProto.protos + 1
  table.insert(parentProto.protos, proto)

  -- OP_CLOSURE [A, Bx]    R(A) := closure(KPROTO[Bx], R(A), ... ,R(A+n))
  -- Proto indices are zero-based.
  self:emitInstruction("CLOSURE", outputRegister, protoIndex - 1, 0)

  -- The following MOVE / GETUPVAL instructions are not arbitrary code.
  -- Together they act as metadata for OP_CLOSURE, telling the VM where each
  -- captured variable should be sourced from.
  for _, upvalueName in ipairs(proto.upvalues) do
    local upvalueType = self:getVariableType(upvalueName, true)

    -- NOTE: The `B` operand is always 0 in these pseudo instructions,
    --       as the upvalue index is determined by the order in which
    --       these instructions appear after the CLOSURE instruction.
    if upvalueType == "Local" then
      -- OP_MOVE [A, B]    R(A) := R(B)
      self:emitInstruction(
        "MOVE",
        0,
        self:getLocalRegister(upvalueName)
      )
    elseif upvalueType == "Upvalue" then
      -- OP_GETUPVAL [A, B]    R(A) := UpValue[B]
      self:emitInstruction(
        "GETUPVAL",
        0,
        self:internUpvalue(upvalueName)
      )
    else
      error(
        "CodeGenerator: The upvalue cannot be a global: "
        .. upvalueName
      )
    end
  end
end

-- Compiles a function AST node into a child prototype.
function CodeGenerator:compileFunction(functionNode, outputRegister)
  local parentProto = self.proto
  local childProto  = self:createPrototype({
    isVararg  = functionNode.isVararg,
    numParams = #functionNode.parameters,
  })
  self.proto = childProto
  self:compileFunctionBody(functionNode)
  self.proto = parentProto

  -- Should we generate specialized instructions to create this closure?
  -- (This is always true except for the top-level main chunk.)
  if outputRegister then
    self:generateClosure(childProto, outputRegister)
  end

  -- Remove internal lookups, as they're no longer needed after codegen.
  childProto.constantLookup = nil
  childProto.upvalueLookup  = nil

  -- Strip the upvalue name list down to a plain count. The names were only
  -- needed during code generation to emit the MOVE/GETUPVAL pseudo-instructions.
  -- After that, the VM and bytecode emitter only need the count. The Lua 5.1
  -- binary format does not store upvalue names, so we match that here.
  childProto.numUpvals = #childProto.upvalues
  childProto.upvalues  = nil

  return childProto
end

--// Entry Point //--

-- Compiles the entire program AST into a top-level function prototype.
--
-- In Lua, the top-level code (the "main chunk") is itself a function.
-- It takes no named parameters but always accepts varargs, which is
-- how command-line arguments are passed (accessible via `...`).
-- This is why `lua -e "print(...)"` can print arguments.
function CodeGenerator:generate()
  return self:compileFunction({
    body       = self.ast.body,
    parameters = {},
    isVararg   = true -- The main chunk always accepts varargs.
                      -- (used for command-line arguments)
  })
end

-- ============================================================================
--                           THE BYTECODE EMITTER!
-- ============================================================================
--
-- The bytecode emitter turns TLC's in-memory function prototypes into a real
-- Lua 5.1 binary chunk.
--
-- In other words, this is the stage that crosses the boundary from structured
-- Lua tables to raw bytes that the standard Lua VM can load.
--
-- At a high level, this section does three jobs:
--   1. Encode scalar values (u8, u32, u64, float64, strings)
--   2. Encode instructions into their packed 32-bit forms
--   3. Serialize a function prototype tree in Lua's chunk layout
--
-- Visual overview of the final chunk layout:
--
--   +--------------------+
--   |   Common Header    |  signature, version, sizes, endianness
--   +--------------------+
--   |   Main Prototype   |
--   |  +--------------+  |
--   |  | Header       |  |  function metadata (param count, vararg flag, etc.)
--   |  +--------------+  |
--   |  | Code         |  |  instruction list
--   |  +--------------+  |
--   |  | Constants    |  |
--   |  +--------------+  |
--   |  | Child Protos |  |  nested functions
--   |  +--------------+  |
--   |  | Debug Info   |  |  omitted by TLC (encoded as empty sections)
--   |  +--------------+  |
--   +--------------------+

--// Utility Functions //--

-- A small compatibility implementation of `frexp`.
--
-- It decomposes a number so that:
--   value = mantissa * 2^exponent
-- with `mantissa` normalized to the range [0.5, 1) for non-zero values.
-- This is needed because newer Lua versions removed `math.frexp`.
local function frexp(value)
  -- Use built-in if available (Lua 5.1, Lua 5.2)
  if math.frexp then return math.frexp(value) end
  if value == 0 then return 0, 0 end

  local absVal = math.abs(value)
  local isNegative = value < 0
  local exponent = 0

  -- Scale down.
  while absVal >= 1 do
    absVal = absVal / 2
    exponent = exponent + 1
  end

  -- Scale up.
  while absVal < 0.5 do
    absVal = absVal * 2
    exponent = exponent - 1
  end

  -- Preserve sign.
  local mantissa = (isNegative and -absVal) or absVal
  return mantissa, exponent
end

-- A bitwise left shift, simulated with multiplication.
local function lshift(value, shift)
  return value * (2 ^ shift)
end

-- Operand types for instruction encoding.
-- https://www.lua.org/source/5.1/lopcodes.h.html#OpArgMask
local OP_ARG_N = 0 -- Operand is not used
local OP_ARG_U = 1 -- Operand is used
local OP_ARG_R = 2 -- Operand is a register or a jump offset
local OP_ARG_K = 3 -- Operand is a constant or register/constant

-- Function Header Flags
-- https://www.lua.org/source/5.1/lobject.h.html#VARARG_HASARG
local VARARG_NONE     = 0
local VARARG_ISVARARG = 2 -- Flag indicating a function accepts '...'

-- Constant type flags.
-- https://www.lua.org/source/5.1/lua.h.html#LUA_TNONE
local LUA_TNUMBER  = 3
local LUA_TSTRING  = 4

-- Operand sizes/positions.
-- https://www.lua.org/source/5.1/lopcodes.h.html#SIZE_C
local POS_A   = 6
local POS_B   = 23
local POS_C   = 14
local POS_Bx  = 14
local POS_sBx = 14

-- Lua 5.1 Binary Chunk Header
-- The header is a fixed sequence of bytes that identifies the file as Lua
-- bytecode and describes the data layout expected by the loader.
--
-- NOTE: We emit sizeof(size_t)=8. This means the generated bytecode
-- requires a 64-bit Lua 5.1 runtime to load. On 32-bit systems, Lua
-- will reject the chunk with "bad header". To support 32-bit, change
-- the size_t byte to 4 and use `:encodeUint32` for string lengths.
local LUA_COMMON_HEADER = "\27Lua" -- Standard magic number for Lua bytecode.
  .. string.char(0x51) -- version             0x51 signifies Lua version 5.1
  .. string.char(0)    -- format              0 = official format version
  .. string.char(1)    -- endianness          1 = little-endian, 0 = big-endian
  .. string.char(4)    -- sizeof(int)         4 bytes
  .. string.char(8)    -- sizeof(size_t)      8 bytes
  .. string.char(4)    -- sizeof(Instruction) 4 bytes
  .. string.char(8)    -- sizeof(lua_Number)  8 bytes
  .. string.char(0)    -- integral flag       0 = lua_Number is floating-point

-- Lua 5.1 Instruction Argument Modes. Instructions have different ways of
-- encoding their arguments (operands) within their 32 bits.
-- https://www.lua.org/source/5.1/lopcodes.h.html#OpMode

--  - iABC:  Uses three arguments: A (8 bits), B (9 bits), C (9 bits).
--           Common for operations involving 3 registers or 2 regs + 1 constant/literal.
--
--           Bit layout (least significant bit on the right):
--
--             31                    23 22              14 13      6 5       0
--             +----------------------+------------------+----------+--------+
--             |        B (9)         |      C (9)       |  A (8)   | Op (6) |
--             +----------------------+------------------+----------+--------+
local MODE_iABC = 0

--  - iABx:  Uses two arguments: A (8 bits), Bx (18 bits).
--           'Bx' combines B and C fields to hold a larger unsigned value,
--           often an index into the constant table or a prototype index.
--
--             31                                        14 13      6 5       0
--             +------------------------------------------+----------+--------+
--             |                 Bx (18)                  |  A (8)   | Op (6) |
--             +------------------------------------------+----------+--------+
local MODE_iABx = 1

--  - iAsBx: Uses two arguments: A (8 bits), sBx (18 bits, signed).
--           'sBx' is like Bx but represents a signed offset, primarily used
--           for jump instructions. The offset is biased, to get the actual value,
--           you need to subtract 131072 (2^17) from the value.
--           This allows for jumps in both directions (forward and backward).
--
--             31                                        14 13      6 5       0
--             +------------------------------------------+----------+--------+
--             |                sBx (18)                  |  A (8)   | Op (6) |
--             +------------------------------------------+----------+--------+
local MODE_iAsBx = 2

-- This table maps all Lua 5.1 opcodes to their respective argument modes,
-- it's used to determine how to encode the instruction arguments.
--
--  Format: [Opname] = {opcodeIndex, argumentMode, argBMode, argCMode}
--   - opcodeIndex: The index (0-based) of the opcode in the Lua 5.1 instruction set.
--   - argumentMode: The mode used to encode the instruction arguments.
--   - argBMode: The type of the B argument (OP_ARG_N, OP_ARG_U, OP_ARG_R, OP_ARG_K).
--   - argCMode: The type of the C argument (OP_ARG_N, OP_ARG_U, OP_ARG_R, OP_ARG_K).
--
-- https://www.lua.org/source/5.1/lopcodes.c.html#luaP_opmodes
local OPCODE_INFO = {
  -- Opname        Idx  Arg Mode    B Mode    C Mode
  ["MOVE"]      = {0,  MODE_iABC,  OP_ARG_R, OP_ARG_N},
  ["LOADK"]     = {1,  MODE_iABx,  OP_ARG_K, OP_ARG_N},
  ["LOADBOOL"]  = {2,  MODE_iABC,  OP_ARG_U, OP_ARG_U},
  ["LOADNIL"]   = {3,  MODE_iABC,  OP_ARG_R, OP_ARG_N},
  ["GETUPVAL"]  = {4,  MODE_iABC,  OP_ARG_U, OP_ARG_N},
  ["GETGLOBAL"] = {5,  MODE_iABx,  OP_ARG_K, OP_ARG_N},
  ["GETTABLE"]  = {6,  MODE_iABC,  OP_ARG_R, OP_ARG_K},
  ["SETGLOBAL"] = {7,  MODE_iABx,  OP_ARG_K, OP_ARG_N},
  ["SETUPVAL"]  = {8,  MODE_iABC,  OP_ARG_U, OP_ARG_N},
  ["SETTABLE"]  = {9,  MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["NEWTABLE"]  = {10, MODE_iABC,  OP_ARG_U, OP_ARG_U},
  ["SELF"]      = {11, MODE_iABC,  OP_ARG_R, OP_ARG_K},
  ["ADD"]       = {12, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["SUB"]       = {13, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["MUL"]       = {14, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["DIV"]       = {15, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["MOD"]       = {16, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["POW"]       = {17, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["UNM"]       = {18, MODE_iABC,  OP_ARG_R, OP_ARG_N},
  ["NOT"]       = {19, MODE_iABC,  OP_ARG_R, OP_ARG_N},
  ["LEN"]       = {20, MODE_iABC,  OP_ARG_R, OP_ARG_N},
  ["CONCAT"]    = {21, MODE_iABC,  OP_ARG_R, OP_ARG_R},
  ["JMP"]       = {22, MODE_iAsBx, OP_ARG_R, OP_ARG_N},
  ["EQ"]        = {23, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["LT"]        = {24, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["LE"]        = {25, MODE_iABC,  OP_ARG_K, OP_ARG_K},
  ["TEST"]      = {26, MODE_iABC,  OP_ARG_R, OP_ARG_U},
  ["TESTSET"]   = {27, MODE_iABC,  OP_ARG_R, OP_ARG_U},
  ["CALL"]      = {28, MODE_iABC,  OP_ARG_U, OP_ARG_U},
  ["TAILCALL"]  = {29, MODE_iABC,  OP_ARG_U, OP_ARG_U},
  ["RETURN"]    = {30, MODE_iABC,  OP_ARG_U, OP_ARG_N},
  ["FORLOOP"]   = {31, MODE_iAsBx, OP_ARG_R, OP_ARG_N},
  ["FORPREP"]   = {32, MODE_iAsBx, OP_ARG_R, OP_ARG_N},
  ["TFORLOOP"]  = {33, MODE_iABC,  OP_ARG_N, OP_ARG_U},
  ["SETLIST"]   = {34, MODE_iABC,  OP_ARG_U, OP_ARG_U},
  ["CLOSE"]     = {35, MODE_iABC,  OP_ARG_N, OP_ARG_N},
  ["CLOSURE"]   = {36, MODE_iABx,  OP_ARG_U, OP_ARG_N},
  ["VARARG"]    = {37, MODE_iABC,  OP_ARG_U, OP_ARG_N}
}

local BytecodeEmitter = {}

-- RK (Register-or-Constant) encoding:
--
-- Lua packs a register index and a constant index into the same 9-bit operand.
-- Values 0..255 refer to registers (stack[operand]).
-- Values >= 256 refer to constants and map to proto.constants[operand - 256].
--
-- TLC represents constants as negative indices (e.g. -1 == first constant).
-- encodeRKOperand converts TLC's negative constant index into the unsigned RK
-- form used in bytecode (for example: -1 -> 256).
--
-- Visual mapping:
--
--   TLC internal operand      Lua RK field      Meaning
--   --------------------      ------------      ---------------------------
--   0                         0                 register 0
--   7                         7                 register 7
--   -1                        256               constant 0
--   -2                        257               constant 1
--   -n                        255 + n           constant (n - 1)
function BytecodeEmitter.encodeRKOperand(value)
  if value < 0 then
    return 255 - value
  end
  return value
end

-- If it's not an RK operand, we need to encode constants differently for
-- instructions that use iABx format (like LOADK, GETGLOBAL, SETGLOBAL).
--
-- Example:
--   LOADK     5, -4 --> LOADK     5, 3
--   GETGLOBAL 2, -1 --> GETGLOBAL 2, 0
--
-- Unlike RK encoding, ABx operands cannot represent registers. They store
-- only a plain zero-based constant/prototype index.
function BytecodeEmitter.encodeConstantIndex(value)
  if value < 0 then
    return -value - 1
  end
  return value
end

function BytecodeEmitter.encodeUint8(value)
  if value < 0 or value > 2^8 - 1 then
    error(
      "BytecodeEmitter: encodeUint8 value out of range of u8: "
      .. tostring(value)
    )
  end

  return string.char(value % 256)
end

function BytecodeEmitter.encodeUint32(value)
  if value < 0 or value > 2^32 - 1 then
    error(
      "BytecodeEmitter: encodeUint32 value out of range of u32: "
      .. value
    )
  end

  local b1 = value % 256
  value = math.floor(value / 256)
  local b2 = value % 256
  value = math.floor(value / 256)
  local b3 = value % 256
  value = math.floor(value / 256)
  local b4 = value % 256

  return string.char(b1, b2, b3, b4)
end

function BytecodeEmitter.encodeUint64(value)
  if value < 0 or value > 2^64 - 1 then
    error(
      "BytecodeEmitter: encodeUint64 value out of range of u64: "
      .. tostring(value)
    )
  end

  local lowWord  = value % 2^32
  local highWord = math.floor(value / 2^32)
  return BytecodeEmitter.encodeUint32(lowWord)
        .. BytecodeEmitter.encodeUint32(highWord)
end

-- Encodes a Lua number (float64) into its IEEE 754 binary representation (8 bytes).
-- Reference: https://en.wikipedia.org/wiki/Double-precision_floating-point_format.
--
-- Note: we emit `lowWord` first since we're emitting in little-endian format.
-- This means the least significant byte comes first in the byte stream.
-- (`BytecodeEmitter.encodeUint32` handles per-byte little-endian encoding for us).
--
-- IEEE 754 double layout:
--
--   63            52 51                                              0
--   +---------------+------------------------------------------------+
--   | sign (1 bit)  | exponent (11 bits) | fraction / mantissa (52)  |
--   +---------------+------------------------------------------------+
--
-- TLC builds this as two 32-bit words:
--
--   highWord = [ sign | exponent | top 20 mantissa bits ]
--   lowWord  = [ remaining 32 mantissa bits ]
function BytecodeEmitter.encodeFloat64(value)
  -- Handle the simple edge cases first.
  if value == 0 then
    if 1/value < 0 then -- Is it a negative zero?
      -- IEEE 754 representation for negative zero:
      --  sign     = 1
      --  exponent = all 0 bits (0)
      --  fraction = all 0 bits (0)
      local lowWord = 0
      local highWord = 0x80000000 -- sign=1, exponent=0, fraction=0
      return BytecodeEmitter.encodeUint32(lowWord)
            .. BytecodeEmitter.encodeUint32(highWord)
    end

    -- This is a normal zero (positive zero), which is just 8 zero-bytes.
    return BytecodeEmitter.encodeUint64(0)
  elseif value == math.huge or value == -math.huge then -- Check for infinity
    -- IEEE 754 representation for positive infinity:
    --  sign     = 0
    --  exponent = all 11 bits set to 1 (2047)
    --  fraction = all 52 bits set to 0
    local lowWord = 0
    local highWord = 0x7FF00000 -- sign=0, exponent=2047, fraction=0
    if value < 0 then
      highWord = highWord + 0x80000000 -- Set sign bit for negative infinity
    end

    return BytecodeEmitter.encodeUint32(lowWord)
          .. BytecodeEmitter.encodeUint32(highWord)
  elseif value ~= value then -- Check for NaN (Not a Number)
    -- IEEE 754 representation for NaN:
    --  exponent = all 11 bits set to 1 (2047)
    --  fraction = any non-zero value
    --
    -- IEEE 754 defines two types of NaN:
    --   - Quiet NaN (qNaN): Most significant bit of fraction is 1.
    --     Used for propagating NaN through operations without raising exceptions.
    --   - Signaling NaN (sNaN): Most significant bit of fraction is 0.
    --     Triggers exceptions when used in operations.
    --
    -- We use a simple quiet NaN with fraction = 1.
    local lowWord = 1
    local highWord = 0x7FF80000 -- exponent=2047, fraction most significant bit=1
    return BytecodeEmitter.encodeUint32(lowWord)
          .. BytecodeEmitter.encodeUint32(highWord)
  end

  -- Deconstruct the number into its core parts.
  local sign = (value < 0 and 1) or 0
  local mantissa, exponent = frexp(math.abs(value))

  -- IEEE 754 uses biased exponent representation. For double-precision, the
  -- bias is 1023. However, we use 1022 here because of how frexp works.
  --
  -- Lua's frexp returns a mantissa in the range [0.5, 1), but IEEE 754 expects
  -- the mantissa to be in the range [1, 2). To fix this, we effectively shift
  -- the mantissa left by 1 (multiplying by 2) and decrease the exponent by 1.
  --
  -- So: actual_exponent = frexp_exponent - 1
  --     biased_exponent = actual_exponent + 1023
  --                     = (frexp_exponent - 1) + 1023
  --                     = frexp_exponent + 1022
  local biasedExponent = exponent + 1022
  if biasedExponent <= 0 then
    -- In some cases (very small numbers), the biased exponent can get set to
    -- zero or negative, which is not representable in the normal normalized
    -- form.
    error(
      "BytecodeEmitter: Number too small to encode as float64: "
      .. tostring(value)
    )
  end

  -- The mantissa from frexp is between 0.5 and 1.0. We need to scale it
  -- to fit into the 52 bits of the fraction field. We also discard the
  -- implicit leading '1' bit of the fraction, hence the `- 1`.
  --
  -- The calculation `(mantissa * 2 - 1)` extracts the fractional part
  -- (removing the implicit leading 1). We then multiply by 2^52 to fill
  -- the 52-bit fraction field.
  local integerMantissa = (mantissa * 2 - 1) * (2^52)

  -- The lower 32 bits are simply the lower 32 bits of our 52-bit mantissa.
  -- We can get this using the modulo operator.
  --
  -- Note: We split the value into two 32-bit words because large numbers lose
  -- precision when represented as a single 64-bit integer in Lua, by storing
  -- the mantissa across two 32-bit integers we can preserve all bits accurately.
  local lowWord = integerMantissa % 2^32

  -- The higher 32 bits are a combination of the sign, exponent, and the
  -- remaining 20 bits of the mantissa.
  --
  -- We integer divide by 2^32 to shift the upper bits down.
  local mantissaHigh = math.floor(integerMantissa / 2^32)

  -- Now we use lshift to put everything in its place for the high word:
  --  - The sign bit is the 31st bit.
  --  - The exponent's 11 bits start at the 20th bit.
  --  - The remaining 20 bits of the mantissa fill the lowest 20 bits.
  local highWord = lshift(sign, 31)
                  + lshift(biasedExponent, 20)
                  + mantissaHigh

  -- Convert the two 32-bit chunks to bytes.
  return BytecodeEmitter.encodeUint32(lowWord)
        .. BytecodeEmitter.encodeUint32(highWord)
end

--// Bytecode Encoding Methods //--
-- These methods serialize various data types into their binary bytecode
-- representations following the Lua 5.1 binary chunk format specification.

function BytecodeEmitter.encodeString(value)
  -- Lua bytecode strings are null-terminated despite having an explicit
  -- length prefix. This redundancy is a legacy quirk from Lua's C
  -- implementation where strings are treated as C-style null-terminated
  -- strings. We must include '\0' for bytecode compatibility.
  value = value .. "\0"
  local size = BytecodeEmitter.encodeUint64(#value)
  return size .. value
end

function BytecodeEmitter.encodeConstant(constantValue)
  local constType = type(constantValue)

  if constType == "number" then
    return BytecodeEmitter.encodeUint8(LUA_TNUMBER)
          .. BytecodeEmitter.encodeFloat64(constantValue)
  elseif constType == "string" then
    return BytecodeEmitter.encodeUint8(LUA_TSTRING)
          .. BytecodeEmitter.encodeString(constantValue)
  end

  error("BytecodeEmitter: Unsupported constant type '" .. constType .. "'")
end

function BytecodeEmitter.assertOperandRange(value, min, max, operandName, opname)
  if value < min or value > max then
    error(string.format(
      "BytecodeEmitter: Operand %s overflow in instruction '%s'."
      .. " Got %d. Valid range is %d to %d.",
      operandName, opname, value, min, max
    ))
  end
end

-- Format: [ Opcode(6) | A(8) | C(9) | B(9) ]
function BytecodeEmitter.encodeInstructionABC(opcode, instruction, argBMode, argCMode)
  local a, b, c = instruction.a, instruction.b, instruction.c

  BytecodeEmitter.assertOperandRange(b, -2^8, 2^8 - 1, "B", instruction.opname)
  BytecodeEmitter.assertOperandRange(c, -2^8, 2^8 - 1, "C", instruction.opname)

  -- If the argument mode is OpArgK, we encode constants specially.
  if argBMode == OP_ARG_K then b = BytecodeEmitter.encodeRKOperand(b) end
  if argCMode == OP_ARG_K then c = BytecodeEmitter.encodeRKOperand(c) end

  -- NOTE: In the binary layout, C (bits 14-22) comes BEFORE B (bits 23-31).
  local shiftedA = lshift(a, POS_A)
  local shiftedB = lshift(b, POS_B)
  local shiftedC = lshift(c, POS_C)
  return BytecodeEmitter.encodeUint32(opcode + shiftedA + shiftedB + shiftedC)
end

-- Format: [ Opcode(6) | A(8) | Bx(18) ]
-- Used for loading constants (LOADK) and global lookups (GETGLOBAL/SETGLOBAL).
function BytecodeEmitter.encodeInstructionABx(opcode, instruction)
  local a, bx = instruction.a, BytecodeEmitter.encodeConstantIndex(instruction.b)

  BytecodeEmitter.assertOperandRange(bx, 0, 2^18 - 1, "Bx", instruction.opname)

  local shiftedA  = lshift(a, POS_A)
  local shiftedBx = lshift(bx, POS_Bx)
  return BytecodeEmitter.encodeUint32(opcode + shiftedA + shiftedBx)
end

-- Format: [ Opcode(6) | A(8) | sBx(18) ]
-- Used for jumps (JMP, FORLOOP).
-- sBx is a signed 18-bit integer represented in "Excess-K" encoding.
function BytecodeEmitter.encodeInstructionAsBx(opcode, instruction)
  local a, b = instruction.a, instruction.b

  BytecodeEmitter.assertOperandRange(b, -2^17 + 1, 2^17, "sBx", instruction.opname)

  -- sBx Encoding (Bias/Excess-K):
  -- Lua uses "bias encoding" for signed jump offsets. We add 131071 (2^17-1)
  -- to shift the range [-131071, 131072] to unsigned [0, 262143]. This lets
  -- the instruction format treat all 18 bits as unsigned while still
  -- representing signed values. The VM subtracts the bias during execution.
  -- This is simpler than two's complement for the instruction decoder.
  local bias = 2^17 - 1
  local sbx  = b + bias

  local shiftedA  = lshift(a, POS_A)
  local shiftedBx = lshift(sbx, POS_sBx)
  return BytecodeEmitter.encodeUint32(opcode + shiftedA + shiftedBx)
end

-- Field     | Size (bits) | Shift Left By | Position (from bit 0)
-- ----------|-------------|---------------|-----------------------
-- Opcode    | 6           | 0             | 0-5
-- A         | 8           | 6             | 6-13
-- C         | 9           | 14            | 14-22 (for iABC)
-- Bx/sBx    | 18          | 14            | 14-31 (for iABx/iAsBx)
-- B         | 9           | 23            | 23-31 (for iABC)
function BytecodeEmitter.encodeInstruction(instruction)
  local opname     = instruction.opname or "<unknown>"
  local opcodeInfo = OPCODE_INFO[opname]

  if not opcodeInfo then
    error("BytecodeEmitter: Unsupported instruction '" .. opname .. "'")
  end

  local opcode, instructionMode = opcodeInfo[1], opcodeInfo[2]
  local argBMode, argCMode      = opcodeInfo[3], opcodeInfo[4]
  local a = instruction.a

  BytecodeEmitter.assertOperandRange(a, 0, 2^8 - 1, "A", opname)

  if instructionMode == MODE_iABC then
    return BytecodeEmitter.encodeInstructionABC(opcode, instruction, argBMode, argCMode)
  elseif instructionMode == MODE_iABx then
    return BytecodeEmitter.encodeInstructionABx(opcode, instruction)
  elseif instructionMode == MODE_iAsBx then
    return BytecodeEmitter.encodeInstructionAsBx(opcode, instruction)
  end

  error(string.format(
    "BytecodeEmitter: Unsupported instruction format for '%s'.",
    opname
  ))
end

function BytecodeEmitter.encodeArray(values, encodeValue)
  -- Every section starts with a 4-byte unsigned integer (size_t) indicating the
  -- number of items in that section, followed by the encoded items themselves.
  local array = { BytecodeEmitter.encodeUint32(#values) }
  for index, value in ipairs(values) do
    array[index + 1] = encodeValue(value)
  end
  return table.concat(array)
end

function BytecodeEmitter.encodeFunctionPrototype(proto)
  if not proto.code or not proto.constants then
    error("BytecodeEmitter: Invalid function prototype")
  end

  -- Note: TLC doesn't implement debug info as that would require tracking
  --       line numbers, local variables, and upvalue names, etc. which will
  --       add a lot of unnecessary complexity to the compiler.
  --
  -- Optimization: Instead of using `..` to concatenate strings repeatedly, we
  --               build a table of string parts and concatenate once at the end.
  --               This is more efficient, as it avoids creating intermediate
  --               strings on each concatenation.
  return table.concat({
    -- Prototype Header --
    BytecodeEmitter.encodeString("@tlc"),         -- Source name.
    BytecodeEmitter.encodeUint32(0),              -- Line defined (debug).
    BytecodeEmitter.encodeUint32(0),              -- Last line defined (debug).
    BytecodeEmitter.encodeUint8(proto.numUpvals), -- Number of upvalues.
    BytecodeEmitter.encodeUint8(proto.numParams), -- Number of parameters.
    BytecodeEmitter.encodeUint8((proto.isVararg and VARARG_ISVARARG) or VARARG_NONE),
    BytecodeEmitter.encodeUint8(proto.maxStackSize),

    -- Sections --
    BytecodeEmitter.encodeArray(proto.code, BytecodeEmitter.encodeInstruction),
    BytecodeEmitter.encodeArray(proto.constants, BytecodeEmitter.encodeConstant),
    BytecodeEmitter.encodeArray(proto.protos, BytecodeEmitter.encodeFunctionPrototype),

    -- Debug info (not implemented) --
    -- These sections would contain source line mappings, local variable names,
    -- and upvalue names for debugging and error reporting. We emit zeros to
    -- indicate empty debug info, which is valid but means stack traces won't
    -- show file:line information. A full implementation would track this data
    -- throughout compilation.

    BytecodeEmitter.encodeUint32(0), -- Line info count.
    BytecodeEmitter.encodeUint32(0), -- Local variable count.
    BytecodeEmitter.encodeUint32(0)  -- Upvalue name count.
  })
end

--// Main //--
function BytecodeEmitter.emit(mainProto)
  assert(
    type(mainProto) == "table",
    "Expected table for 'mainProto', got " .. type(mainProto)
  )
  assert(
    mainProto.code and mainProto.constants,
    "Expected a valid Lua function prototype for 'mainProto'"
  )

  return LUA_COMMON_HEADER .. BytecodeEmitter.encodeFunctionPrototype(mainProto)
end

-- ============================================================================
--                                   (/^▽^)/
--                             THE VIRTUAL MACHINE!
-- ============================================================================
--
-- The final stage: execution.
--
-- At this point the compiler has produced a function prototype - a table
-- containing instructions, constants, and nested child prototypes. This VM
-- interprets those prototypes using Lua 5.1's register model.
--
-- TLC runs the in-memory prototype tables directly rather than decoding binary
-- `.luac` chunks. This keeps the implementation focused on what matters:
-- instruction semantics.
--
-- This VM is intentionally not complete. It skips metamethod dispatch, proper
-- garbage collection, and a handful of other functionality, relying instead on
-- the host Lua environment for those. This is what lets the whole VM fit in a
-- few hundred readable lines.
--
-- For a more complete Lua-in-Lua VM, see FiOne:
--   https://github.com/Rerumu/FiOne

local unpack = (unpack or table.unpack) -- Lua 5.1/5.2+ compatibility.

-- Special value for CALL/TAILCALL/RETURN/etc. to indicate "all arguments/results".
local USE_CURRENT_TOP = 0

-- Maximum amount of implicit (numeric) table fields to set when executing a
-- single SETLIST instruction. This constant is set to the Lua 5.1 limit of 50.
local FIELDS_PER_FLUSH = SETLIST_MAX

-- Counts vararg arguments accurately, including any trailing nils.
-- `#{...}` would be wrong because Lua's # operator is undefined for
-- tables with holes (nils in the middle or at the end). `select("#", ...)`
-- counts ALL arguments including trailing nils, which is what the VM
-- needs to correctly populate vararg registers.
local function pack(...)
  return select("#", ...), { ... }
end

local VirtualMachine = {}

--// Execution //--

-- Runtime frame layout:
--
--   stack[0..maxStackSize-1]   -> registers for the current closure
--   constants[]                -> immutable constant pool from the prototype
--   upvalues[]                 -> captured variables from enclosing scopes
--   stackTop                   -> a value that the VM stores which tells it up
--                                 until which register (stack slot) to return or
--                                 take values from when the USE_CURRENT_TOP is used
--
-- Reference: https://www.lua.org/source/5.1/lvm.c.html#luaV_execute
function VirtualMachine.executeClosure(closure, ...)
  --
  -- In our implementation, registers used in a function are stored in a unique
  -- table (`stack`) that is created when the function is called. The original
  -- Lua VM's implementation has a single large stack for all function calls
  -- that grows and shrinks as functions are called and return. However,
  -- creating a new table per function call is simpler to implement in Lua and
  -- doesn't have significant performance implications for our educational
  -- purposes, so we take that approach for clarity.
  --
  -- In the VM we use closures to represent function values, and each closure has
  -- a reference to its prototype (which contains the instructions and constants),
  -- the reason why we don't use prototypes instead is because closures allow us
  -- to hold runtime information like upvalues and the environment, which is
  -- necessary for executing the function correctly.
  --
  -- Think of prototypes as the compile-time blueprint for a function's structure,
  -- and closures as the runtime instances that carry the actual state needed for
  -- execution.

  local stack = {}

  -- stackTop is a useful abstraction that can hold a runtime information of how
  -- many registers are currently in use. This is especially important for instructions
  -- like CALL and RETURN where the number of arguments or return values may not be
  -- know until runtime (when USE_CURRENT_TOP is used). By tracking stackTop, the VM can
  -- correctly handle these cases without needing to hardcode limits or make assumptions.
  local stackTop = 0

  local openUpvalues = {}
  local highestOpenUpvalue = -1

  local closureEnv    = closure.env
  local closureProto  = closure.proto
  local closureUpvals = closure.upvalues

  local protoCode      = closureProto.code
  local protoConstants = closureProto.constants
  local protoNumParams = closureProto.numParams
  local protoIsVararg  = closureProto.isVararg

  -- Gets a value from either the stack or constants table. Register or constant.
  --
  -- NOTE: Constant indices are represented as negative numbers
  -- in our instruction operands, so in order to get the correct constant,
  -- we negate the index when accessing the constants table.
  local function rk(index)
    if index < 0 then
      return protoConstants[-index]
    end
    return stack[index]
  end

  local vararg, varargCount
  local fixedArguments = { ... }

  -- Named parameters are copied into R(0) .. R(numparams - 1).
  for paramIdx = 1, protoNumParams do
    stack[paramIdx - 1] = fixedArguments[paramIdx]
  end
  if protoIsVararg then
    varargCount, vararg = pack(select(protoNumParams + 1, ...))
  end

  -- The current instruction index (also referred to as PC).
  local programCounter = 1
  local codeLength     = #protoCode
  while programCounter <= codeLength do
    local instruction = protoCode[programCounter]
    local opname  = instruction.opname
    local a, b, c = instruction.a, instruction.b, instruction.c

    -- Many Lua VM implementations written in Lua use binary search on sorted
    -- opcode tables for faster dispatch, reducing lookup time from O(n) to
    -- O(log n). For clarity and maintainability, TLC uses a simple linear
    -- if-else chain instead. This makes the code easier to read and modify.

    -- OP_MOVE [A, B]    R(A) := R(B)
    -- Copy a value between registers.
    if opname == "MOVE" then
      stack[a] = stack[b]

    -- OP_LOADK [A, Bx]    R(A) = Kst(Bx)
    -- Load a constant into a register.
    elseif opname == "LOADK" then
      stack[a] = protoConstants[-b]

    -- OP_LOADBOOL [A, B, C]    R(A) := (Bool)B; if (C) pc++
    -- Load a boolean into a register. Skip next instruction if C is non-zero.
    elseif opname == "LOADBOOL" then
      stack[a] = (b ~= 0)
      if c ~= 0 then
        programCounter = programCounter + 1
      end

    -- OP_LOADNIL [A, B]    R(A) := ... := R(B) := nil
    -- Load nil values into a range of registers.
    elseif opname == "LOADNIL" then
      for register = a, b do
        stack[register] = nil
      end

    -- OP_GETUPVAL [A, B]    R(A) := UpValue[B]
    -- Read an upvalue into a register.
    elseif opname == "GETUPVAL" then
      local upvalue = closureUpvals[b + 1]
      stack[a] = upvalue.stack[upvalue.index]

    -- OP_GETGLOBAL [A, Bx]    R(A) := Gbl[Kst(Bx)]
    -- Read a global variable into a register.
    elseif opname == "GETGLOBAL" then
      stack[a] = closureEnv[protoConstants[-b]]

    -- OP_GETTABLE [A, B, C]    R(A) := R(B)[RK(C)]
    -- Read a table element into a register.
    elseif opname == "GETTABLE" then
      stack[a] = stack[b][rk(c)]

    -- OP_SETGLOBAL [A, Bx]    Gbl[Kst(Bx)] := R(A)
    -- Write a register value into a global variable.
    elseif opname == "SETGLOBAL" then
      closureEnv[rk(b)] = stack[a]

    -- OP_SETUPVAL [A, B]    UpValue[B] := R(A)
    -- Write a register value into an upvalue.
    elseif opname == "SETUPVAL" then
      local upvalue = closureUpvals[b + 1]
      upvalue.stack[upvalue.index] = stack[a]

    -- OP_SETTABLE [A, B, C]    R(A)[RK(B)] := RK(C)
    -- Write a register value into a table element.
    elseif opname == "SETTABLE" then
      stack[a][rk(b)] = rk(c)

    -- OP_NEWTABLE [A, B, C]    R(A) := {} (size = B,C)
    -- Create a new table and store it in a register.
    -- B = array part size hint, C = hash part size hint (ignored here).
    elseif opname == "NEWTABLE" then
      stack[a] = {}

    -- OP_SELF [A, B, C]    R(A+1) := R(B); R(A) := R(B)[RK(C)]
    -- Prepare an object method for calling.
    elseif opname == "SELF" then
      local rb = stack[b]
      stack[a + 1] = rb
      stack[a] = rb[rk(c)]

    -- OP_ADD [A, B, C]    R(A) := RK(B) + RK(C)
    -- Add two values and store the result in a register.
    elseif opname == "ADD" then
      stack[a] = rk(b) + rk(c)

    -- OP_SUB [A, B, C]    R(A) := RK(B) - RK(C)
    -- Subtract two values and store the result in a register.
    elseif opname == "SUB" then
      stack[a] = rk(b) - rk(c)

    -- OP_MUL [A, B, C]    R(A) := RK(B) * RK(C)
    -- Multiply two values and store the result in a register.
    elseif opname == "MUL" then
      stack[a] = rk(b) * rk(c)

    -- OP_DIV [A, B, C]    R(A) := RK(B) / RK(C)
    -- Divide two values and store the result in a register.
    elseif opname == "DIV" then
      stack[a] = rk(b) / rk(c)

    -- OP_MOD [A, B, C]    R(A) := RK(B) % RK(C)
    -- Calculate the modulus of two values and store the result in a register.
    elseif opname == "MOD" then
      stack[a] = rk(b) % rk(c)

    -- OP_POW [A, B, C]    R(A) := RK(B) ^ RK(C)
    -- Raise a value to the power of another and store the result in a register.
    elseif opname == "POW" then
      stack[a] = rk(b) ^ rk(c)

    -- OP_UNM [A, B]    R(A) := -R(B)
    -- Negate a value and store the result in a register.
    elseif opname == "UNM" then
      stack[a] = -stack[b]

    -- OP_NOT [A, B]    R(A) := not R(B)
    -- Logical NOT of a value and store the result in a register.
    elseif opname == "NOT" then
      stack[a] = not stack[b]

    -- OP_LEN [A, B]    R(A) := length of R(B)
    -- Get the length of a value and store it in a register.
    elseif opname == "LEN" then
      stack[a] = #stack[b]

    -- OP_CONCAT [A, B, C]    R(A) := R(B).. ... ..R(C)
    -- Concatenate a range of registers and store the result in a register.
    elseif opname == "CONCAT" then
      for reg = c - 1, b, -1 do
        -- NOTE: We cannot use table.concat as it does not call metamethods
        stack[reg] = stack[reg] .. stack[reg + 1]
      end
      stack[a] = stack[b]

    -- OP_JMP [A, sBx]    pc+=sBx
    -- Jump to a new instruction offset.
    elseif opname == "JMP" then
      programCounter = programCounter + b

    -- OP_EQ [A, B, C]    if ((RK(B) == RK(C)) ~= A) then pc++
    -- Conditional jump based on equality.
    elseif opname == "EQ" then
      if (rk(b) == rk(c)) ~= (a ~= 0) then
        programCounter = programCounter + 1
      end

    -- OP_LT [A, B, C]    if ((RK(B) < RK(C)) ~= A) then pc++
    -- Conditional jump based on less-than comparison.
    elseif opname == "LT" then
      if (rk(b) < rk(c)) ~= (a ~= 0) then
        programCounter = programCounter + 1
      end

    -- OP_LE [A, B, C]    if ((RK(B) <= RK(C)) ~= A) then pc++
    -- Conditional jump based on less-than-or-equal comparison.
    elseif opname == "LE" then
      if (rk(b) <= rk(c)) ~= (a ~= 0) then
        programCounter = programCounter + 1
      end

    -- OP_TEST [A, C]    if not (R(A) <=> C) then pc++
    -- Boolean test, with a conditional jump.
    -- NOTE: After this instruction, the next instruction will always be a jump.
    elseif opname == "TEST" then
      if (not stack[a]) == (c ~= 0) then
        programCounter = programCounter + 1
      else
        -- The Lua 5.1 spec guarantees that TEST/TFORLOOP is always followed by JMP.
        -- Since TLC's own compiler always emits this sequence, we trust the bytecode
        -- and skip the validation check for performance.
        local nextInstruction = protoCode[programCounter + 1]
        local jumpDistance = nextInstruction.b
        programCounter = programCounter + 1 + jumpDistance
      end

    -- OP_TESTSET [A, B, C]    if (R(B) <=> C) then R(A) := R(B) else pc++
    -- Boolean test, with conditional assignment and jump.
    -- NOTE: After this instruction, the next instruction will always be a jump.
    elseif opname == "TESTSET" then
      if (not stack[b]) == (c ~= 0) then
        stack[a] = stack[b]

        -- The Lua 5.1 spec guarantees that TEST/TFORLOOP is always followed by JMP.
        -- Since TLC's own compiler always emits this sequence, we trust the bytecode
        -- and skip the validation check for performance.
        local nextInstruction = protoCode[programCounter + 1]
        local jumpDistance = nextInstruction.b
        programCounter = programCounter + 1 + jumpDistance
      else
        programCounter = programCounter + 1
      end

    -- OP_CALL [A, B, C]    R(A), ... ,R(A+C-2) := R(A)(R(A+1), ... ,R(A+B-1))
    -- Call a closure (function) with arguments and handle returns.
    -- Call frame convention:
    --
    --   stack[A]     = callee
    --   stack[A + 1] = first argument
    --   ...
    --
    -- After CALL returns:
    --   stack[A] ... stack[A + C - 2] = return values
    elseif opname == "CALL" then
      local func = stack[a]
      if b ~= USE_CURRENT_TOP then
        stackTop = a + b
      end

      local nReturns, returns = pack(func(unpack(stack, a + 1, stackTop - 1)))

      if c ~= USE_CURRENT_TOP then
        nReturns = c - 1
      else
        -- Multi-return: push all results onto the stack.
        stackTop = a + nReturns
      end
      for index = 1, nReturns do
        stack[a + index - 1] = returns[index]
      end

    -- OP_TAILCALL [A, B, C]    return R(A)(R(A+1), ... ,R(A+B-1))
    -- Perform a tail call to a closure (function).
    elseif opname == "TAILCALL" then
      local func = stack[a]
      if b ~= USE_CURRENT_TOP then
        stackTop = a + b
      end

      -- Optimization:
      --   Since TAILCALL always comes before RETURN,
      --   we can just return the function call results directly.
      return func(unpack(stack, a + 1, stackTop - 1))

    -- OP_RETURN [A, B]    return R(A), ... ,R(A+B-2)
    -- Return values from function call.
    elseif opname == "RETURN" then
      if b ~= USE_CURRENT_TOP then
        stackTop = a + b - 1
      end

      return unpack(stack, a, stackTop - 1)

    -- OP_FORLOOP [A, sBx]   R(A)+=R(A+2)
    --                       if R(A) <?= R(A+1) then { pc+=sBx R(A+3)=R(A) }
    -- Iterate a numeric for loop.
    elseif opname == "FORLOOP" then
      local limit = stack[a + 1]
      local step  = stack[a + 2]
      local idx   = stack[a] + step

      local shouldContinue = false
      if step >= 0 then
        shouldContinue = (idx <= limit)
      else
        shouldContinue = (idx >= limit)
      end

      if shouldContinue then
        programCounter = programCounter + b
        stack[a] = idx -- The internal index register.
        stack[a + 3] = idx -- The user-visible control variable register.
      end

    -- OP_FORPREP [A, sBx]    R(A)-=R(A+2) pc+=sBx
    -- Initialize a numeric for loop.
    elseif opname == "FORPREP" then
      local init  = stack[a]
      local limit = stack[a + 1]
      local step  = stack[a + 2]

      if type(init) ~= "number" then
        error("'for' initial value must be a number")
      elseif type(limit) ~= "number" then
        error("'for' limit must be a number")
      elseif type(step) ~= "number" then
        error("'for' step must be a number")
      end

      stack[a] = init - step
      programCounter = programCounter + b

    -- OP_TFORLOOP [A, C]    R(A+3), ... ,R(A+2+C) := R(A)(R(A+1), R(A+2))
    --                       if R(A+3) ~= nil then R(A+2)=R(A+3) else pc++
    -- Iterate a generic for loop.
    -- NOTE: After this instruction, the next instruction will always be a jump.
    --
    -- Register layout:
    --
    --   R(A)   = iterator function
    --   R(A+1) = state
    --   R(A+2) = control variable
    --   R(A+3) ... = yielded values
    elseif opname == "TFORLOOP" then
      local cb = a + 3

      -- Copy function and arguments for the call.
      local control  = stack[a + 2]
      local state    = stack[a + 1]
      local iterator = stack[a]
      local results = { iterator(state, control) }

      -- Place results on the stack, starting at cb.
      -- The number of results to keep is specified by `c`.
      for i = cb, cb + c - 1 do
        stack[i] = results[i - cb + 1]
      end

      -- Check if the new control variable is nil.
      if stack[cb] ~= nil then
        -- If not nil, copy it to R(A+2) (the control variable).
        stack[cb - 1] = stack[cb]

        -- The Lua 5.1 spec guarantees that TEST/TFORLOOP is always followed by JMP.
        -- Since TLC's own compiler always emits this sequence, we trust the program
        -- and skip the validation check for performance.
        local nextInstruction = protoCode[programCounter + 1]
        local jumpDistance = nextInstruction.b
        programCounter = programCounter + 1 + jumpDistance
      else
        -- Otherwise, skip the next instruction.
        -- (Should be the jump back to the loop start.)
        programCounter = programCounter + 1
      end

    -- OP_SETLIST [A, B, C]    R(A)[(C-1)*FPF+i] := R(A+i), 1 <= i <= B
    -- Set a range of array elements in a table.
    elseif opname == "SETLIST" then
      local targetTable = stack[a]
      if type(targetTable) ~= "table" then
        error("SETLIST target is not a table")
      end

      local len = b
      if len == USE_CURRENT_TOP then
        len = stackTop - a - 1
      end

      -- If C is 0, it indicates that the next instruction is the C value itself.
      -- We can't implement it here without making the VM execute actual bytecode,
      -- so we will just throw an error if we encounter this case.
      if c == 0 then
        error("SETLIST with C=0 is not supported in this VM implementation")
      end

      local offset = (c - 1) * FIELDS_PER_FLUSH
      for i = 1, len do
        targetTable[offset + i] = stack[a + i]
      end

    -- OP_CLOSE [A]    close all variables in the stack up to (>=) R(A)
    elseif opname == "CLOSE" then
      for upvalIdx = a, highestOpenUpvalue do
        local upvalue = openUpvalues[upvalIdx]
        if upvalue then
          upvalue.stack = { stack[upvalue.index] }
          upvalue.index = 1
          openUpvalues[upvalIdx] = nil
        end
      end
      highestOpenUpvalue = a - 1

    -- OP_CLOSURE [A, Bx]    R(A) := closure(KPROTO[Bx], R(A), ... ,R(A+n))
    -- Create a new closure (function) and store it in a register.
    elseif opname == "CLOSURE" then
      local tProto = closureProto.protos[b + 1]
      local tProtoUpvalues = {}

      -- Consume pseudo-instructions that bind closure upvalues.
      -- These pseudo-instructions can only be:
      --    `MOVE` (for capturing a local variable)
      --    `GETUPVAL` (for capturing an enclosing upvalue)
      for upvalueIndex = 1, tProto.numUpvals do
        programCounter = programCounter + 1

        local nextInstruction = protoCode[programCounter]
        local nextOpname = nextInstruction.opname
        local source = nextInstruction.b

        if nextOpname == "MOVE" then
          local upvalue = openUpvalues[source] or {
            index = source,
            stack = stack
          }
          openUpvalues[source] = upvalue
          if source > highestOpenUpvalue then
            highestOpenUpvalue = source
          end
          tProtoUpvalues[upvalueIndex] = upvalue
        elseif nextOpname == "GETUPVAL" then
          local upvalue = closureUpvals[source + 1]
          tProtoUpvalues[upvalueIndex] = upvalue
        else
          error(
            "Unexpected instruction while capturing upvalues: "
            .. tostring(nextOpname)
          )
        end
      end

      local newClosure = {
        proto    = tProto,
        upvalues = tProtoUpvalues,
        env      = closureEnv
      }

      stack[a] = function(...)
        return VirtualMachine.executeClosure(newClosure, ...)
      end

    -- OP_VARARG [A, B]    R(A), R(A+1), ..., R(A+B-1) = vararg
    -- Load vararg function arguments into registers.
    elseif opname == "VARARG" then
      local upperBound = b - 1
      if b == USE_CURRENT_TOP then
        stackTop = a + varargCount
        upperBound = stackTop
      end

      for i = 1, upperBound do
        stack[a + i - 1] = vararg[i]
      end
    else
      error("Unimplemented instruction: " .. tostring(opname))
    end

    -- Advance to the next instruction.
    programCounter = programCounter + 1
  end

  -- Should never be reached, as every valid Lua function prototype should
  -- always end with a RETURN instruction that exits the function.
  return
end

function VirtualMachine.execute(proto, env, ...)
  assert(type(proto) == "table", "VirtualMachine.execute: proto must be a table")
  assert(type(env) == "table" or env == nil,
    "VirtualMachine.execute: env must be a table or nil"
  )

  return VirtualMachine.executeClosure({
    env      = env or _G,
    proto    = proto,
    upvalues = {}
  }, ...)
end

-- Code (string) -> Tokens (list)
local function tokenize(code)
  return Tokenizer.new(code):tokenize()
end

-- Tokens (list) -> AST (table)
local function parseTokens(tokens)
  return Parser.new(tokens):parse()
end

-- AST (table) -> Proto (table)
local function generate(ast)
  return CodeGenerator.new(ast):generate()
end

-- Proto (table) -> Bytecode (string)
local function emit(proto)
  return BytecodeEmitter.emit(proto)
end

-- Proto (table) -> Execution (results...)
local function execute(proto, env, ...)
  return VirtualMachine.execute(proto, env, ...)
end

-- Code -> AST
local function parse(code)
  return parseTokens(tokenize(code))
end

-- Code -> Proto
local function compileToProto(code)
  return generate(parse(code))
end

-- Code -> Bytecode
local function compile(code)
  return emit(compileToProto(code))
end

-- Code -> Execution
local function run(code, env, ...)
  return execute(compileToProto(code), env, ...)
end

-- Public API module table.
-- High-level functions wrap the full pipeline for common use cases.
-- Low-level functions expose individual stages for composition.

return {
  -- ████████╗  ██╗        ██████╗
  -- ╚══██╔══╝  ██║       ██╔════╝
  --    ██║     ██║       ██║
  --    ██║     ██║       ██║
  --    ██║     ███████╗  ╚██████╗
  --    ╚═╝     ╚══════╝   ╚═════╝
  --       Tiny Lua Compiler
  --          MIT License

  -- Raw components (classes)
  Tokenizer       = Tokenizer,
  Parser          = Parser,
  CodeGenerator   = CodeGenerator,
  BytecodeEmitter = BytecodeEmitter,
  VirtualMachine  = VirtualMachine,

  -- Standard API (High-Level)
  tokenize       = tokenize,       -- Code -> Tokens
  parse          = parse,          -- Code -> AST
  compileToProto = compileToProto, -- Code -> Proto
  compile        = compile,        -- Code -> Bytecode
  run            = run,            -- Code -> Execution

  -- Low-Level API (Step-by-Step)
  tokenize    = tokenize,    -- Code   -> Tokens
  parseTokens = parseTokens, -- Tokens -> AST
  generate    = generate,    -- AST    -> Proto
  emit        = emit,        -- Proto  -> Bytecode
  execute     = execute,     -- Proto  -> Execution
}