#!/usr/bin/python
# -*- coding: iso-8859-2 -*-

###########################################################################
# Interpreter and IA-32 assembly compiler for a Simple Programming Language
#
# Copyright (C) 2006-2008 Gergõ ÉRDI
# http://cactus.rulez.org/
#
# Published under the terms of the GNU General Public License 2.0
# $Id: sp.py 188 2006-03-12 11:59:46Z cactus $
###########################################################################
#
# SP is a toy language allowing the following simple statements:
#
# READ var              Read an int from the user, and store it in 'var'
# WRITE expr            Print the value of 'expr'
# GOTO lbl [IF expr]    Jump to the specified line. If 'expr' is present,
#                       jump only if the result of 'expr' >= 0
# LET var = expr        Set value of 'var' to 'expr'
#
# There are 33 variables, named X, Y, Z, X0, ..., X9, Y0, ...
# The type of variables and expressions are all 32-bit signed integers
# 
###########################################################################

class Interpreter:
    def __init__ (self):
        self.vars = {}
        self.stmts = []
        self.lbl_map = {}
        
    def get_variable (self, var):
        if not var in self.vars: return 0
        return self.vars[var]

    def set_variable (self, var, val):
        self.vars[var] = val

    def load (self, program):
        for (lbl, stmt) in program:
            self.stmts.append(stmt)

            pos = len(self.stmts) - 1
            if not (lbl is None): self.lbl_map[lbl] = pos

    def run(self):
        self.pc = 0
        end = len(self.stmts)
        while self.pc < end:
            self.step()

    def step(self):
        stmt = self.stmts[self.pc]
        self.pc += 1
        stmt.evaluate (self)

    def jump(self, lbl):
        self.pc = self.lbl_map[lbl]

    def write(self, val):
        print val

    def read(self):
        try:
            s = raw_input('> ')
            return int(s)
        except EOFError:
            print
            return self.read()
        except ValueError:
            print "'%s': Invalid number" % s
            return self.read()

class Compiler:
    def __init__ (self):
        self.variables = {}

    def writeline (self, s = ""):
        print s
    
    def var_address (self, varname):
        if not varname in self.variables:
            self.variables[varname] = len (self.variables)
        return '_sp_v + %d' % (self.variables[varname] * 4)

    def label (self, lbl):
        return '_sp_l_%s' % lbl

    def compile (self, program):
        def create_var (varname):
            self.writeline ('%s\tresd 1' % self.var_address (varname))

        # String literals
        self.writeline ('\tsection .data')
        self.writeline ('_sp_fmt\t\tdb "%d", 0x0A, 0x00')
        self.writeline ('_sp_prompt\tdb "> ", 0x00')
        self.writeline ('_sp_infmt\tdb "%d", 0x00')
        self.writeline ()

        # Asm code of the program itself ('main' is exported, so you can link it with GCC)
        self.writeline ('\tsection .text')
        self.writeline ('\tglobal main')
        self.writeline ('\textern printf, scanf')
        self.writeline ('main:')
        
        for (lbl, stmt) in program:
            if not (lbl is None):
                self.writeline ('%s:' % self.label (lbl))
            for asmline in stmt.compile_to_asm (self):
                self.writeline ('\t%s' % asmline)
            self.writeline ()
        self.writeline ('\tret')

        # Implementation of _sp_input
        self.writeline ("""
_sp_input:
	push ebp
	mov ebp, esp
	push dword 0

_sp_input_loop:
	push _sp_prompt
	call printf
	sub esp, 4

	push ebp
	push _sp_infmt
	call scanf
	sub esp, 8

	mov edx, [ebp]
	cmp eax, 1 		; TODO: real input checking
	jne _sp_input_loop

	mov esp, ebp
	pop ebp
	ret
	""")

        # Space for SP variables
        self.writeline ('\tsection .bss')
        self.writeline ('_sp_v\tresd %d' % len (self.variables))


class Stmt:
    class Read:
        def __init__ (self, var):
            self.var = var

        def evaluate (self, interpreter):
            interpreter.set_variable(self.var, interpreter.read ())

        def compile_to_asm (self, compiler):
            return ['call _sp_input',
                    'mov [%s], edx' % compiler.var_address (self.var)]
            pass
        
    class Write:
        def __init__ (self, expr):
            self.expr = expr

        def evaluate (self, interpreter):
            interpreter.write(self.expr.evaluate(interpreter))

        def compile_to_asm (self, compiler):
            return self.expr.push_to_stack (compiler) + \
                   ['push dword _sp_fmt',
                    'call printf',
                    'sub esp, 8']

    class Assign:
        def __init__(self, var, expr):
            self.var = var
            self.expr = expr

        def evaluate(self, interpreter):
            interpreter.set_variable(self.var, self.expr.evaluate(interpreter))

        def compile_to_asm (self, compiler):
            return self.expr.push_to_stack (compiler) + \
                   ['pop edx',
                    'mov [%s], edx' % compiler.var_address (self.var)]

    class Goto:
        def __init__ (self, lbl, cond = None):
            self.lbl = lbl
            self.cond = cond

        def evaluate (self, interpreter):
            if not self.cond or self.cond.evaluate(interpreter) >= 0:
                interpreter.jump (self.lbl)

        def compile_to_asm (self, compiler):
            if self.cond:
                return self.cond.push_to_stack (compiler) + \
                       ['pop ecx',
                        'cmp ecx, 0',
                        'jge %s' % compiler.label (self.lbl)]
            else:
                return ['jmp %s' % compiler.label (self.lbl)]
           

class Expr:
    class NumLit:
        def __init__ (self, num):
            self.num = num

        def evaluate (self, interpreter):
            return self.num

        def push_to_stack (self, compiler):
            return ['push dword %s' % self.num]        

    class VariableVal:
        def __init__ (self, var):
            self.var = var

        def evaluate (self, interpreter):
            return interpreter.get_variable (self.var)

        def push_to_stack (self, compiler):
            return ['push dword [%s]' % compiler.var_address (self.var)]

    class BinaryFunction:
        def __init__ (self, fun, asm_opcode, left, right):
            self.fun = fun
            self.asm_opcode = asm_opcode
            self.left = left
            self.right = right

        def evaluate (self, interpreter):
            return self.fun (self.left.evaluate(interpreter),
                             self.right.evaluate(interpreter))

        def push_to_stack (self, compiler):
            return self.left.push_to_stack (compiler) + \
                   ['pop eax'] + \
                   self.right.push_to_stack (compiler) + \
                   ['pop ebx',
                    '%s eax, ebx' % self.asm_opcode,
                    'push eax']

    class Plus (BinaryFunction):
        def __init__ (self, left, right):
            Expr.BinaryFunction.__init__ (self, lambda x, y: x + y, 'add', left, right)

    class Minus (BinaryFunction):
        def __init__ (self, left, right):
            Expr.BinaryFunction.__init__ (self, lambda x, y: x - y, 'sub', left, right)

    class Mul (BinaryFunction):
        def __init__ (self, left, right):
            Expr.BinaryFunction.__init__ (self, lambda x, y: x * y, 'mul', left, right)

    class Div (BinaryFunction):
        def __init__ (self, left, right):
            Expr.BinaryFunction.__init__ (self, lambda x, y: x / y, 'div', left, right)


########################        
# Create lexer
########################        

def token(name, fun = lambda x: x):
    return lambda x: (name, fun(x))


lexer = [('\s+',          ''),
         ('[+\-*/:;()=]', lambda x: (x, x)),
         
         ('(?i)GOTO',     token('GOTO')),
         ('(?i)IF',       token('IF')),
         ('(?i)READ',     token('READ')),
         ('(?i)WRITE',    token('WRITE')),
         ('(?i)LET',      token('LET')),
         
         ('(?i)[a-z]',    token('LETTER', lambda x: x.lower())),
         ('(?i)[0-9]',    token('DIGIT',  lambda x: int(x)))]


########################        
# Create parser
########################        

grammar = [('program', ['progline'],                      lambda (l,), c:      [l]),
           ('program', ['program', 'progline'],           lambda (prog, l), c: prog + [l]),
           
           ('progline', ['number', ':', 'stmt', ';'],     lambda (lbl, col, stmt, scol), c: (lbl, stmt)),
           ('progline', ['stmt', ';'],                    lambda (stmt, scol), c:           (None, stmt)),
           
           ('stmt',   ['READ', 'varname'],                lambda (read, id), c:             Stmt.Read(id)),
           ('stmt',   ['WRITE', 'expr'],                  lambda (write, expr), c:          Stmt.Write(expr)),
           ('stmt',   ['LET', 'varname', '=', 'expr'],    lambda (let, id, eq, expr), c:    Stmt.Assign(id, expr)),
           ('stmt',   ['GOTO', 'number'],                 lambda (goto, lbl), c:            Stmt.Goto(lbl)),
           ('stmt',   ['GOTO', 'number', 'IF', 'expr'],   lambda (goto, lbl, if_, cond), c: Stmt.Goto(lbl, cond)),
           
           ('expr',   ['term'],                           lambda (t,),c:        t),
           ('expr',   ['term', '+', 'term'],              lambda (a, op, b), c: Expr.Plus (a, b)),
           ('expr',   ['term', '-', 'term'],              lambda (a, op, b), c: Expr.Minus (a, b)),
           
           ('term',   ['factor'],                         lambda (f,),c:        f),
           ('term',   ['factor', '*', 'factor'],          lambda (a, op, b), c: Expr.Mul (a, b)),
           ('term',   ['factor', '/', 'factor'],          lambda (a, op, b), c: Expr.Div (a, b)),
           
           ('factor', ['(', 'expr', ')'],                 lambda (lp, e, rp),c: e),
           ('factor', ['number'],                         lambda (num,),c:      Expr.NumLit(num)),
           ('factor', ['varname'],                        lambda (id,),c:       Expr.VariableVal(id)),
           
           ('varname', ['LETTER'],                        lambda (l,), c:        l),
           ('varname', ['LETTER', 'DIGIT'],               lambda (l, d), c:     '%s%d' % (l, d)),
           
           ('number', ['DIGIT'],                          lambda (d,),c:     d),
           ('number', ['number', 'DIGIT'],                lambda (num, d),c: num * 10 + d)]

           
import yappy.parser
parser = yappy.parser.Yappy (lexer, grammar)

########################        
# Compile programs
########################
def load_sp (filename):
    f = open (filename, 'r')
    return parser.input (f.read ())

def run_sp (filename):
    prog = load_sp (filename)
    interpreter = Interpreter ()
    interpreter.load(prog)
    interpreter.run()

def compile_sp (filename):
    prog = load_sp (filename)
    compiler = Compiler ()
    compiler.compile (prog)

if __name__ == '__main__':
    import sys

    for filename in sys.argv[1:]:
        #run_sp (filename)
        compile_sp (filename)