# 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)
    getFunction = findFunction

    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))
        for f in self.Functions:
            f.check()