#!/usr/bin/python

import unittest
import os

from compiler.compiler import Compiler
from compiler.errors import CompilerException, printError
from compiler import lexer
from compiler.parser import Parser
from compiler import assembler
from compiler.codegenerator import CodeGenerator
from project import Project

class CompilerTestCase(unittest.TestCase):
   """ test methods start with 'test*' """
   def testSource1(self):
      source = """
      module lcfos;
      var  
        a : integer;

      procedure putchar(num : integer);
      begin
      end putchar;

      procedure WriteNum( num: integer);
        var 
          d, base :  integer;
          dgt : integer;
        begin
          d := 1;
          base := 10;
          while num div d >= base do
            d := d * base
          end;
          while d <> 0 do
             dgt := num div d;
             num := num mod d;
             d   := d div base;
             putchar(48 + dgt)
           end
        end WriteNum;

      begin
        a := 1;
        while a < 26
         do
           putchar(65+a);
           a := a * 2
         end;
      end lcfos.
      """
      pc = Compiler()
      pc.compilesource(source)
   def testSource2(self):
      source = """
      module lcfos;
      var  
        a, b : integer;
        arr: array 30 of integer;
        arr2: array 10, 12 of integer;
        procedure t2*() : integer;
        begin
        a := 2;
        while a < 5 do
           b := arr[a-1] + arr[a-2];
           arr2[a,2] := b;
           arr2[a,3] := arr2[a,2] + arr2[a,2]*3 + b;
           arr[a] := b;
           a := a  + 1;
         end;
         return b
        end t2;
        begin
         b := 12;
         arr[0] := 1;
         arr[1] := 1;
      end lcfos.
      """
      pc = Compiler()
      mod = pc.compilesource(source)
   def testSource5(self):
      source = """
      module lcfos;
      procedure WriteLn() : integer;
        const zzz = 13;
        var
          a, b, c: integer;
        begin
         a := 2;
         b := 7;
         c := 10 * a + b*10*a;
         return c
        end WriteLn;
      begin  end lcfos.
      """
      pc = Compiler()
      pc.compilesource(source)
   def tstForStatement(self):
      source = """
      module fortest;
      var  
        a,b,c : integer;
      begin
         c := 0;
         for a := 1 to 10 by 1 do
            b := a + 15;
            c := c + b * a;
         end;
      end fortest.
      """
      pc = Compiler()
      pc.compilesource(source)
   def testSourceIfAndWhilePattern(self):
      source = """
      module lcfos;
      procedure WriteLn() : integer;
        const zzz = 13;
        var
          a, b, c: integer;
        begin
         a := 1;
         b := 2;
         if a * 3 > b then
            c := 10*a + b*10*a*a*a*b;
         else
            c := 13;
         end;
         while a < 101 do
            a := a + 1;
            c := c + 2;
         end;
         return c
        end WriteLn;
      begin end lcfos.
      """
      pc = Compiler()
      pc.compilesource(source)

   def testPattern1(self):
      """ Test if expression can be compiled into byte code """
      src = "12*13+33-12*2*3"
      tokens = lexer.tokenize(src)
      ast = Parser(tokens).parseExpression()
      code = CodeGenerator().genexprcode(ast)

   def testAssembler(self):
      """ Check all kind of assembler cases """
      assert(assembler.shortjump(5) == [0xeb, 0x5])
      assert(assembler.shortjump(-2) == [0xeb, 0xfc])
      assert(assembler.shortjump(10,'GE') == [0x7d, 0xa])
      assert(assembler.nearjump(5) == [0xe9, 0x5,0x0,0x0,0x0])
      assert(assembler.nearjump(-2) == [0xe9, 0xf9, 0xff,0xff,0xff])
      assert(assembler.nearjump(10,'LE') == [0x0f, 0x8e, 0xa,0x0,0x0,0x0])

   def testCall(self):
      assert(assembler.call('r10') == [0x41, 0xff, 0xd2])
      assert(assembler.call('rcx') == [0xff, 0xd1])
   def testXOR(self):
      assert(assembler.xorreg64('rax', 'rax') == [0x48, 0x31, 0xc0])
      assert(assembler.xorreg64('r9', 'r8') == [0x4d, 0x31, 0xc1])
      assert(assembler.xorreg64('rbx', 'r11') == [0x4c, 0x31, 0xdb])

   def testINC(self):
      assert(assembler.increg64('r11') == [0x49, 0xff, 0xc3])
      assert(assembler.increg64('rcx') == [0x48, 0xff, 0xc1])

   def testPush(self):
      assert(assembler.push('rbp') == [0x55])
      assert(assembler.push('rbx') == [0x53])
      assert(assembler.push('r12') == [0x41, 0x54])
   def testPop(self):
      assert(assembler.pop('rbx') == [0x5b])
      assert(assembler.pop('rbp') == [0x5d])
      assert(assembler.pop('r12') == [0x41, 0x5c])

   def testAsmLoads(self):
      # TODO constant add testcases
      assert(assembler.mov('rbx', 'r14') == [0x4c, 0x89, 0xf3])
      assert(assembler.mov('r12', 'r8')  == [0x4d, 0x89, 0xc4])
      assert(assembler.mov('rdi', 'rsp') == [0x48, 0x89, 0xe7])

   def testAsmMemLoads(self):
      assert(assembler.mov('rax', ['r8','r15',0x11]) == [0x4b,0x8b,0x44,0x38,0x11])
      assert(assembler.mov('r13', ['rbp','rcx',0x23]) == [0x4c,0x8b,0x6c,0xd,0x23])

      assert(assembler.mov('r9', ['rbp',-0x33]) == [0x4c,0x8b,0x4d,0xcd])
      #assert(assembler.movreg64('rbx', ['rax']) == [0x48, 0x8b,0x18])

      assert(assembler.mov('rax', [0xb000]) == [0x48,0x8b,0x4,0x25,0x0,0xb0,0x0,0x0])
      assert(assembler.mov('r11', [0xa0]) == [0x4c,0x8b,0x1c,0x25,0xa0,0x0,0x0,0x0])

      assert(assembler.mov('r11', ['RIP', 0xf]) == [0x4c,0x8b,0x1d,0x0f,0x0,0x0,0x0])

   def testAsmMemStores(self):
      assert(assembler.mov(['rbp', 0x13],'rbx') == [0x48,0x89,0x5d,0x13])
      assert(assembler.mov(['r12', 0x12],'r9') == [0x4d,0x89,0x4c,0x24,0x12])
      assert(assembler.mov(['rcx', 0x11],'r14') == [0x4c,0x89,0x71,0x11])


      assert(assembler.mov([0xab], 'rbx') == [0x48,0x89,0x1c,0x25,0xab,0x0,0x0,0x0])
      assert(assembler.mov([0xcd], 'r13') == [0x4c,0x89,0x2c,0x25,0xcd,0x0,0x0,0x0])

      assert(assembler.mov(['RIP', 0xf], 'r9') == [0x4c,0x89,0x0d,0x0f,0x0,0x0,0x0])

   def testAsmMOV8(self):
      assert(assembler.mov(['rbp', -8], 'al') == [0x88, 0x45, 0xf8])
      assert(assembler.mov(['r11', 9], 'cl') == [0x41, 0x88, 0x4b, 0x09])

      assert(assembler.mov(['rbx'], 'al') == [0x88, 0x03])
      assert(assembler.mov(['r11'], 'dl') == [0x41, 0x88, 0x13])

   def testAsmLea(self):
      assert(assembler.leareg64('r11', ['RIP', 0xf]) == [0x4c,0x8d,0x1d,0x0f,0x0,0x0,0x0])
      assert(assembler.leareg64('rsi', ['RIP', 0x7]) == [0x48,0x8d,0x35,0x07,0x0,0x0,0x0])

      assert(assembler.leareg64('rcx', ['rbp', -8]) == [0x48,0x8d,0x4d,0xf8])

   def testAssemblerCMP(self):
      assert(assembler.cmpreg64('rdi', 'r13') == [0x4c, 0x39, 0xef])
      assert(assembler.cmpreg64('rbx', 'r14') == [0x4c, 0x39, 0xf3])
      assert(assembler.cmpreg64('r12', 'r9')  == [0x4d, 0x39, 0xcc])

      assert(assembler.cmpreg64('rdi', 1)  == [0x48, 0x83, 0xff, 0x01])
      assert(assembler.cmpreg64('r11', 2)  == [0x49, 0x83, 0xfb, 0x02])
   def testAssemblerADD(self):
      assert(assembler.addreg64('rbx', 'r13') == [0x4c, 0x01, 0xeb])
      assert(assembler.addreg64('rax', 'rbx') == [0x48, 0x01, 0xd8])
      assert(assembler.addreg64('r12', 'r13') == [0x4d, 0x01, 0xec])

      assert(assembler.addreg64('rbx', 0x13) == [0x48, 0x83, 0xc3, 0x13])
      assert(assembler.addreg64('r11', 0x1234567) == [0x49, 0x81, 0xc3, 0x67, 0x45,0x23,0x1])
      assert(assembler.addreg64('rsp', 0x33) == [0x48, 0x83, 0xc4, 0x33])

   def testAssemblerSUB(self):
      assert(assembler.subreg64('rdx', 'r14') == [0x4c, 0x29, 0xf2])
      assert(assembler.subreg64('r15', 'rbx') == [0x49, 0x29, 0xdf])
      assert(assembler.subreg64('r8', 'r9') == [0x4d, 0x29, 0xc8])

      assert(assembler.subreg64('rsp', 0x123456) == [0x48, 0x81, 0xec, 0x56,0x34,0x12,0x0])
      assert(assembler.subreg64('rsp', 0x12) == [0x48, 0x83, 0xec, 0x12])

   def testAssemblerIDIV(self):
      assert(assembler.idivreg64('r11') == [0x49, 0xf7, 0xfb])
      assert(assembler.idivreg64('rcx') == [0x48, 0xf7, 0xf9])
      assert(assembler.idivreg64('rsp') == [0x48, 0xf7, 0xfc])

   def testAssemblerIMUL(self):
      assert(assembler.imulreg64_rax('rdi') == [0x48, 0xf7, 0xef])
      assert(assembler.imulreg64_rax('r10') == [0x49, 0xf7, 0xea])
      assert(assembler.imulreg64_rax('rdx') == [0x48, 0xf7, 0xea])

      assert(assembler.imulreg64('r11', 'rdi') == [0x4c, 0xf, 0xaf, 0xdf])
      assert(assembler.imulreg64('r12', 'rbx') == [0x4c, 0xf, 0xaf, 0xe3])
      # nasm generates this machine code: 0x4d, 0x6b, 0xff, 0xee
      # This also works: 4D0FAFFE (another variant?? )
      assert(assembler.imulreg64('r15', 'r14') == [0x4d, 0x0f, 0xaf, 0xfe])
   def testProject(self):
      p = Project('test.xml', isnew=True)
      p.name = "Test project"
      p.files.append('main.mod')
      p.files.append('test.mod')
      p.save('test.xml')

      q = Project('test.xml')

      assert(p.name == q.name)
      assert(p.files == q.files)
      # TODO: remove test.xml test file
      os.remove('test.xml')

if __name__ == '__main__':
   unittest.main()