# The parser: tokens in, an abstract syntax tree out.
#
# A recursive-descent parser: one def per rule of the grammar, each reading
# tokens from the front and returning the node it built. The nodes live in
# one array, the arena; a node names its children by their index in it.
# Every parse def takes the parser state and returns the new state and the
# index of the node it made: `p, n = expr(p)`.
#
# The grammar so far:
#
#     program     = function* eof
#     function    = "int" ident "(" params? ")" "{" compound_stmt
#     params      = "int" ident ("," "int" ident)*
#     declaration = "int" ident ("=" expr)? ("," ident ("=" expr)?)* ";"
#     stmt        = "return" expr ";"
#                 | "if" "(" expr ")" stmt ("else" stmt)?
#                 | "for" "(" expr? ";" expr? ";" expr? ")" stmt
#                 | "while" "(" expr ")" stmt
#                 | "{" compound_stmt
#                 | expr? ";"
#     compound_stmt = (declaration | stmt)* "}"
#     expr        = assign
#     assign      = logor ("=" assign)?
#     logor       = logand ("||" logand)*
#     logand      = equality ("&&" equality)*
#     equality   = relational ("==" relational | "!=" relational)*
#     relational = add ("<" add | "<=" add | ">" add | ">=" add)*
#     add     = mul ("+" mul | "-" mul)*
#     mul     = unary ("*" unary | "/" unary | "%" unary)*
#     unary   = ("+" | "-" | "!") unary | primary
#     primary = "(" expr ")" | ident "(" args? ")" | ident | num
#     args    = assign ("," assign)*
#
# Precedence comes from the nesting: mul is called from add, so * binds
# tighter than +.

from cc.lexer import Token, TokKind, tokenize

enum NodeKind {
    Num         # an integer literal: val
    Add         # lhs + rhs
    Sub         # lhs - rhs
    Mul         # lhs * rhs
    Div         # lhs / rhs
    Mod         # lhs % rhs
    Neg         # -lhs
    Not         # !lhs: 1 when lhs is 0, else 0
    LogAnd      # lhs && rhs: rhs is computed only when lhs is not 0
    LogOr       # lhs || rhs: rhs is computed only when lhs is 0
    Eq          # lhs == rhs: 1 or 0
    Ne          # lhs != rhs
    Lt          # lhs < rhs
    Le          # lhs <= rhs (a > b and a >= b are made as b < a and b <= a)
    Var         # a local variable: val is its slot in the function's locals
    Assign      # lhs = rhs, lhs a Var
    Call        # name(args): a call; body holds the argument expressions
    ExprStmt    # lhs ";": an expression computed for its effect
    Return      # "return" lhs ";"
    If          # "if" (cond) then (else els)?
    For         # "for" (init; cond; inc) then, and "while" (cond) then
    Block       # a list of statements: body
}

#| One node of the tree. Fields a kind does not use keep their defaults.
record Node {
    kind: NodeKind
    tok: Int            # the token it was made from, for error messages
    val: Int = 0        # Num: the value
    name: Str = ""      # Call: the function called
    lhs: Int = -1       # the left operand
    rhs: Int = -1       # the right operand
    body: [Int] = []    # Block: the statements; Call: the arguments
    cond: Int = -1      # If: the condition
    then: Int = -1      # If: the statement run when it holds
    els: Int = -1       # If: the statement run when it does not, or -1
    init: Int = -1      # For: run once before the loop, or -1
    inc: Int = -1       # For: run after each pass, or -1 (cond -1: loop forever)
}

#| A function: its name and the node of its body.
record Function {
    name: Str
    body: Int
    locals: [Str]       # the local variables' names; a variable's slot is its index here
    nparams: Int = 0    # the first nparams locals are the parameters
}

#| The parsed program: the node arena and the functions.
record Program {
    nodes: [Node]
    funcs: [Function]
}

#| The parser's state: the tokens, the next token's index, and the nodes
#| made so far.
record Parser {
    toks: [Token]
    pos: Int
    nodes: [Node]
    locals: [Str]       # the variables declared so far in the current function
}

# ---- reading tokens

#| The token the parser stands on, not consumed.
def peek(p: Parser) -> Token {
    return p.toks[p.pos]
}

#| The parser moved past one token.
def advance(p: Parser) -> Parser {
    q = p
    q.pos = p.pos + 1
    return q
}

#| Is the current token the punctuator or keyword s?
def equal(p: Parser, s: Str) -> Bool {
    t = peek(p)
    match t.kind {
        Num | Eof | Ident => { return false }
        _ => { return t.text == s }
    }
    return false
}

#| The parser past the punctuator or keyword s, which must be next.
def skip(p: Parser, s: Str) -> Parser {
    if not equal(p, s) {
        error_tok(p, f"expected `{s}`")
    }
    return advance(p)
}

#| The slot of local variable `name`, or -1 when it is not declared.
def find_local(p: Parser, name: Str) -> Int {
    for k to len(p.locals) {
        if p.locals[k] == name {
            return int(k)
        }
    }
    return -1
}

#| A node of kind k with operands lhs and rhs, made at token `tok`.
def binary(k: NodeKind, lhs: Int, rhs: Int, tok: Int) -> Node {
    return Node { kind: k, tok: tok, lhs: lhs, rhs: rhs }
}

#| Stop with a message pointing at the parser's current token.
def error_tok(p: Parser, msg: Str) {
    t = peek(p)
    stop(f"{t.line}:{t.col}: {msg} (at `{t.text}`)")
}

#| Add a node to the arena; the new state and the node's index.
def add_node(p: Parser, n: Node) {
    q = p
    index: Int = len(p.nodes)
    q.nodes = append(p.nodes, n)
    return q, index
}

# ---- the grammar, one def per rule

#| declaration = "int" ident ("=" expr)? ("," ident ("=" expr)?)* ";".
#| Each name gets the next slot; each initializer becomes an assignment
#| statement, and the declaration is the Block of them.
def declaration(p: Parser) {
    at = p.pos
    p = skip(p, "int")
    inits: [Int] = []
    first = true
    while not equal(p, ";") {
        if not first {
            p = skip(p, ",")
        }
        first = false
        t = peek(p)
        match t.kind {
            Ident => { p = p }
            _ => { error_tok(p, "expected a variable name") }
        }
        if find_local(p, t.text) >= 0 {
            error_tok(p, f"`{t.text}` is already declared")
        }
        slot: Int = len(p.locals)
        q = advance(p)
        q.locals = append(p.locals, t.text)
        p = q
        if equal(p, "=") {
            p, var = add_node(p, Node { kind: NodeKind.Var, tok: p.pos - 1, val: slot })
            eq = p.pos
            p, value = assign(advance(p))
            p, set = add_node(p, binary(NodeKind.Assign, var, value, eq))
            p, s = add_node(p, Node { kind: NodeKind.ExprStmt, tok: eq, lhs: set })
            inits = append(inits, s)
        }
    }
    return add_node(skip(p, ";"), Node { kind: NodeKind.Block, tok: at, body: inits })
}

#| compound_stmt = (declaration | stmt)* "}": the statements of a block,
#| after its "{".
def compound_stmt(p: Parser) {
    at = p.pos
    stmts: [Int] = []
    while not equal(p, "}") {
        if equal(p, "int") {
            p, s = declaration(p)
        }
        else {
            p, s = stmt(p)
        }
        stmts = append(stmts, s)
    }
    return add_node(advance(p), Node { kind: NodeKind.Block, tok: at, body: stmts })
}

#| stmt = "return" expr ";" | "{" compound_stmt | expr? ";"
def stmt(p: Parser) {
    at = p.pos
    if equal(p, "return") {
        p, value = expr(advance(p))
        return add_node(skip(p, ";"), Node { kind: NodeKind.Return, tok: at, lhs: value })
    }
    if equal(p, "if") {
        p = skip(advance(p), "(")
        p, c = expr(p)
        p, t = stmt(skip(p, ")"))
        e: Int = -1
        if equal(p, "else") {
            p, e = stmt(advance(p))
        }
        return add_node(p, Node { kind: NodeKind.If, tok: at, cond: c, then: t, els: e })
    }
    if equal(p, "for") {
        #: each of the three parts may be left out
        p = skip(advance(p), "(")
        i: Int = -1
        c: Int = -1
        n: Int = -1
        if not equal(p, ";") {
            p, e = expr(p)
            p, i = add_node(p, Node { kind: NodeKind.ExprStmt, tok: at, lhs: e })
        }
        p = skip(p, ";")
        if not equal(p, ";") {
            p, c = expr(p)
        }
        p = skip(p, ";")
        if not equal(p, ")") {
            p, e = expr(p)
            p, n = add_node(p, Node { kind: NodeKind.ExprStmt, tok: at, lhs: e })
        }
        p, body = stmt(skip(p, ")"))
        return add_node(p, Node { kind: NodeKind.For, tok: at, init: i, cond: c, inc: n, then: body })
    }
    if equal(p, "while") {
        #: a while loop is a for loop with only a condition
        p = skip(advance(p), "(")
        p, c = expr(p)
        p, body = stmt(skip(p, ")"))
        return add_node(p, Node { kind: NodeKind.For, tok: at, cond: c, then: body })
    }
    if equal(p, "{") {
        return compound_stmt(advance(p))
    }
    if equal(p, ";") {
        #: the null statement: an empty block
        return add_node(advance(p), Node { kind: NodeKind.Block, tok: at })
    }
    p, e = expr(p)
    return add_node(skip(p, ";"), Node { kind: NodeKind.ExprStmt, tok: at, lhs: e })
}

#| expr = assign
def expr(p: Parser) {
    return assign(p)
}

#| assign = logor ("=" assign)?. Recursing on the right makes = group
#| to the right: a = b = 3 is a = (b = 3).
def assign(p: Parser) {
    p, node = logor(p)
    if equal(p, "=") {
        match p.nodes[node].kind {
            Var => { p = p }
            _ => { error_tok(p, "the left side of = is not a variable") }
        }
        at = p.pos
        p, value = assign(advance(p))
        p, node = add_node(p, binary(NodeKind.Assign, node, value, at))
    }
    return p, node
}

#| logor = logand ("||" logand)*
def logor(p: Parser) {
    p, node = logand(p)
    while equal(p, "||") {
        at = p.pos
        p, rhs = logand(advance(p))
        p, node = add_node(p, binary(NodeKind.LogOr, node, rhs, at))
    }
    return p, node
}

#| logand = equality ("&&" equality)*
def logand(p: Parser) {
    p, node = equality(p)
    while equal(p, "&&") {
        at = p.pos
        p, rhs = equality(advance(p))
        p, node = add_node(p, binary(NodeKind.LogAnd, node, rhs, at))
    }
    return p, node
}

#| equality = relational ("==" relational | "!=" relational)*
def equality(p: Parser) {
    p, node = relational(p)
    while equal(p, "==") or equal(p, "!=") {
        at = p.pos
        kind = NodeKind.Eq
        if equal(p, "!=") {
            kind = NodeKind.Ne
        }
        p, rhs = relational(advance(p))
        p, node = add_node(p, binary(kind, node, rhs, at))
    }
    return p, node
}

#| relational = add ("<" add | "<=" add | ">" add | ">=" add)*. There are
#| no Gt and Ge nodes: a > b is built as b < a, with the operands swapped.
def relational(p: Parser) {
    p, node = add(p)
    while equal(p, "<") or equal(p, "<=") or equal(p, ">") or equal(p, ">=") {
        at = p.pos
        op = peek(p).text
        p, rhs = add(advance(p))
        if op == "<" {
            p, node = add_node(p, binary(NodeKind.Lt, node, rhs, at))
        }
        else if op == "<=" {
            p, node = add_node(p, binary(NodeKind.Le, node, rhs, at))
        }
        else if op == ">" {
            p, node = add_node(p, binary(NodeKind.Lt, rhs, node, at))
        }
        else {
            p, node = add_node(p, binary(NodeKind.Le, rhs, node, at))
        }
    }
    return p, node
}

#| add = mul ("+" mul | "-" mul)*. The loop makes the operators group to
#| the left: 1 - 2 - 3 is (1 - 2) - 3.
def add(p: Parser) {
    p, node = mul(p)
    while equal(p, "+") or equal(p, "-") {
        at = p.pos
        kind = NodeKind.Add
        if equal(p, "-") {
            kind = NodeKind.Sub
        }
        p, rhs = mul(advance(p))
        p, node = add_node(p, binary(kind, node, rhs, at))
    }
    return p, node
}

#| mul = unary ("*" unary | "/" unary | "%" unary)*
def mul(p: Parser) {
    p, node = unary(p)
    while equal(p, "*") or equal(p, "/") or equal(p, "%") {
        at = p.pos
        kind = NodeKind.Mul
        if equal(p, "/") {
            kind = NodeKind.Div
        }
        if equal(p, "%") {
            kind = NodeKind.Mod
        }
        p, rhs = unary(advance(p))
        p, node = add_node(p, binary(kind, node, rhs, at))
    }
    return p, node
}

#| unary = ("+" | "-" | "!") unary | primary. Unary plus changes nothing.
def unary(p: Parser) {
    if equal(p, "+") {
        return unary(advance(p))
    }
    if equal(p, "-") {
        at = p.pos
        p, operand = unary(advance(p))
        return add_node(p, Node { kind: NodeKind.Neg, tok: at, lhs: operand })
    }
    if equal(p, "!") {
        at = p.pos
        p, operand = unary(advance(p))
        return add_node(p, Node { kind: NodeKind.Not, tok: at, lhs: operand })
    }
    return primary(p)
}

#| primary = "(" expr ")" | num
def primary(p: Parser) {
    if equal(p, "(") {
        p, node = expr(advance(p))
        return skip(p, ")"), node
    }
    t = peek(p)
    match t.kind {
        Num => { return add_node(advance(p), Node { kind: NodeKind.Num, tok: p.pos, val: t.val }) }
        Ident => {
            if equal(advance(p), "(") {
                #: a name followed by ( is a call; its arguments are assignments, so , separates them
                at = p.pos
                p = advance(advance(p))
                args: [Int] = []
                while not equal(p, ")") {
                    if len(args) > 0 {
                        p = skip(p, ",")
                    }
                    p, a = assign(p)
                    args = append(args, a)
                }
                return add_node(advance(p), Node { kind: NodeKind.Call, tok: at, name: t.text, body: args })
            }
            slot = find_local(p, t.text)
            if slot < 0 {
                error_tok(p, f"undeclared variable `{t.text}`")
            }
            return add_node(advance(p), Node { kind: NodeKind.Var, tok: p.pos, val: slot })
        }
        _ => { error_tok(p, "expected an expression") }
    }
    return p, -1
}

#| Has the parser consumed every token but the end marker?
def at_eof(p: Parser) -> Bool {
    match peek(p).kind {
        Eof => { return true }
        _ => { return false }
    }
    return false
}

#| The parser past a name, which must be next; the name.
def ident(p: Parser) {
    t = peek(p)
    match t.kind {
        Ident => { return advance(p), t.text }
        _ => { error_tok(p, "expected a name") }
    }
    return p, ""
}

#| function = "int" ident "(" params? ")" "{" compound_stmt. Each function
#| starts with no local variables; its parameters are its first locals.
def function(p: Parser) {
    p = skip(p, "int")
    p, name = ident(p)
    p = skip(p, "(")
    p = Parser { locals: ["x"][:0], ..p }
    while not equal(p, ")") {
        if len(p.locals) > 0 {
            p = skip(p, ",")
        }
        p, param = ident(skip(p, "int"))
        if find_local(p, param) >= 0 {
            error_tok(p, f"parameter `{param}` appears twice")
        }
        names = p.locals
        names = append(names, param)
        p = Parser { locals: names, ..p }
    }
    nparams: Int = len(p.locals)
    p, body = compound_stmt(skip(advance(p), "{"))
    return p, Function { name: name, body: body, locals: p.locals, nparams: nparams }
}

#| program = function* eof
def parse(toks: [Token]) -> Program {
    no_nodes: [Node] = []
    p = Parser { toks: toks, pos: 0, nodes: no_nodes, locals: ["x"][:0] }
    funcs: [Function] = []
    while not at_eof(p) {
        p, f = function(p)
        funcs = append(funcs, f)
    }
    return Program { nodes: p.nodes, funcs: funcs }
}

test "numbers and left-grouped + and -" {
    prog = parse(tokenize("int main() { 42; }"))
    assert_true(len(prog.nodes) == 3)
    assert_true(prog.nodes[0].val == 42)
    prog = parse(tokenize("int main() { 1 - 2 + 3; }"))
    #: nodes: 1, 2, Sub, 3, Add, ExprStmt, Block
    top = prog.nodes[prog.nodes[5].lhs]
    match top.kind {
        Add => { assert_true(prog.nodes[prog.nodes[top.lhs].lhs].val == 1) }
        _ => { stop("1 - 2 + 3 should be (1 - 2) + 3") }
    }
}

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: 2019 Rui Ueyama
# SPDX-FileCopyrightText: 2026 The Viv authors