"""bytecode-vm -- one expression tree, three back ends, the same work.

`tiny-interpreter` held the evaluator fixed and varied the tree. This toy
holds the *tree* fixed and varies what consumes it: `evaluate` walks it,
`compile_expr` + `vm_run` flatten it to opcodes and run a dispatch loop, and
`closure_compile` turns it into nested closures with no opcodes at all.

One instruction per AST node means *how* the work is represented varies while
*how much* work there is is held fixed. The only thing here that changes the
instruction COUNT is `fold_constants`, which removes nodes before compiling.

The language is `tiny-interpreter`'s plus variables -- which is what makes
"compile once, run many times" askable. No clock and no randomness live here:
a result is a function of (source, env).
"""

import re

# ---------------------------------------------------------------- tokenizer

TOKEN = re.compile(r"\s*(?:(\d+)|([A-Za-z]\w*)|([-+*/()]))")


def tokenize(source):
    """`2-x` -> [('num', 2), ('-', '-'), ('name', 'x'), ('eof', None)].

    A symbol carries itself as its kind, so the parser can test a token
    straight against `LEVELS`, and the trailing `eof` makes `peek` total.
    """
    source, tokens, pos = source.strip(), [], 0
    while pos < len(source):
        match = TOKEN.match(source, pos)
        if match is None:
            raise SyntaxError(f"stray {source[pos]!r} at offset {pos}")
        number, name, symbol = match.groups()
        pos = match.end()
        tokens.append(("num", int(number)) if number is not None else
                      ("name", name) if name is not None else (symbol, symbol))
    tokens.append(("eof", None))
    return tokens


# ---------------------------------------------------------------- the tree

class Num:
    __slots__ = ("value",)

    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return str(self.value)


class Var:
    __slots__ = ("name",)

    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return self.name


class Neg:
    __slots__ = ("operand",)

    def __init__(self, operand):
        self.operand = operand

    def __repr__(self):
        return f"(-{self.operand!r})"


class Bin:
    """`left` and `right` are NAMED fields -- which is exactly what postfix
    code stops being able to say."""
    __slots__ = ("op", "left", "right")

    def __init__(self, op, left, right):
        self.op, self.left, self.right = op, left, right

    def __repr__(self):
        return f"({self.left!r} {self.op} {self.right!r})"


def count_nodes(node):
    """The number the headline compares `len(code)` against."""
    if isinstance(node, (Num, Var)):
        return 1
    if isinstance(node, Neg):
        return 1 + count_nodes(node.operand)
    return 1 + count_nodes(node.left) + count_nodes(node.right)


# ---------------------------------------------------------------- parser

LEVELS = (("+", "-"), ("*", "/"))   # loosest first: the index IS precedence


class Parser:
    """Recursive descent, left-folding at every level. Associativity is
    `tiny-interpreter`'s subject, and is deliberately fixed here."""

    def __init__(self, tokens):
        self.tokens, self.pos = tokens, 0

    def peek(self):
        return self.tokens[self.pos][0]

    def take(self):
        token = self.tokens[self.pos]
        self.pos += 1
        return token

    def parse(self):
        node = self.binary(0)
        if self.peek() != "eof":
            raise SyntaxError(f"trailing {self.peek()!r}")
        return node

    def tighter(self, level):
        if level + 1 < len(LEVELS):
            return self.binary(level + 1)
        return self.unary()

    def binary(self, level):
        node = self.tighter(level)
        while self.peek() in LEVELS[level]:
            op = self.take()[0]
            node = Bin(op, node, self.tighter(level))
        return node

    def unary(self):
        if self.peek() == "-":
            self.take()
            return Neg(self.unary())
        return self.atom()

    def atom(self):
        kind, value = self.take()
        if kind == "num":
            return Num(value)
        if kind == "name":
            return Var(value)
        if kind == "(":
            node = self.binary(0)
            if self.take()[0] != ")":
                raise SyntaxError("unclosed '('")
            return node
        raise SyntaxError(f"expected a number, name or '(', got {kind!r}")


def parse(source):
    return Parser(tokenize(source)).parse()


# ------------------------------------------------- back end 1: walk the tree

def evaluate(node, env):
    """Post-order walk: one Python call per AST node, nothing to compile."""
    if isinstance(node, Num):
        return node.value
    if isinstance(node, Var):
        return env[node.name]
    if isinstance(node, Neg):
        return -evaluate(node.operand, env)
    left, right = evaluate(node.left, env), evaluate(node.right, env)
    if node.op == "+":
        return left + right
    if node.op == "-":
        return left - right
    if node.op == "*":
        return left * right
    return left / right


# ------------------------------------------- back end 2: compile, then loop

PUSH, LOAD, ADD, SUB, MUL, DIV, NEG = range(7)
BINOP = {"+": ADD, "-": SUB, "*": MUL, "/": DIV}


def fold_constants(node):
    """A pass over the TREE, before any code exists: every subtree with no
    variable in it collapses to one `Num`. The only thing here that changes
    how many operations there are, and a separate pass on purpose -- fold and
    emit in one recursion and a node whose left operand folded appends that
    operand's PUSH *after* the right subtree's code (demo.py, section 7).
    """
    if isinstance(node, (Num, Var)):
        return node
    if isinstance(node, Neg):
        inner = fold_constants(node.operand)
        return Num(-inner.value) if isinstance(inner, Num) else Neg(inner)
    lhs, rhs = fold_constants(node.left), fold_constants(node.right)
    if isinstance(lhs, Num) and isinstance(rhs, Num) and not (
            node.op == "/" and rhs.value == 0):
        return Num(evaluate(Bin(node.op, lhs, rhs), {}))   # a compiler that
    return Bin(node.op, lhs, rhs)                          # contains a walker


def compile_expr(node, fold=False):
    """AST -> (code, consts, names). One instruction per AST node, always:
    `len(code) == count_nodes(node)` for every tree this can be handed."""
    if fold:
        node = fold_constants(node)
    code, consts, names = [], [], []
    seen = {}

    def intern(pool, key, value):
        """A constant pool and a name table -- CPython's co_consts/co_names.
        The key carries the type, so 1 and Fraction(1) stay two constants."""
        if key not in seen:
            seen[key] = len(pool)
            pool.append(value)
        return seen[key]

    def emit(node):
        if isinstance(node, Num):
            key = ("const", type(node.value), node.value)
            code.append((PUSH, intern(consts, key, node.value)))
        elif isinstance(node, Var):
            code.append((LOAD, intern(names, ("name", node.name), node.name)))
        elif isinstance(node, Neg):
            emit(node.operand)
            code.append((NEG, 0))
        else:
            emit(node.left)     # THE LINE: left before right, and nothing
            emit(node.right)    # downstream can tell if you swap them
            code.append((BINOP[node.op], 0))

    emit(node)
    return code, consts, names


def vm_run(code, consts, names, env):
    """The dispatch loop. No step counter inside: the code is straight-line,
    so instructions executed is exactly `len(code)` -- which demo.py checks
    with a counted copy rather than assuming."""
    stack = []
    push = stack.append
    pop = stack.pop
    pc, n = 0, len(code)
    while pc < n:
        op, arg = code[pc]
        pc += 1
        if op == PUSH:
            push(consts[arg])
        elif op == LOAD:
            push(env[names[arg]])
        elif op == ADD:
            b = pop()
            push(pop() + b)
        elif op == SUB:
            b = pop()
            push(pop() - b)
        elif op == MUL:
            b = pop()
            push(pop() * b)
        elif op == DIV:
            b = pop()
            push(pop() / b)
        else:
            push(-pop())         # NEG, the one opcode that pops just one
    return stack[-1]


# ----------------------------------------- back end 3: compile to closures

def closure_compile(node):
    """The same tree as nested closures: no opcodes, no stack, no dispatch
    loop. The tree's shape is baked into the closures' call graph -- and this
    beats every bytecode variant in the fairness grid (demo.py, section 3)."""
    if isinstance(node, Num):
        value = node.value
        return lambda env: value
    if isinstance(node, Var):
        name = node.name
        return lambda env: env[name]
    if isinstance(node, Neg):
        inner = closure_compile(node.operand)
        return lambda env: -inner(env)
    lhs, rhs = closure_compile(node.left), closure_compile(node.right)
    if node.op == "+":
        return lambda env: lhs(env) + rhs(env)
    if node.op == "-":
        return lambda env: lhs(env) - rhs(env)
    if node.op == "*":
        return lambda env: lhs(env) * rhs(env)
    return lambda env: lhs(env) / rhs(env)