# 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.funcs = []
      self.variables = []

   def __repr__(self):
      return 'IR-module [{0}]'.format(self.name)

   def getInstructions(self):
      ins = []
      for bb in self.BasicBlocks:
         ins += bb.Instructions
      return ins
   Instructions = property(getInstructions)

   def getBBs(self):
      bbs = []
      for f in self.Functions:
         bbs += f.BasicBlocks
      return bbs

   BasicBlocks = property(getBBs)
   def addFunc(self, f):
      self.funcs.append(f)
   addFunction = addFunc

   def addVariable(self, v):
        self.variables.append(v)

   def getVariables(self):
        return self.variables
   Variables = property(getVariables)

   def getFunctions(self):
      return self.funcs
   Functions = property(getFunctions)

   def findFunction(self, name):
        for f in self.funcs:
            if f.name == name:
                return f
        raise KeyError(name)

   def dump(self):
      print(self)
      for v in self.Variables:
            print('   ', v)
      for fn in self.Functions:
         print(fn)
         for bb in fn.BasicBlocks:
            print('   ', bb)
            for ins in bb.Instructions:
               print('      ', ins)

   def dumpgv(self, outf):
      outf.write('digraph G \n{\n')
      for f in self.Functions:
         outf.write('{0} [label="{1}" shape=box3d]\n'.format(id(f), f))
         for bb in f.BasicBlocks:
            contents = str(bb) + '\n'
            contents += '\n'.join([str(i) for i in bb.Instructions])
            outf.write('{0} [shape=note label="{1}"];\n'.format(id(bb), contents))
            for successor in bb.Successors:
               outf.write('"{0}" -> "{1}"\n'.format(id(bb), id(successor)))
            # Draw dags if any:
            if hasattr(bb, 'dag'):
               outf.write('{0} -> {1}\n'.format(id(bb), id(bb.dag)))
               bb.dag.dumpgv(outf)

         outf.write('"{0}" -> "{1}" [label="entry"]\n'.format(id(f), id(f.entry)))
      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 Use, "use must be Value, not {0}".format(type(t))