implement the code generator

2 files changed, 182 insertions, 3 deletions

diff --git a/src/back/generator.py b/src/back/generator.py
index 851149b..b743684 100644
--- a/src/back/generator.py
+++ b/src/back/generator.py
@@ -1,5 +1,10 @@
+from back.tac import *
+from front.scope import Scope
+
 class Register(object):
     __shared_state = {}
+    __current_function = None
+    __function_registers = {}
     count = 0
 
     def __init__(self):
@@ -7,4 +12,173 @@ class Register(object):
 
     def new(self):
         self.count += 1
+        self.__function_registers[self.__current_function].append(self.count)
         return self.count
+
+    def set_function(self, name):
+        self.__current_function = name
+        self.__function_registers[self.__current_function] = []
+
+    def get_registers(self, name):
+        return self.__function_registers[name]
+
+class Generator(object):
+    def __init__(self):
+        self.tac_list = TACList()
+        self.__op_list = []
+
+    def emit(self, op, tac = None):
+        if tac:
+            self.__op_list[self.__op_list.index(tac)] = op
+        else:
+            self.__op_list.append(op)
+
+    def generate_prologue(self, name):
+        Scope().set_function(name)
+        self.emit(Label(name))
+        self.emit("PUSH bp")
+        self.emit("ADD bp, r0, sp")
+        for variable in Scope().get_variables()[1]:
+            self.emit("PUSH r0")
+        for register in Register().get_registers(name):
+            self.emit("PUSH r%d" % register)
+
+    def generate_epilogue(self, name):
+        for variable in Scope().get_variables()[1]:
+            self.emit("POP r0")
+        Register().get_registers(name).reverse()
+        for register in Register().get_registers(name):
+            self.emit("POP r%d" % register)
+        Register().get_registers(name).reverse()
+        self.emit("ADD sp, r0, bp")
+        self.emit("POP bp")
+        if name == "main":
+            self.emit("PUSH rv")
+            self.emit("PUSH r0")
+            self.emit("SYS")
+        else:
+            self.emit("RET")
+
+    def generate(self):
+        # pass 1 - generate non-label ops
+        for tac in self.tac_list:
+            if isinstance(tac, TAC):
+                if tac.op == Op.ADD:
+                    self.emit("ADD r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.SUB:
+                    self.emit("SUB r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.MUL:
+                    self.emit("MUL r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.DIV:
+                    self.emit("DIV r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.MOD:
+                    self.emit("MOD r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.AND:
+                    self.emit("AND r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.OR:
+                    self.emit("OR  r%d, r%d, r%d" % (tac.arg1, tac.arg1, tac.arg2))
+                elif tac.op == Op.NOT:
+                    r = Register().new()
+                    self.emit("MOV r%d, 0" % r)
+                    self.emit("CMP r%d, r%d" % (tac.arg1, r))
+                    self.emit("EQ  r%d" % (tac.arg1))
+                elif tac.op == Op.MINUS:
+                    self.emit("SUB r%d, r0, r%d" % (tac.arg1, tac.arg1))
+                elif tac.op == Op.STORE:
+                    offset = Scope().get_variable_offset(tac.arg1)
+                    self.emit("SW bp, r%d, %d" % (tac.arg2, offset * 4))
+                elif tac.op == Op.LOAD:
+                    offset = Scope().get_variable_offset(tac.arg1)
+                    self.emit("LW bp, r%d, %d" % (tac.arg2, offset * 4))
+                elif tac.op == Op.MOV:
+                    self.emit("MOV r%d, %d" % (tac.arg1, tac.arg2))
+                elif tac.op == Op.CMP:
+                    self.emit("CMP r%d, r%d" % (tac.arg1, tac.arg2))
+                elif tac.op == Op.EQ:
+                    self.emit("EQ  r%d" % tac.arg1)
+                elif tac.op == Op.NE:
+                    self.emit("NE  r%d" % tac.arg1)
+                elif tac.op == Op.LT:
+                    self.emit("LT  r%d" % tac.arg1)
+                elif tac.op == Op.LE:
+                    self.emit("LE  r%d" % tac.arg1)
+                elif tac.op == Op.GE:
+                    self.emit("GE  r%d" % tac.arg1)
+                elif tac.op == Op.GT:
+                    self.emit("GT  r%d" % tac.arg1)
+                elif tac.op == Op.BEZ:
+                    self.emit(tac)
+                elif tac.op == Op.JMP:
+                    self.emit(tac)
+                elif tac.op == Op.PUSH:
+                    self.emit("PUSH r%d" % tac.arg1)
+                elif tac.op == Op.POP:
+                    self.emit("POP r%d" % tac.arg1)
+                elif tac.op == Op.CALL:
+                    self.emit(tac)
+                    self.emit("ADD r%d, r0, rv" % tac.arg2)
+                elif tac.op == Op.RETURN:
+                    self.emit("ADD rv, r0, r%d" % tac.arg1)
+                    self.emit(tac)
+            elif isinstance(tac, FunctionPrologue):
+                self.generate_prologue(tac.name)
+            elif isinstance(tac, FunctionEpilogue):
+                self.generate_epilogue(tac.name)
+            elif isinstance(tac, Label):
+                self.emit(Label(tac.name))
+            else:
+                raise Exception("%s is not a valid TACList element", repr(tac))
+
+        # pass 2 - generate label ops
+        for op in self.__op_list:
+            if isinstance(op, TAC):
+                if op.op == Op.BEZ:
+                    offset = self.get_label_offset(op.arg2, op)
+                    self.emit("BEZ r%d, %d" % (op.arg1, offset), op)
+                elif op.op == Op.JMP:
+                    offset = self.get_label_offset(op.arg1, op)
+                    self.emit("JMP %d" % offset, op)
+                elif op.op == Op.CALL:
+                    offset = self.get_label_offset(op.arg1, op)
+                    self.emit("CALL %d" % offset, op)
+                elif op.op == Op.RETURN:
+                    offset = self.get_label_offset("L%d" % op.arg2, op)
+                    self.emit("JMP %d" % offset, op)
+
+        # pass 3 - remove labels
+        self.__op_list = filter(lambda x: not isinstance(x, Label), self.__op_list)
+
+        # pass 4 - replace pseudo registers and insert number of registers
+        def replace_pseudo_regs(op):
+            regs = Register().count
+            bp = "r%d" % (regs + 1)
+            sp = "r%d" % (regs + 2)
+            rv = "r%d" % (regs + 3)
+            return op.replace("bp", bp).replace("sp", sp).replace("rv", rv)
+
+        self.__op_list = map(replace_pseudo_regs, self.__op_list)
+        self.__op_list = [".REGS %d" % Register().count] + self.__op_list
+
+        # print code
+        #for i in range(len(self.__op_list)):
+        #    print "%04d: %s" % (i, str(self.__op_list[i]))
+        print "\n".join(self.__op_list)
+
+    def get_label_offset(self, name, tac):
+        tac_index = 0
+        for op in self.__op_list:
+            if isinstance(op, Label):
+                continue
+            if op == tac:
+                break
+            tac_index += 1
+
+        name_index = 0
+        for op in self.__op_list:
+            if isinstance(op, Label):
+                if op.name == name:
+                    return name_index - tac_index
+                continue
+            name_index += 1
+
+        raise Exception("label %s does not exist" % name)
diff --git a/src/back/tac.py b/src/back/tac.py
index bff5944..20b14d9 100644
--- a/src/back/tac.py
+++ b/src/back/tac.py
@@ -37,7 +37,8 @@ Op.GT       = Op("GT")     # x = !Z && !N
 Op.BEZ      = Op("BEZ")    # if x == 0 goto y
 Op.JMP      = Op("JMP")    # goto x
 
-Op.PARAM    = Op("PARAM")  # param x
+Op.PUSH     = Op("PUSH")   # push x
+Op.POP      = Op("POP")    # pop x
 Op.CALL     = Op("CALL")   # call x return in y
 Op.RETURN   = Op("RETURN") # return x
 
@@ -59,9 +60,13 @@ class Label(object):
     def __repr__(self):
         return "<Label, name: %s>" % self.name
 
-class FunctionLabel(Label):
+class FunctionPrologue(Label):
     def __repr__(self):
-        return "<FunctionLabel, name: %s>" % self.name
+        return "<FunctionPrologue, name: %s>" % self.name
+
+class FunctionEpilogue(Label):
+    def __repr__(self):
+        return "<FunctionEpilogue, name: %s>" % self.name
 
 class TACList(object):
     __shared_state = {}