# IR-Structures:
from .instruction import *
from .basicblock import BasicBlock

class Module:
   """ Main container for a piece of code. """
   def __init__(self, name):
      self.name = name
      self.bbs = []
   def __repr__(self):
      return 'IR-module [{0}]'.format(self.name)
   def getInstructions(self):
      ins = []
      for bb in self.bbs:
         ins += bb.Instructions
      return ins
   Instructions = property(getInstructions)
   def addBB(self, bb):
      self.bbs.append(bb)
   def getBBs(self):
      return self.bbs
   BasicBlocks = property(getBBs)
   def dump(self):
      print(self)
      for i in self.Instructions:
         print(i, 'live vars:', list(i.live_in), 'uses', list(i.uses), 'defs', list(i.defs))
      print('END')
   def dumpgv(self, outf):
      outf.write('digraph G \n{\n')
      for i in self.Instructions:
         outf.write('{0} [label="{1}"];\n'.format(id(i), i))
         for succ in i.succ:
            outf.write('"{0}" -> "{1}" [label="{2}"];\n'.format(id(i), id(succ), succ.live_in))
      outf.write('}\n')

   # Analysis functions:
   def check(self):
      """ Perform sanity check on module """
      for i in self.Instructions:
         for t in i.defs:
            assert type(t) is Value, "def must be Value, not {0}".format(type(t))
         for t in i.uses:
            assert type(t) is Value, "use must be Value, not {0}".format(type(t))
   def analyze(self):
      # Determine pred and succ:
      def link(a, b):
         a.succ.add(b)
         b.pred.add(a)
      for bb in self.bbs:
         if not bb.Empty:
            if len(bb.Instructions) > 1:
               for i1, i2 in zip(bb.Instructions[:-1], bb.Instructions[1:]):
                  link(i1, i2)
            else:
               print('Only 1 long!', bb, bb.Instructions)
            if type(bb.LastIns) is ConditionalBranch:
               link(bb.LastIns, bb.LastIns.lab1.FirstIns)
               link(bb.LastIns, bb.LastIns.lab2.FirstIns)
            if type(bb.LastIns) is Branch:
               link(bb.LastIns, bb.LastIns.target.FirstIns)
         else:
            print('Empty!', bb)
      # We now have cfg

      # Determine liveness:
      for i in self.Instructions:
         i.live_in = set()
         i.live_out = set()
      for z in range(50):
         # TODO iterate until converge
         for i in self.Instructions:
            lo_mk = i.live_out.difference(i.defs)
            i.live_in = i.uses.union(lo_mk)
            lo = set()
            for s in i.succ:
               lo = lo.union(s.live_in)
            i.live_out = lo
   def constantProp(self):
      """ Constant propagation. Pre-calculate constant values """
      for i in self.Instructions:
         if type(i) is ImmLoad:
            i.target.constval = i.value
         elif type(i) is BinaryOperator:
            a = i.value1
            b = i.value2
            if i.value1.constval and i.value2.constval:
               op = i.operation
               if op == '+':
                  i.result.constval = a + b
               else:
                  raise NotImplementedError(op)
            else:
               i.result.constval = None

   def registerAllocate(self, regs):
      print(regs)
      allVals = []
      # construct interference:
      for i in self.Instructions:
         for v in i.live_in:
            allVals.append(v)
            for v2 in i.live_in:
               if v != v2:
                  v.interferes.add(v2)
      # assign random registers:
      print(allVals)
      regs = set(regs)
      for v in allVals:
         takenregs = set([iv.reg for iv in v.interferes])
         r2 = list(regs.difference(takenregs))
         # Pick next available:
         v.reg = r2[0]