from ir import * 
# Standard passes:

class FunctionPass:
   def run(self, ir):
      """ Main entry point for the pass """
      self.prepare()
      for f in ir.Functions:
         self.onFunction(f)
   def onFunction(self, f):
      """ Override this virtual method """
      raise NotImplementedError()
   def prepare(self):
      pass

class BasicBlockPass(FunctionPass):
   def onFunction(self, f):
      for bb in f.BasicBlocks:
         self.onBasicBlock(bb)
   def onBasicBlock(self, bb):
      """ Override this virtual method """
      raise NotImplementedError()
      
class InstructionPass(BasicBlockPass):
   def onBasicBlock(self, bb):
      for ins in iter(bb.Instructions):
         self.onInstruction(ins)
   def onInstruction(self, ins):
      """ Override this virtual method """
      raise NotImplementedError()

# Usefull transforms:
class ConstantFolder(InstructionPass):
   def prepare(self):
      self.constMap = {}
   def onInstruction(self, i):
      if type(i) is ImmLoad:
         self.constMap[i.target] = i.value
      elif type(i) is BinaryOperator:
         if i.value1 in self.constMap and i.value2 in self.constMap:
            op = i.operation
            va = self.constMap[i.value1]
            vb = self.constMap[i.value2]
            if op == '+':
               vr = va + vb
            elif op == '*':
               vr = va * vb
            elif op == '-':
               vr = va - vb
            else:
               vr = None
               return
            self.constMap[i.result] = vr
            i2 = ImmLoad(i.result, vr)
            i.Parent.replaceInstruction(i, i2)

class DeadCodeDeleter(BasicBlockPass):
   def onBasicBlock(self, bb):
      def instructionUsed(ins):
         if len(ins.defs) == 0:
            # In case this instruction does not define any variables, assume it is usefull.
            return True
         for d in ins.defs:
            if d.IsUsed:
               return True
         return False
      bb.Instructions = list(filter(instructionUsed, bb.Instructions))

def isAllocPromotable(allocinst):
   # Check if alloc value is only used by load and store operations.
   assert type(allocinst) is Alloc
   for use in ai.value.used_by:
      print(use.user, use)
      if not type(use.user) in [Load, Store]:
         # TODO: check volatile
         return False
         otherUse = True
   return True

class CleanPass(FunctionPass):
    def onFunction(self, f):
      bbs = list(f.BasicBlocks)
      for bb in bbs:
         # TODO: determine check for 'empty'

         # If a block only contains a branch, it can be removed:
         if len(bb.Instructions) == 1:
            # This block is empty.
            # find predecessors of this block and replace this block reference with the jumped reference.
            ins = bb.LastInstruction
            if type(ins) is Branch:
                print('Removing block {}'.format(bb))
                #print(ins, bb.Predecessors)
                preds = bb.Predecessors
                if bb in preds:
                    # Do not remove if preceeded by itself
                    pass
                else:
                    for pred in bb.Predecessors:
                          print('predecessor: {}'.format(pred))
                          pred.LastInstruction.changeTarget(bb, ins.target)
                    f.removeBasicBlock(bb)

class Mem2RegPromotor(FunctionPass):
   def onFunction(self, f):
      # TODO
      print(f)