#!/usr/bin/env python3
"""
Unit tests for x86 AT&T assembly parser
"""

import os
import unittest

from pyparsing import ParseException

from osaca.parser import AttrDict, ParserX86ATT


class TestParserX86ATT(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        self.parser = ParserX86ATT()
        with open(self._find_file('triad_x86_iaca.s')) as f:
            self.triad_code = f.read()

    ##################
    # Test
    ##################

    def test_comment_parser(self):
        self.assertEqual(self._get_comment(self.parser, '# some comments'), 'some comments')
        self.assertEqual(self._get_comment(self.parser, '\t\t#AA BB CC \t end \t'), 'AA BB CC end')
        self.assertEqual(
            self._get_comment(self.parser, '\t## comment ## comment'), '# comment ## comment'
        )

    def test_label_parser(self):
        self.assertEqual(self._get_label(self.parser, 'main:').name, 'main')
        self.assertEqual(self._get_label(self.parser, '..B1.10:').name, '..B1.10')
        self.assertEqual(self._get_label(self.parser, '.2.3_2_pack.3:').name, '.2.3_2_pack.3')
        self.assertEqual(self._get_label(self.parser, '.L1:\t\t\t#label1').name, '.L1')
        self.assertEqual(
            ' '.join(self._get_label(self.parser, '.L1:\t\t\t#label1').comment), 'label1'
        )
        with self.assertRaises(ParseException):
            self._get_label(self.parser, '\t.cfi_startproc')

    def test_directive_parser(self):
        self.assertEqual(self._get_directive(self.parser, '\t.text').name, 'text')
        self.assertEqual(len(self._get_directive(self.parser, '\t.text').parameters), 0)
        self.assertEqual(self._get_directive(self.parser, '\t.align\t16,0x90').name, 'align')
        self.assertEqual(len(self._get_directive(self.parser, '\t.align\t16,0x90').parameters), 2)
        self.assertEqual(len(self._get_directive(self.parser, '.text').parameters), 0)
        self.assertEqual(
            len(self._get_directive(self.parser, '.file\t1 "path/to/file.c"').parameters), 2
        )
        self.assertEqual(
            self._get_directive(self.parser, '.file\t1 "path/to/file.c"').parameters[1],
            '"path/to/file.c"',
        )
        self.assertEqual(
            self._get_directive(self.parser, '\t.set\tL$set$0,LECIE1-LSCIE1').parameters,
            ['L$set$0', 'LECIE1-LSCIE1'],
        )
        self.assertEqual(
            self._get_directive(
                self.parser,
                '\t.section __TEXT,__eh_frame,coalesced,no_toc+strip_static_syms+live_support',
            ).parameters,
            ['__TEXT', '__eh_frame', 'coalesced', 'no_toc+strip_static_syms+live_support'],
        )
        self.assertEqual(
            self._get_directive(
                self.parser, '\t.section\t__TEXT,__literal16,16byte_literals'
            ).parameters,
            ['__TEXT', '__literal16', '16byte_literals'],
        )
        self.assertEqual(
            self._get_directive(self.parser, '\t.align\t16,0x90').parameters[1], '0x90'
        )
        self.assertEqual(
            self._get_directive(self.parser, '        .byte 100,103,144       #IACA START')[
                'name'
            ],
            'byte',
        )
        self.assertEqual(
            self._get_directive(self.parser, '        .byte 100,103,144       #IACA START')[
                'parameters'
            ][2],
            '144',
        )
        self.assertEqual(
            ' '.join(
                self._get_directive(self.parser, '        .byte 100,103,144       #IACA START')[
                    'comment'
                ]
            ),
            'IACA START',
        )

    def test_parse_instruction(self):
        instr1 = '\t\tvcvtsi2ss %edx, %xmm2, %xmm2\t\t\t#12.27'
        instr2 = 'jb        ..B1.4 \t'
        instr3 = '        movl $222,%ebx          #IACA END'
        instr4 = 'vmovss    %xmm4, -4(%rsp,%rax,8) #12.9'
        instr5 = 'mov %ebx,var(,1)'
        instr6 = 'lea (,%rax,8),%rbx'
        instr7 = 'vinsertf128 $0x1, %xmm0, %ymm1, %ymm1'

        parsed_1 = self.parser.parse_instruction(instr1)
        parsed_2 = self.parser.parse_instruction(instr2)
        parsed_3 = self.parser.parse_instruction(instr3)
        parsed_4 = self.parser.parse_instruction(instr4)
        parsed_5 = self.parser.parse_instruction(instr5)
        parsed_6 = self.parser.parse_instruction(instr6)
        parsed_7 = self.parser.parse_instruction(instr7)

        self.assertEqual(parsed_1.instruction, 'vcvtsi2ss')
        self.assertEqual(parsed_1.operands[0].register.name, 'edx')
        self.assertEqual(parsed_1.operands[1].register.name, 'xmm2')
        self.assertEqual(parsed_1.comment, '12.27')

        self.assertEqual(parsed_2.instruction, 'jb')
        self.assertEqual(parsed_2.operands[0].identifier.name, '..B1.4')
        self.assertEqual(len(parsed_2.operands), 1)
        self.assertIsNone(parsed_2.comment)

        self.assertEqual(parsed_3.instruction, 'movl')
        self.assertEqual(parsed_3.operands[0].immediate.value, '222')
        self.assertEqual(parsed_3.operands[1].register.name, 'ebx')
        self.assertEqual(parsed_3.comment, 'IACA END')

        self.assertEqual(parsed_4.instruction, 'vmovss')
        self.assertEqual(parsed_4.operands[1].memory.offset.value, '-4')
        self.assertEqual(parsed_4.operands[1].memory.base.name, 'rsp')
        self.assertEqual(parsed_4.operands[1].memory.index.name, 'rax')
        self.assertEqual(parsed_4.operands[1].memory.scale, 8)
        self.assertEqual(parsed_4.operands[0].register.name, 'xmm4')
        self.assertEqual(parsed_4.comment, '12.9')

        self.assertEqual(parsed_5.instruction, 'mov')
        self.assertEqual(parsed_5.operands[1].memory.offset.identifier.name, 'var')
        self.assertIsNone(parsed_5.operands[1].memory.base)
        self.assertIsNone(parsed_5.operands[1].memory.index)
        self.assertEqual(parsed_5.operands[1].memory.scale, 1)
        self.assertEqual(parsed_5.operands[0].register.name, 'ebx')

        self.assertEqual(parsed_6.instruction, 'lea')
        self.assertIsNone(parsed_6.operands[0].memory.offset)
        self.assertIsNone(parsed_6.operands[0].memory.base)
        self.assertEqual(parsed_6.operands[0].memory.index.name, 'rax')
        self.assertEqual(parsed_6.operands[0].memory.scale, 8)
        self.assertEqual(parsed_6.operands[1].register.name, 'rbx')

        self.assertEqual(parsed_7.operands[0].immediate.value, '0x1')
        self.assertEqual(parsed_7.operands[1].register.name, 'xmm0')
        self.assertEqual(parsed_7.operands[2].register.name, 'ymm1')
        self.assertEqual(parsed_7.operands[3].register.name, 'ymm1')

    def test_parse_line(self):
        line_comment = '# -- Begin  main'
        line_label = '..B1.7:                         # Preds ..B1.6'
        line_directive = '.quad   .2.3_2__kmpc_loc_pack.2 #qed'
        line_instruction = 'lea       2(%rax,%rax), %ecx #12.9'

        instruction_form_1 = {
            'instruction': None,
            'operands': [],
            'directive': None,
            'comment': '-- Begin main',
            'label': None,
            'line': '# -- Begin  main',
            'line_number': 1,
        }
        instruction_form_2 = {
            'instruction': None,
            'operands': [],
            'directive': None,
            'comment': 'Preds ..B1.6',
            'label': '..B1.7',
            'line': '..B1.7:                         # Preds ..B1.6',
            'line_number': 2,
        }
        instruction_form_3 = {
            'instruction': None,
            'operands': [],
            'directive': {'name': 'quad', 'parameters': ['.2.3_2__kmpc_loc_pack.2']},
            'comment': 'qed',
            'label': None,
            'line': '.quad   .2.3_2__kmpc_loc_pack.2 #qed',
            'line_number': 3,
        }
        instruction_form_4 = {
            'instruction': 'lea',
            'operands': [
                {
                    'memory': {
                        'offset': {'value': '2'},
                        'base': {'name': 'rax'},
                        'index': {'name': 'rax'},
                        'scale': 1,
                    }
                },
                {'register': {'name': 'ecx'}},
            ],
            'directive': None,
            'comment': '12.9',
            'label': None,
            'line': 'lea       2(%rax,%rax), %ecx #12.9',
            'line_number': 4,
        }

        parsed_1 = self.parser.parse_line(line_comment, 1)
        parsed_2 = self.parser.parse_line(line_label, 2)
        parsed_3 = self.parser.parse_line(line_directive, 3)
        parsed_4 = self.parser.parse_line(line_instruction, 4)

        self.assertEqual(parsed_1, instruction_form_1)
        self.assertEqual(parsed_2, instruction_form_2)
        self.assertEqual(parsed_3, instruction_form_3)
        self.assertEqual(parsed_4, instruction_form_4)

    def test_parse_file(self):
        parsed = self.parser.parse_file(self.triad_code)
        self.assertEqual(parsed[0].line_number, 1)
        self.assertEqual(len(parsed), 353)

    def test_parse_register(self):
        register_str_1 = '%rax'
        register_str_2 = '%r9'
        register_str_3 = '%xmm1'
        register_str_4 = '%rip'

        parsed_reg_1 = {'register': {'name': 'rax'}}
        parsed_reg_2 = {'register': {'name': 'r9'}}
        parsed_reg_3 = {'register': {'name': 'xmm1'}}
        parsed_reg_4 = {'register': {'name': 'rip'}}

        self.assertEqual(self.parser.parse_register(register_str_1), parsed_reg_1)
        self.assertEqual(self.parser.parse_register(register_str_2), parsed_reg_2)
        self.assertEqual(self.parser.parse_register(register_str_3), parsed_reg_3)
        self.assertEqual(self.parser.parse_register(register_str_4), parsed_reg_4)
        self.assertIsNone(self.parser.parse_register('rax'))

    def test_normalize_imd(self):
        imd_decimal_1 = {'value': '79'}
        imd_hex_1 = {'value': '0x4f'}
        imd_decimal_2 = {'value': '8'}
        imd_hex_2 = {'value': '0x8'}
        self.assertEqual(
            self.parser.normalize_imd(imd_decimal_1), self.parser.normalize_imd(imd_hex_1)
        )
        self.assertEqual(
            self.parser.normalize_imd(imd_decimal_2), self.parser.normalize_imd(imd_hex_2)
        )

    def test_reg_dependency(self):
        reg_a1 = AttrDict({'name': 'rax'})
        reg_a2 = AttrDict({'name': 'eax'})
        reg_a3 = AttrDict({'name': 'ax'})
        reg_a4 = AttrDict({'name': 'al'})
        reg_r11 = AttrDict({'name': 'r11'})
        reg_r11b = AttrDict({'name': 'r11b'})
        reg_r11d = AttrDict({'name': 'r11d'})
        reg_r11w = AttrDict({'name': 'r11w'})
        reg_xmm1 = AttrDict({'name': 'xmm1'})
        reg_ymm1 = AttrDict({'name': 'ymm1'})
        reg_zmm1 = AttrDict({'name': 'zmm1'})

        reg_b1 = AttrDict({'name': 'rbx'})
        reg_r15 = AttrDict({'name': 'r15'})
        reg_xmm2 = AttrDict({'name': 'xmm2'})
        reg_ymm3 = AttrDict({'name': 'ymm3'})

        reg_a = [reg_a1, reg_a2, reg_a3, reg_a4]
        reg_r = [reg_r11, reg_r11b, reg_r11d, reg_r11w]
        reg_vec_1 = [reg_xmm1, reg_ymm1, reg_zmm1]
        reg_others = [reg_b1, reg_r15, reg_xmm2, reg_ymm3]
        regs = reg_a + reg_r + reg_vec_1 + reg_others

        # test each register against each other
        for ri in reg_a:
            for rj in regs:
                assert_value = True if rj in reg_a else False
                with self.subTest(reg_a=ri, reg_b=rj, assert_val=assert_value):
                    self.assertEqual(self.parser.is_reg_dependend_of(ri, rj), assert_value)
        for ri in reg_r:
            for rj in regs:
                assert_value = True if rj in reg_r else False
                with self.subTest(reg_a=ri, reg_b=rj, assert_val=assert_value):
                    self.assertEqual(self.parser.is_reg_dependend_of(ri, rj), assert_value)
        for ri in reg_vec_1:
            for rj in regs:
                assert_value = True if rj in reg_vec_1 else False
                with self.subTest(reg_a=ri, reg_b=rj, assert_val=assert_value):
                    self.assertEqual(self.parser.is_reg_dependend_of(ri, rj), assert_value)
        for ri in reg_others:
            for rj in regs:
                assert_value = True if rj == ri else False
                with self.subTest(reg_a=ri, reg_b=rj, assert_val=assert_value):
                    self.assertEqual(self.parser.is_reg_dependend_of(ri, rj), assert_value)

    ##################
    # Helper functions
    ##################
    def _get_comment(self, parser, comment):
        return ' '.join(
            AttrDict.convert_dict(
                parser.process_operand(parser.comment.parseString(comment, parseAll=True).asDict())
            ).comment
        )

    def _get_label(self, parser, label):
        return AttrDict.convert_dict(
            parser.process_operand(parser.label.parseString(label, parseAll=True).asDict())
        ).label

    def _get_directive(self, parser, directive):
        return AttrDict.convert_dict(
            parser.process_operand(parser.directive.parseString(directive, parseAll=True).asDict())
        ).directive

    @staticmethod
    def _find_file(name):
        testdir = os.path.dirname(__file__)
        name = os.path.join(testdir, 'test_files', name)
        assert os.path.exists(name)
        return name


if __name__ == '__main__':
    suite = unittest.TestLoader().loadTestsFromTestCase(TestParserX86ATT)
    unittest.TextTestRunner(verbosity=2).run(suite)
