#!/usr/bin/python

import unittest, cProfile
from ppci import CompilerError
from asmnodes import AInstruction, ABinop, AUnop, ASymbol, ALabel, ANumber
from asm import tokenize, Assembler
import msp430
import cortexm3 as arm
import outstream
from target import Label

class AssemblerLexingCase(unittest.TestCase):
    """ Tests the assemblers lexer """

    def testLex0(self):
        """ Check if the lexer is OK """
        asmline, toks = 'mov rax, rbx ', ['ID', 'ID', ',', 'ID']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

    def testLex1(self):
        """ Test if lexer correctly maps some tokens """
        asmline, toks = 'lab1: mov rax, rbx ', ['ID', ':', 'ID', 'ID', ',', 'ID']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

    def testLex1(self):
        """ Test if lexer correctly maps some tokens """
        asmline, toks = 'mov 3.13 0xC 13', ['ID', 'REAL', 'NUMBER', 'NUMBER']
        self.assertSequenceEqual([tok.typ for tok in tokenize(asmline)], toks)

    def testLex2(self):
        """ Test if lexer fails on a token that is invalid """
        asmline = '0z4: mov rax, rbx $ '
        with self.assertRaises(CompilerError):
            list(tokenize(asmline))

class AssemblerParsingTestCase(unittest.TestCase):
    """ 
        Tests the assembler parts
    """
    def setUp(self):
        self.a = Assembler()

    def testParse(self):
        asmline = 'lab1: mov rax, rbx'
        self.a.parse_line(asmline)

    def expectTree(self, asmline, stack):
        self.a.parse_line(asmline)
        self.assertSequenceEqual(stack, self.a.stack)

    def testParse2(self):
        asmline = 'a: mov rax, [rbx + 2]'
        output = []
        output.append(ALabel('a'))
        output.append(AInstruction('mov', [ASymbol('rax'), AUnop('[]', ASymbol('rbx') + ANumber(2))]))
        self.expectTree(asmline, output)

    def testParse3(self):
        # A label must be optional:
        asmline = 'mov rax, 1'
        output = [AInstruction('mov', [ASymbol('rax'), ANumber(1)])]
        self.expectTree(asmline, output)

    def testParse4(self):
        # Test 3 operands:
        asmline = 'add rax, [4*rbx + 22], rcx'
        ops = []
        ops.append(ASymbol('rax'))
        ops.append(AUnop('[]', ANumber(4) * ASymbol('rbx') + ANumber(22)))
        ops.append(ASymbol('rcx'))
        output = [AInstruction('add', ops)]
        self.expectTree(asmline, output)

    def testParse5(self):
        # An instruction must be optional:
        asmline = 'lab1:'
        output = []
        output.append(ALabel('lab1'))
        self.expectTree(asmline, output)

    def testParse6(self):
        # A line can be empty
        self.a.parse_line('')
    
class AssemblerOtherTestCase(unittest.TestCase):
    def testWithoutTarget(self):
        a = Assembler()
        with self.assertRaises(CompilerError):
            a.assemble_line('')

    @unittest.skip 
    def testX86(self):
        testsrc = """ ; tst
        begin:
        mov rax, rbx ; 0x48, 0x89, 0xd8
        xor rcx, rbx ; 0x48, 0x31, 0xd9
        inc rcx ; 0x48 0xff 0xc1
        """
        a = Assembler()
        a.assemble(testsrc)
        # Compare with nasm output:
        nasmbytes = [0x48, 0x89, 0xd8, 0x48, 0x31, 0xd9, 0x48, 0xff, 0xc1]


class OustreamTestCase(unittest.TestCase):
    def test1(self):
        o = outstream.BinOutputStream()
        o.selectSection('.text')
        o.emit(Label('a'))
        self.assertSequenceEqual(bytes(), o.Data)


class AsmTestCaseBase(unittest.TestCase):
    def feed(self, line):
        self.a.assemble(line)

    def check(self, hexstr):
        self.assertSequenceEqual(bytes.fromhex(hexstr), self.o.Data)


class AssemblerMSP430TestCase(AsmTestCaseBase):
    def setUp(self):
        self.t = msp430.msp430target
        self.o = outstream.BinOutputStream()
        self.o.selectSection('.text')
        self.a = Assembler(target=self.t, stream=self.o)

    def testMapMovInstruction(self):
        i = AInstruction('mov', [ASymbol('r14'), ASymbol('r15')])
        ri = self.t.mapInstruction(i)

    def testMapRetiInstruction(self):
        i = AInstruction('reti', [])
        ri = self.t.mapInstruction(i)

    @unittest.skip
    def testMapOperand(self):
        o = ASymbol('r14')
        mo = self.t.mapOperand(o)
        self.assertEqual(mo, msp430.r14)

    @unittest.skip
    def testMapOperandIndirection(self):
        o = AUnop('[]', ASymbol('r14'))
        mo = self.t.mapOperand(o)

    def testMov(self):
        line1 = "mov r14, r15"
        self.feed(line1)
        self.check('0F4E')

    def testMov1337(self):
        line1 = "mov 0x1337, r12"
        self.feed(line1)
        self.check('3C403713')

    def testAdd(self):
        line1 = "add r15, r13"
        self.feed(line1)
        self.check('0D5F')

    def testReti(self):
        line1 = "reti"
        self.feed(line1)
        self.check('0013')

    def testMSPinstructionCount(self):
        """ Check that there are 27 instructions """
        self.assertEqual(27, len(self.t.instructions))


class AssemblerARMTestCase(AsmTestCaseBase):
    def setUp(self):
        self.t = arm.armtarget
        self.o = outstream.BinOutputStream()
        self.o.selectSection('.text')
        self.a = Assembler(target=self.t, stream=self.o)

    def testMapOperand(self):
        pass

    def testMovImm8(self):
        self.feed('mov r4, 100')
        self.check('6424')

    def testYield(self):
        self.feed('yield')
        self.check('10bf')

    def testPush(self):
        self.feed('push {r2,r3,lr}')
        self.check('0cb5')

    def testPop(self):
        self.feed('pop {r4-r6, pc}')
        self.check('70bd')

    def testStr5(self):
        self.feed('str r4, [r1 + 0]')
        self.check('0c60')

    def testLdr5(self):
        self.feed('ldr r4, [r0 + 0]')
        self.check('0468')

    def testLdrSpRel(self):
        self.feed('ldr r0, [sp + 4]')
        self.check('0198')

    def testStrSpRel(self):
        self.feed('str r0, [sp + 4]')
        self.check('0190')

    def testLdrPcRel(self):
        self.feed('ldr r7, henkie')
        self.feed('ldr r6, henkie')
        self.feed('ldr r1, henkie')
        self.feed('align 4')
        self.feed('dcd 1')
        self.feed('henkie: dcd 2')
        self.check('024F024E 01490000 01000000 02000000')

    def testBranch(self):
        self.feed('start: b henkie')
        self.feed('beq henkie')
        self.feed('bne henkie')
        self.feed('henkie: b start')
        self.feed('eof: b eof')
        self.check('01e000d0 ffd1fbe7 fee7')

    def testConditions(self):
        self.feed('blt x')
        self.feed('bgt x')
        self.feed('x:')
        self.check('00dbffdc')

    def testBoff(self):
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('henkie:')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.feed('b henkie')
        self.check('05e004e0 03e002e0 01e000e0 ffe7fee7 fde7fce7 fbe7')

    def testBl(self):
        self.feed('bl henkie')
        self.feed('bl henkie')
        self.feed('henkie:')
        self.feed('bl henkie')
        self.feed('bl henkie')
        self.check('00f0 02f8 00f0 00f8 fff7 feff fff7 fcff')

    def testCmpRegReg(self):
        self.feed('cmp r0, r1')
        self.check('8842')

    def testLeftShit(self):
        self.feed('lsl r3, r5')
        self.check('ab40')

    def testSequence1(self):
        self.feed('mov r5, 3')
        self.feed('add r4, r5, 0')
        self.feed('loop: add r6, r4, 7')
        self.feed('cmp r6, 5')
        self.check('0325 2c1c e61d 052e')

    def testSequence2(self):
        self.feed('henkie:')
        self.feed('push {r1,r4,r5}')
        self.feed('add r5, r2, r4')
        self.feed('cmp r4, r2')
        self.feed('ldr r0, [sp + 4]')
        self.feed('str r3, [sp + 16]')
        self.feed('pop {r1, r4, r5}')
        self.feed('lsl r3, r4')
        self.feed('cmp r3, r5')
        self.feed('beq henkie')
        self.feed('bne henkie')
        self.feed('b henkie')
        self.check('32b41519 94420198 049332bc a340ab42 f6d0f5d1 f4e7')

if __name__ == '__main__':
    #cProfile.run('unittest.main()')
    unittest.main()