import ppci
import ir
import registerallocator

# Instruction selection with DAG (Directed Acyclic Graph)
class DagLeaf:
   def __init__(self, v):
      self.v = v

class DagNode:
   def __init__(self, name):
      self.name = name
      self.children = []
   def __repr__(self):
      return str(self.name)

class Dag:
   def __init__(self, bb):
      self.mapping = {}
      self.buildFromBB(bb)
   def buildFromBB(self, bb):
      for ins in bb.Instructions:
         if type(ins) is ir.BinaryOperator:
            if not ins.value1 in self.mapping:
               self.mapping[ins.value1] = DagNode(ins.value1)
            if not ins.value2 in self.mapping:
               self.mapping[ins.value2] = DagNode(ins.value2)
            # look for op with left and right operand the same:
            N = None
            lnode = self.mapping[ins.value1]
            rnode = self.mapping[ins.value2]
            for node in self.mapping.values():
               if node.name == ins.operation:
                  if node.children[0] == lnode and node.children[1] == rnode:
                     N = node
                     break
            if not N:
               # Create a node.
               N = DagNode(ins.operation)
               N.children.append(lnode)
               N.children.append(rnode)
            self.mapping[ins.result] = N
         else:
            pass
   def dumpgv(self, outf):
      outf.write('subgraph {0} {{\n'.format(id(self)))
      for node in self.mapping.values():
         outf.write('{0} [label="{1}"];\n'.format(id(node), node.name))
         for c in node.children:
            outf.write('{0} -> {1};\n'.format(id(node), id(c)))
      outf.write('label="dag"}\n')

def insSelect(mod):
   """ Create DAG from ir-code """
   for bb in mod.BasicBlocks:
      print(bb)
      dag = Dag(bb)
      print(dag.mapping)
      bb.dag = dag

# Machine code interface:
class MachineOperand:
   """ Single machine operand """
   pass

class MachineInstruction:
   def __init__(self, opcode):
      self.opcode = opcode
      self.operands = []


# x86 specific:
class AsmLabel:
   def __init__(self, lab):
      self.lab = lab
   def __repr__(self):
      return '{0}:'.format(self.lab)

class Op:
   def __init__(self, op, dst, src):
      self.op = op
      self.src = src
      self.dst = dst
   def __repr__(self):
      return '{0} {1}, {2}'.format(self.op, self.dst, self.src)

class Jmp:
   def __init__(self, j, target):
      self.j = j
      self.target = target
   def __repr__(self):
      return '{0} {1}'.format(self.j, self.target)

class X86CodeGen:
   def __init__(self, diag):
      self.diag = diag
      self.regs = ['rax', 'rbx', 'rcx', 'rdx']

   def emit(self, i):
      self.asm.append(i)

   def genBin(self, ir):
      self.asm = []
      # Allocate registers:
      insSelect(ir)
      ra = registerallocator.RegisterAllocator()
      # TODO: do not register allocate on intermediate code:
      ra.registerAllocate(ir, self.regs)
      self.genModule(ir)
      return self.asm

   def genModule(self, ir):
      for f in ir.Functions:
         self.genFunction(f)
   def genFunction(self, f):
      self.emit('global {0}'.format(f.name))
      self.emit(AsmLabel(f.name))
      self.emit(Jmp('jmp', f.entry.name))
      for bb in f.BasicBlocks:
         self.genBB(bb)
   def genBB(self, bb):
      self.emit(AsmLabel(bb.name))
      for i in bb.Instructions:
         self.genIns(i)
   def genIns(self, i):
      if type(i) is ir.BinaryOperator:
         ops = {'+':'add', '-':'sub', '*':'mul'}
         if i.operation in ops:
            self.emit(Op('mov', i.result.reg, i.value1.reg))
            self.emit(Op(ops[i.operation], i.result.reg, i.value2.reg))
         else:
            raise NotImplementedError('op {0}'.format(i.operation))
      elif type(i) is ir.Load:
         self.emit(Op('mov', i.value, '[{0}]'.format(i.location)))
      elif type(i) is ir.Return:
         self.emit('ret')
      elif type(i) is ir.Call:
         self.emit('call')
      elif type(i) is ir.ImmLoad:
         self.emit(Op('mov', i.target, i.value))
      elif type(i) is ir.Store:
         self.emit(Op('mov', '[{0}]'.format(i.location), i.value))
      elif type(i) is ir.ConditionalBranch:
         self.emit(Op('cmp', i.a, i.b))
         jmps = {'>':'jg', '<':'jl', '==':'je'}
         if i.cond in jmps:
            j = jmps[i.cond]
            self.emit(Jmp(j, i.lab1.name))
         else:
            raise NotImplementedError('condition {0}'.format(i.cond))
         self.emit(Jmp('jmp', i.lab2.name))
      elif type(i) is ir.Branch:
         self.emit(Jmp('jmp', i.target.name))
      elif type(i) is ir.Alloc:
         pass
      else:
         raise NotImplementedError('{0}'.format(i))