#!/usr/bin/python3
"""Simple stack-based virtual machine.

What I have in mind is to compile this stack code to RTL quads.

But I haven’t done that yet.

"""
from collections import namedtuple


### Stack code ###

class Const(namedtuple('Const', ('val',))):
    def run(self, machine):
        machine.push(self.val)

class Store(namedtuple('Store', ('index',))):
    def run(self, machine):
        machine.store(self.index, machine.pop())

class Label(namedtuple('Label', ('id',))):
    def run(self, machine):
        pass

class Load(namedtuple('Load', ('index',))):
    def run(self, machine):
        machine.push(machine.load(self.index))

class BranchLessEqual(namedtuple('BranchLessEqual', ('dest',))):
    def run(self, machine):
        a = machine.pop()
        if machine.pop() <= a:
            machine.jump(self.dest)

class Add(namedtuple('Add', ())):
    def run(self, machine):
        machine.push(machine.pop() + machine.pop())

class Subtract(namedtuple('Subtract', ())):
    def run(self, machine):
        a = machine.pop()
        machine.push(machine.pop() - a)

class Jump(namedtuple('Jump', ('dest',))):
    def run(self, machine):
        machine.jump(self.dest)

class Return(namedtuple('Return', ())):
    def run(self, machine):
        machine.do_return()

# Program to compute a triangular number.  Seems to work.
program = [ Const(0), Store(1),                                # t := 0
            Label(0), Load(0), Const(0), BranchLessEqual(1),   # while n > 0 {
            Load(0), Load(1), Add(), Store(1),                 # t <- t + n
            Load(0), Const(1), Subtract(), Store(0),           # n <- n - 1
            Jump(0),                                           # }
            Label(1), Load(1), Return(),                       # return t
          ]

class Machine:
    def execute(self, code, args, trace=lambda *args: None):
        self.stack = []
        self.locals = dict(enumerate(args))
        self.pc = 0
        self.running = True

        self.labels = {insn.id: i
                       for i, insn in enumerate(code)
                       if isinstance(insn, Label)}

        while self.running:
            insn = code[self.pc]
            trace(self.pc, insn, self.stack)
            self.pc += 1
            insn.run(self)
        rv = self.stack.pop()
        del self.stack, self.running
        return rv

    def push(self, val):
        self.stack.append(val)

    def pop(self):
        return self.stack.pop()

    def load(self, index):
        return self.locals[index]

    def store(self, index, val):
        self.locals[index] = val

    def jump(self, dest):
        self.pc = self.labels[dest]

    def do_return(self):
        self.running = False


def generate_ssa_rtl(code, n_args):
    "Not working yet; doesn't insert phony functions yet."
    labels = {insn.id: i for i, insn in enumerate(code)
              if isinstance(insn, Label)}
    rtl = []
    next_reg = n_args  # We start with the args
    locals = {i: 'v%d' % i for i in range(n_args)}
    stack = []
    for insn in code:
        if isinstance(insn, Label):
            rtl.append(insn.id)
        elif isinstance(insn, Const):
            dest = 'v%d' % next_reg
            next_reg += 1
            rtl.append(('lc', dest, insn.val))
            stack.append(dest)
        elif isinstance(insn, Load):
            stack.append('v%d' % insn.index)
        elif isinstance(insn, Store):
            locals[insn.index] = stack.pop()
        elif isinstance(insn, BranchLessEqual):
            a = stack.pop()
            rtl.append(('ble', stack.pop(), a, insn.dest))
        elif isinstance(insn, Jump):
            rtl.append(('jmp', insn.dest))
        elif isinstance(insn, Add):
            dest = 'v%d' % next_reg
            next_reg += 1
            rtl.append(('add', dest, stack.pop(), stack.pop()))
            stack.append(dest)
        elif isinstance(insn, Subtract):
            dest = 'v%d' % next_reg
            next_reg += 1
            a = stack.pop()
            rtl.append(('sub', dest, stack.pop(), a))
            stack.append(dest)
        elif isinstance(insn, Return):
            rtl.append(('ret', stack.pop()))
        else:
            pass
            raise ValueError(insn, rtl)

    return rtl