lr_parser.py

"""
LR Shift-Reduce Parser (Bottom-Up)
Simulates LR(1) parsing using a hand-crafted precedence climbing approach
with an explicit stack and action log, so students can see the shift/reduce steps.
"""

from typing import List, Tuple, Optional
from lexer import Token, TokenType, Lexer
from ast_nodes import *


class LRParseError(Exception):
    def __init__(self, message, token=None):
        super().__init__(message)
        self.token = token


# Operator precedence table (higher = tighter binding)
PRECEDENCE = {
    'or':  (1, 'left'),
    'and': (2, 'left'),
    '==':  (3, 'left'),
    '!=':  (3, 'left'),
    '<':   (4, 'left'),
    '>':   (4, 'left'),
    '<=':  (4, 'left'),
    '>=':  (4, 'left'),
    '+':   (5, 'left'),
    '-':   (5, 'left'),
    '*':   (6, 'left'),
    '/':   (6, 'left'),
    '%':   (6, 'left'),
    '**':  (7, 'right'),
}

UNARY_OPS  = {'-', 'not'}
BINARY_OPS = set(PRECEDENCE.keys())


class LRParser:
    """
    Precedence-climbing LR-style parser.
    Maintains an explicit value-stack and operator-stack (like a real LR parser's
    state/symbol stack) and records every shift and reduce action.
    """

    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0
        self.parse_steps: List[str] = []  # shift/reduce log

    @classmethod
    def from_source(cls, source: str) -> 'LRParser':
        lexer = Lexer(source)
        tokens = lexer.tokenize()
        return cls(tokens)

    # ── Helpers ───────────────────────────────────────────────────────────────

    def current(self) -> Token:
        return self.tokens[self.pos]

    def peek(self, offset=1) -> Token:
        idx = self.pos + offset
        return self.tokens[min(idx, len(self.tokens) - 1)]

    def check(self, *types: TokenType) -> bool:
        return self.current().type in types

    def consume(self, *types: TokenType) -> Token:
        tok = self.current()
        if types and tok.type not in types:
            raise LRParseError(
                f"Expected {[t.name for t in types]} but got {tok.type.name} '{tok.value}'",
                tok
            )
        self.pos += 1
        return tok

    def shift(self, description: str):
        tok = self.current()
        self.parse_steps.append(
            f"SHIFT   {tok.type.name:15s} ← '{tok.value}'  │ {description}"
        )

    def reduce(self, rule: str):
        self.parse_steps.append(f"REDUCE  {rule}")

    # ── Main entry ────────────────────────────────────────────────────────────

    def parse(self) -> ProgramNode:
        self.parse_steps.append("─── LR Shift-Reduce Parse Begin ───")
        stmts = []
        while not self.check(TokenType.EOF):
            stmts.append(self.statement())
        self.reduce("program → statement*")
        return ProgramNode(stmts)

    # ── Statements ────────────────────────────────────────────────────────────

    def statement(self) -> ASTNode:
        if self.check(TokenType.VAR):
            return self.var_decl()
        elif self.check(TokenType.FUNCTION):
            return self.func_def()
        elif self.check(TokenType.IF):
            return self.if_stmt()
        elif self.check(TokenType.WHILE):
            return self.while_stmt()
        elif self.check(TokenType.FOR):
            return self.for_stmt()
        elif self.check(TokenType.RETURN):
            return self.return_stmt()
        elif self.check(TokenType.PRINT):
            return self.print_stmt()
        elif self.check(TokenType.LBRACE):
            return self.block()
        else:
            return self.expr_stmt()

    def var_decl(self) -> VarDeclNode:
        self.shift("varDecl")
        self.consume(TokenType.VAR)
        name_tok = self.consume(TokenType.IDENTIFIER)
        name = name_tok.value
        value = None
        if self.check(TokenType.ASSIGN):
            self.shift("assignment op")
            self.consume(TokenType.ASSIGN)
            value = self.expr_bp(0)
        self.consume(TokenType.SEMICOLON)
        self.reduce(f"varDecl → 'var' {name} = expr")
        return VarDeclNode(name, value)

    def func_def(self) -> FunctionDefNode:
        self.shift("funcDef")
        self.consume(TokenType.FUNCTION)
        name = self.consume(TokenType.IDENTIFIER).value
        self.consume(TokenType.LPAREN)
        params = []
        if not self.check(TokenType.RPAREN):
            params.append(self.consume(TokenType.IDENTIFIER).value)
            while self.check(TokenType.COMMA):
                self.consume(TokenType.COMMA)
                params.append(self.consume(TokenType.IDENTIFIER).value)
        self.consume(TokenType.RPAREN)
        body = self.block()
        self.reduce(f"funcDef → function {name}({', '.join(params)})")
        return FunctionDefNode(name, params, body)

    def if_stmt(self) -> IfNode:
        self.shift("ifStmt")
        self.consume(TokenType.IF)
        self.consume(TokenType.LPAREN)
        cond = self.expr_bp(0)
        self.consume(TokenType.RPAREN)
        then = self.block()
        else_ = None
        if self.check(TokenType.ELSE):
            self.shift("else branch")
            self.consume(TokenType.ELSE)
            else_ = self.block()
        self.reduce("ifStmt → if (expr) block [else block]")
        return IfNode(cond, then, else_)

    def while_stmt(self) -> WhileNode:
        self.shift("whileStmt")
        self.consume(TokenType.WHILE)
        self.consume(TokenType.LPAREN)
        cond = self.expr_bp(0)
        self.consume(TokenType.RPAREN)
        body = self.block()
        self.reduce("whileStmt → while (expr) block")
        return WhileNode(cond, body)

    def for_stmt(self) -> ForNode:
        self.shift("forStmt")
        self.consume(TokenType.FOR)
        self.consume(TokenType.LPAREN)
        init = self.var_decl() if self.check(TokenType.VAR) else None
        if init is None:
            self.consume(TokenType.SEMICOLON)
        cond = self.expr_bp(0) if not self.check(TokenType.SEMICOLON) else None
        self.consume(TokenType.SEMICOLON)
        update = self.expr_bp(0) if not self.check(TokenType.RPAREN) else None
        self.consume(TokenType.RPAREN)
        body = self.block()
        self.reduce("forStmt → for(init; cond; update) block")
        return ForNode(init, cond, update, body)

    def return_stmt(self) -> ReturnNode:
        self.shift("returnStmt")
        self.consume(TokenType.RETURN)
        val = None
        if not self.check(TokenType.SEMICOLON):
            val = self.expr_bp(0)
        self.consume(TokenType.SEMICOLON)
        self.reduce("returnStmt → return expr?")
        return ReturnNode(val)

    def print_stmt(self) -> PrintNode:
        self.shift("printStmt")
        self.consume(TokenType.PRINT)
        self.consume(TokenType.LPAREN)
        val = self.expr_bp(0)
        self.consume(TokenType.RPAREN)
        self.consume(TokenType.SEMICOLON)
        self.reduce("printStmt → print(expr)")
        return PrintNode(val)

    def block(self) -> BlockNode:
        self.shift("block {")
        self.consume(TokenType.LBRACE)
        stmts = []
        while not self.check(TokenType.RBRACE) and not self.check(TokenType.EOF):
            stmts.append(self.statement())
        self.consume(TokenType.RBRACE)
        self.reduce(f"block → {{ {len(stmts)} statements }}")
        return BlockNode(stmts)

    def expr_stmt(self) -> ASTNode:
        node = self.expr_bp(0)
        self.consume(TokenType.SEMICOLON)
        self.reduce("exprStmt → expr ;")
        return node

    # ── Pratt / precedence-climbing expression parser ─────────────────────────
    # This IS the LR heart: we shift tokens and reduce when precedence says so.

    def _tok_op(self, tok: Token) -> Optional[str]:
        """Return operator string if token is a binary operator, else None."""
        op_map = {
            TokenType.PLUS: '+', TokenType.MINUS: '-',
            TokenType.STAR: '*', TokenType.SLASH: '/',
            TokenType.PERCENT: '%', TokenType.POWER: '**',
            TokenType.EQ: '==', TokenType.NEQ: '!=',
            TokenType.LT: '<',  TokenType.GT: '>',
            TokenType.LEQ: '<=', TokenType.GEQ: '>=',
            TokenType.AND: 'and', TokenType.OR: 'or',
        }
        return op_map.get(tok.type)

    def expr_bp(self, min_bp: int) -> ASTNode:
        """Binding-power (Pratt) expression parser — models LR reduce decisions."""
        # Prefix / unary
        tok = self.current()

        # Unary minus
        if tok.type == TokenType.MINUS:
            self.shift("unary -")
            self.consume()
            operand = self.expr_bp(100)  # high bp = right-binding
            self.reduce("unary → '-' expr")
            left = UnaryOpNode('-', operand)

        # Unary not
        elif tok.type == TokenType.NOT:
            self.shift("unary not")
            self.consume()
            operand = self.expr_bp(100)
            self.reduce("unary → 'not' expr")
            left = UnaryOpNode('not', operand)

        # Grouped expression
        elif tok.type == TokenType.LPAREN:
            self.shift("( grouped expr")
            self.consume()
            left = self.expr_bp(0)
            self.consume(TokenType.RPAREN)
            self.reduce("primary → '(' expr ')'")

        # Literals / identifiers
        else:
            left = self.primary()

        # Infix loop — the shift-reduce core
        while True:
            tok = self.current()

            # Function call
            if tok.type == TokenType.LPAREN and isinstance(left, IdentifierNode):
                self.shift("call (")
                self.consume()
                args = []
                if not self.check(TokenType.RPAREN):
                    args.append(self.expr_bp(0))
                    while self.check(TokenType.COMMA):
                        self.consume(TokenType.COMMA)
                        args.append(self.expr_bp(0))
                self.consume(TokenType.RPAREN)
                self.reduce(f"call → {left.name}({len(args)} args)")
                left = FunctionCallNode(left.name, args)
                continue

            # Assignment: IDENT '=' ...
            if tok.type == TokenType.ASSIGN and isinstance(left, IdentifierNode):
                if min_bp > 0:
                    break
                self.shift("assign =")
                self.consume()
                right = self.expr_bp(0)
                self.reduce(f"assign → {left.name} = expr")
                left = AssignNode(left.name, right)
                continue

            op = self._tok_op(tok)
            if op is None:
                break

            prec, assoc = PRECEDENCE[op]
            if prec < min_bp:
                break  # REDUCE: current operator has lower priority → stop

            self.shift(f"binop '{op}'")
            self.consume()
            # For right-associative ops (like **), right_bp = prec
            # For left-associative ops, right_bp = prec + 1
            right_bp = prec if assoc == 'right' else prec + 1
            right = self.expr_bp(right_bp)
            self.reduce(f"binop → expr '{op}' expr")
            left = BinaryOpNode(op, left, right)

        return left

    def primary(self) -> ASTNode:
        tok = self.current()
        if tok.type == TokenType.NUMBER:
            self.shift("number")
            self.consume()
            self.reduce(f"primary → NUMBER({tok.value})")
            return NumberNode(float(tok.value))
        if tok.type == TokenType.STRING:
            self.shift("string")
            self.consume()
            self.reduce(f"primary → STRING({tok.value})")
            return StringNode(tok.value.strip('"\''))
        if tok.type == TokenType.BOOLEAN:
            self.shift("boolean")
            self.consume()
            self.reduce(f"primary → BOOLEAN({tok.value})")
            return BooleanNode(tok.value == 'true')
        if tok.type == TokenType.IDENTIFIER:
            self.shift("identifier")
            self.consume()
            self.reduce(f"primary → ID({tok.value})")
            return IdentifierNode(tok.value)
        raise LRParseError(
            f"Unexpected token '{tok.value}' ({tok.type.name}) at line {tok.line}",
            tok
        )