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
            i.removeDef(i.result)
            i2 = ImmLoad(i.result, vr)
            i.Parent.replaceInstruction(i, i2)


class ConstantMerge(InstructionPass):
    def prepare(self):
        self.constMap = {}
    def onInstruction(self, i):
        if type(i) is ImmLoad:
            v = i.value
            if v in self.constMap:
                # v is already defined, re-use the imm-load from elsewhere
                pass
            else:
                self.constMap[v] = i
        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))


class SameImmLoadDeletePass(BasicBlockPass):
    def onBasicBlock(self, bb):
        constMap = {}
        imms = filter(lambda i: isinstance(i, ImmLoad), bb.Instructions)
        for ins in list(imms):
            if ins.value in constMap:
                # remove this immload and update all references to the target
                t_old = ins.target
                if not t_old.onlyUsedInBlock(bb):
                    continue
                # update all references:
                t_new = constMap[ins.value]
                print('Replace {} with {}'.format(t_old, t_new))
                for use in t_old.used_by:
                    use.replaceValue(t_old, t_new)
                bb.removeInstruction(ins)
            else:
                constMap[ins.value] = ins.target
            

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:
      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 and type(bb.LastInstruction) is Branch:
            # This block is empty.
            # find predecessors of this block and replace this block reference with the jumped reference.
            ins = bb.LastInstruction
            preds = bb.Predecessors
            if bb in preds:
                    # Do not remove if preceeded by itself
                    pass
            else:
                    for pred in bb.Predecessors:
                          pred.LastInstruction.changeTarget(bb, ins.target)
                    f.removeBasicBlock(bb)

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