#!/usr/bin/env python3
"""
Unit tests for IACA/OSACA marker utilities
"""
import os
import unittest
from collections import OrderedDict

from osaca.semantics import (
    reduce_to_section,
    find_basic_blocks,
    find_jump_labels,
    find_basic_loop_bodies,
)
from osaca.parser import ParserAArch64, ParserX86ATT


class TestMarkerUtils(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        self.parser_AArch = ParserAArch64()
        self.parser_x86 = ParserX86ATT()
        with open(self._find_file("triad_arm_iaca.s")) as f:
            triad_code_arm = f.read()
        with open(self._find_file("triad_x86_iaca.s")) as f:
            triad_code_x86 = f.read()
        self.parsed_AArch = self.parser_AArch.parse_file(triad_code_arm)
        self.parsed_x86 = self.parser_x86.parse_file(triad_code_x86)

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

    def test_marker_detection_AArch64(self):
        kernel = reduce_to_section(self.parsed_AArch, "AArch64")
        self.assertEqual(len(kernel), 138)
        self.assertEqual(kernel[0].line_number, 307)
        self.assertEqual(kernel[-1].line_number, 444)

    def test_marker_detection_x86(self):
        kernel = reduce_to_section(self.parsed_x86, "x86")
        self.assertEqual(len(kernel), 9)
        self.assertEqual(kernel[0].line_number, 146)
        self.assertEqual(kernel[-1].line_number, 154)

    def test_marker_matching_AArch64(self):
        # preparation
        bytes_1_line = ".byte     213,3,32,31\n"
        bytes_2_lines_1 = ".byte     213,3,32\n" + ".byte 31\n"
        bytes_2_lines_2 = ".byte     213,3\n" + ".byte 32,31\n"
        bytes_2_lines_3 = ".byte     213\n" + ".byte 3,32,31\n"
        bytes_3_lines_1 = ".byte     213,3\n" + ".byte     32\n" + ".byte     31\n"
        bytes_3_lines_2 = ".byte     213\n" + ".byte     3,32\n" + ".byte     31\n"
        bytes_3_lines_3 = ".byte     213\n" + ".byte     3\n" + ".byte     32,31\n"
        bytes_4_lines = ".byte     213\n" + ".byte     3\n" + ".byte     32\n" + ".byte     31\n"
        mov_start_1 = "mov      x1, #111\n"
        mov_start_2 = "mov      x1, 111  // should work as well\n"
        mov_end_1 = "mov      x1, #222 // preferred way\n"
        mov_end_2 = "mov      x1, 222\n"
        prologue = (
            "mov x12, xzr\n"
            + "\tldp x9, x10, [sp, #16]      // 8-byte Folded Reload\n"
            + "     .p2align    6\n"
        )
        kernel = (
            ".LBB0_28:\n"
            + "fmul    v7.2d, v7.2d, v19.2d\n"
            + "stp q0, q1, [x10, #-32]\n"
            + "b.ne    .LBB0_28\n"
        )
        epilogue = ".LBB0_29:   //   Parent Loop BB0_20 Depth=1\n" + "bl    dummy\n"
        kernel_length = len(list(filter(None, kernel.split("\n"))))

        bytes_variations = [
            bytes_1_line,
            bytes_2_lines_1,
            bytes_2_lines_2,
            bytes_2_lines_3,
            bytes_3_lines_1,
            bytes_3_lines_2,
            bytes_3_lines_3,
            bytes_4_lines,
        ]
        mov_start_variations = [mov_start_1, mov_start_2]
        mov_end_variations = [mov_end_1, mov_end_2]
        # actual tests
        for mov_start_var in mov_start_variations:
            for bytes_var_1 in bytes_variations:
                for mov_end_var in mov_end_variations:
                    for bytes_var_2 in bytes_variations:
                        sample_code = (
                            prologue
                            + mov_start_var
                            + bytes_var_1
                            + kernel
                            + mov_end_var
                            + bytes_var_2
                            + epilogue
                        )
                        with self.subTest(
                            mov_start=mov_start_var,
                            bytes_start=bytes_var_1,
                            mov_end=mov_end_var,
                            bytes_end=bytes_var_2,
                        ):
                            sample_parsed = self.parser_AArch.parse_file(sample_code)
                            sample_kernel = reduce_to_section(sample_parsed, "AArch64")
                            self.assertEqual(len(sample_kernel), kernel_length)
                            kernel_start = len(
                                list(
                                    filter(
                                        None,
                                        (prologue + mov_start_var + bytes_var_1).split("\n"),
                                    )
                                )
                            )
                            parsed_kernel = self.parser_AArch.parse_file(
                                kernel, start_line=kernel_start
                            )
                            self.assertEqual(sample_kernel, parsed_kernel)

    def test_marker_matching_x86(self):
        # preparation
        bytes_1_line = ".byte     100,103,144\n"
        bytes_2_lines_1 = ".byte     100,103\n" + ".byte 144\n"
        bytes_2_lines_2 = ".byte     100\n" + ".byte 103,144\n"
        bytes_3_lines = (
            ".byte     100 # IACA MARKER UTILITY\n"
            + ".byte     103 # IACA MARKER UTILITY\n"
            + ".byte     144 # IACA MARKER UTILITY\n"
        )
        mov_start_1 = "movl      $111, %ebx # IACA START\n"
        mov_start_2 = "mov      $111, %ebx # IACA START\n"
        mov_end_1 = "movl      $222, %ebx # IACA END\n"
        mov_end_2 = "mov      $222, %ebx # IACA END\n"
        prologue = "movl    -92(%rbp), %r11d\n" + "movl      $111, %ebx\n"
        kernel = (
            "vfmadd132sd (%r15,%rcx,8), %xmm5, %xmm0\n"
            + "vmovsd  %xmm0, (%r14,%rcx,8)\n"
            + "cmpl    %ebx, %ecx\n"
            + "jge .L8\n"
        )
        epilogue = ".LE9:\t\t#12.2\n" "call    dummy\n"
        kernel_length = len(list(filter(None, kernel.split("\n"))))

        bytes_variations = [
            bytes_1_line,
            bytes_2_lines_1,
            bytes_2_lines_2,
            bytes_3_lines,
        ]
        mov_start_variations = [mov_start_1, mov_start_2]
        mov_end_variations = [mov_end_1, mov_end_2]
        # actual tests
        for mov_start_var in mov_start_variations:
            for bytes_var_1 in bytes_variations:
                for mov_end_var in mov_end_variations:
                    for bytes_var_2 in bytes_variations:
                        sample_code = (
                            prologue
                            + mov_start_var
                            + bytes_var_1
                            + kernel
                            + mov_end_var
                            + bytes_var_2
                            + epilogue
                        )
                        with self.subTest(
                            mov_start=mov_start_var,
                            bytes_start=bytes_var_1,
                            mov_end=mov_end_var,
                            bytes_end=bytes_var_2,
                        ):
                            sample_parsed = self.parser_x86.parse_file(sample_code)
                            sample_kernel = reduce_to_section(sample_parsed, "x86")
                            self.assertEqual(len(sample_kernel), kernel_length)
                            kernel_start = len(
                                list(
                                    filter(
                                        None,
                                        (prologue + mov_start_var + bytes_var_1).split("\n"),
                                    )
                                )
                            )
                            parsed_kernel = self.parser_x86.parse_file(
                                kernel, start_line=kernel_start
                            )
                            self.assertEqual(sample_kernel, parsed_kernel)

    def test_marker_special_cases_AArch(self):
        bytes_line = ".byte     213,3,32,31\n"
        start_marker = "mov      x1, #111\n" + bytes_line
        end_marker = "mov      x1, #222\n" + bytes_line
        prologue = "dup v0.2d, x14\n" "neg x9, x9\n" ".p2align    6\n"
        kernel = (
            ".LBB0_28:\n"
            + "fmul    v7.2d, v7.2d, v19.2d\n"
            + "stp q0, q1, [x10, #-32]\n"
            + "b.ne    .LBB0_28\n"
        )
        epilogue = ".LBB0_29:   //   Parent Loop BB0_20 Depth=1\n" "bl    dummy\n"

        samples = [
            # (test name,
            #  ignored prologue, section to be extraced, ignored epilogue)
            ("markers", prologue + start_marker, kernel, end_marker + epilogue),
            ("marker at file start", start_marker, kernel, end_marker + epilogue),
            ("no start marker", "", prologue + kernel, end_marker + epilogue),
            ("marker at file end", prologue + start_marker, kernel, end_marker),
            ("no end marker", prologue + start_marker, kernel + epilogue, ""),
            ("empty kernel", prologue + start_marker, "", end_marker + epilogue),
        ]

        for test_name, pro, kernel, epi in samples:
            code = pro + kernel + epi
            parsed = self.parser_AArch.parse_file(code)
            test_kernel = reduce_to_section(parsed, "AArch64")
            if kernel:
                kernel_length = len(kernel.strip().split("\n"))
            else:
                kernel_length = 0
            self.assertEqual(
                len(test_kernel),
                kernel_length,
                msg="Invalid exctracted kernel length on {!r} sample".format(test_name),
            )
            if pro:
                kernel_start = len((pro).strip().split("\n"))
            else:
                kernel_start = 0
            parsed_kernel = self.parser_AArch.parse_file(kernel, start_line=kernel_start)
            self.assertEqual(
                test_kernel,
                parsed_kernel,
                msg="Invalid exctracted kernel on {!r}".format(test_name),
            )

    def test_marker_special_cases_x86(self):
        bytes_line = ".byte     100\n" ".byte     103\n" ".byte     144\n"
        start_marker = "movl     $111, %ebx\n" + bytes_line
        end_marker = "movl     $222, %ebx\n" + bytes_line
        prologue = "movl    -88(%rbp), %r10d\n" "xorl    %r11d, %r11d\n" ".p2align 4,,10\n"
        kernel = (
            ".L3: #L3\n"
            "vmovsd  .LC1(%rip), %xmm0\n"
            "vmovsd  %xmm0, (%r15,%rcx,8)\n"
            "cmpl    %ecx, %ebx\n"
            "jle .L3\n"
        )
        epilogue = "leaq    -56(%rbp), %rsi\n" "movl    %r10d, -88(%rbp)\n" "call    timing\n"
        samples = [
            # (test name,
            #  ignored prologue, section to be extraced, ignored epilogue)
            ("markers", prologue + start_marker, kernel, end_marker + epilogue),
            ("marker at file start", start_marker, kernel, end_marker + epilogue),
            ("no start marker", "", prologue + kernel, end_marker + epilogue),
            ("marker at file end", prologue + start_marker, kernel, end_marker),
            ("no end marker", prologue + start_marker, kernel + epilogue, ""),
            ("empty kernel", prologue + start_marker, "", end_marker + epilogue),
        ]

        for test_name, pro, kernel, epi in samples:
            code = pro + kernel + epi
            parsed = self.parser_x86.parse_file(code)
            test_kernel = reduce_to_section(parsed, "x86")
            if kernel:
                kernel_length = len(kernel.strip().split("\n"))
            else:
                kernel_length = 0
            self.assertEqual(
                len(test_kernel),
                kernel_length,
                msg="Invalid exctracted kernel length on {!r} sample".format(test_name),
            )
            if pro:
                kernel_start = len((pro).strip().split("\n"))
            else:
                kernel_start = 0
            parsed_kernel = self.parser_x86.parse_file(kernel, start_line=kernel_start)
            self.assertEqual(
                test_kernel,
                parsed_kernel,
                msg="Invalid exctracted kernel on {!r}".format(test_name),
            )

    def test_find_jump_labels(self):
        self.assertEqual(
            find_jump_labels(self.parsed_x86),
            OrderedDict(
                [
                    (".LFB24", 10),
                    (".L4", 65),
                    (".L3", 79),
                    (".L2", 102),
                    (".L13", 111),
                    (".L12", 120),
                    (".L6", 132),
                    (".L10", 145),
                    (".L9", 161),
                    (".L8", 183),
                    (".L15", 252),
                    (".L26", 256),
                    (".L14", 259),
                    (".LFB25", 277),
                    (".L28", 289),
                ]
            ),
        )

        self.assertEqual(
            find_jump_labels(self.parsed_AArch),
            OrderedDict(
                [
                    ("triad", 18),
                    (".LBB0_3", 71),
                    (".LBB0_4", 76),
                    (".LBB0_5", 84),
                    (".LBB0_7", 92),
                    (".LBB0_8", 95),
                    (".LBB0_9", 106),
                    (".LBB0_11", 118),
                    (".LBB0_12", 133),
                    (".LBB0_14", 177),
                    (".LBB0_15", 190),
                    (".LBB0_16", 205),
                    (".LBB0_17", 208),
                    (".LBB0_18", 221),
                    (".LBB0_19", 228),
                    (".LBB0_20", 260),
                    (".LBB0_22", 272),
                    (".LBB0_24", 283),
                    (".LBB0_26", 290),
                    (".LBB0_28", 298),
                    (".LBB0_29", 306),
                    (".LBB0_31", 448),
                    (".LBB0_32", 458),
                    (".LBB0_33", 480),
                    (".LBB0_34", 484),
                    (".LBB0_35", 493),
                    (".LBB0_36", 504),
                    (".LBB0_37", 508),
                    (".LBB0_38", 518),
                    ("main", 574),
                ]
            ),
        )

    def test_find_basic_blocks(self):
        self.assertEqual(
            [
                (k, v[0]["line_number"], v[-1]["line_number"])
                for k, v in find_basic_blocks(self.parsed_x86).items()
            ],
            [
                (".LFB24", 11, 56),
                (".L4", 66, 74),
                (".L3", 80, 89),
                (".L2", 103, 112),
                (".L13", 112, 121),
                (".L12", 121, 125),
                (".L6", 133, 135),
                (".L10", 146, 154),
                (".L9", 162, 170),
                (".L8", 184, 187),
                (".L15", 253, 256),
                (".L26", 257, 259),
                (".L14", 260, 262),
                (".LFB25", 278, 290),
                (".L28", 290, 300),
            ],
        )

        self.assertEqual(
            [
                (k, v[0]["line_number"], v[-1]["line_number"])
                for k, v in find_basic_blocks(self.parsed_AArch).items()
            ],
            [
                ("triad", 19, 64),
                (".LBB0_3", 72, 77),
                (".LBB0_4", 77, 83),
                (".LBB0_5", 85, 89),
                (".LBB0_7", 93, 95),
                (".LBB0_8", 96, 105),
                (".LBB0_9", 107, 114),
                (".LBB0_11", 119, 134),
                (".LBB0_12", 134, 173),
                (".LBB0_14", 178, 191),
                (".LBB0_15", 191, 205),
                (".LBB0_16", 206, 208),
                (".LBB0_17", 209, 222),
                (".LBB0_18", 222, 228),
                (".LBB0_19", 229, 261),
                (".LBB0_20", 261, 269),
                (".LBB0_22", 273, 280),
                (".LBB0_24", 284, 286),
                (".LBB0_26", 291, 293),
                (".LBB0_28", 299, 307),
                (".LBB0_29", 307, 444),
                (".LBB0_31", 449, 459),
                (".LBB0_32", 459, 480),
                (".LBB0_33", 481, 484),
                (".LBB0_34", 485, 494),
                (".LBB0_35", 494, 504),
                (".LBB0_36", 505, 508),
                (".LBB0_37", 509, 518),
                (".LBB0_38", 519, 568),
                ("main", 575, 590),
            ],
        )

    def test_find_basic_loop_body(self):
        self.assertEqual(
            [
                (k, v[0]["line_number"], v[-1]["line_number"])
                for k, v in find_basic_loop_bodies(self.parsed_x86).items()
            ],
            [(".L4", 66, 74), (".L10", 146, 154), (".L28", 290, 300)],
        )

        self.assertEqual(
            [
                (k, v[0]["line_number"], v[-1]["line_number"])
                for k, v in find_basic_loop_bodies(self.parsed_AArch).items()
            ],
            [
                (".LBB0_12", 134, 173),
                (".LBB0_15", 191, 205),
                (".LBB0_18", 222, 228),
                (".LBB0_29", 307, 444),
                (".LBB0_32", 459, 480),
                (".LBB0_35", 494, 504),
            ],
        )

    ##################
    # Helper functions
    ##################

    @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(TestMarkerUtils)
    unittest.TextTestRunner(verbosity=2).run(suite)
