import unittest
import dis
from opcode import *


class TestFindLabels(unittest.TestCase):
    def test_findlabels(self):
        # Issue 26448 - findlabels ignoring extended arguments.

        # Note: Do not duplicate values in args, as we calculate *all*
        #       labels for our expected results, whereas findlabels
        #       returns a unique set.
        args = [
                # TODO figure out a good set of edge cases here.
                0x0000,
                0xFF00,
                0x0003,
                0x00FF,
                0xFFFF,
                0x00010100,
                0x0006,
                0x00020000,
                0xFFFF0000,
                0xFFFFFFFF,
                0xFFFFFF00,
                0x0001FFFFFFFF,
                0xFFFF0001FFFFFFFF,
                0x0009,
        ]

        templates = []
        offsets = []
        prev_offset = 0
        for arg in args:
            template = b"%%b%2b" % (arg & 0xFFFF).to_bytes(2, "little")
            while arg > 0xFFFF:
                arg >>= 16
                template = b"%b%2b" % (
                        EXTENDED_ARG.to_bytes(1, "little"),
                        (arg & 0xFFFF).to_bytes(2, "little"),
                ) + template

            templates.append(template)
            offsets.append(prev_offset + len(template) - template.count(b"%b"))
            prev_offset = offsets[-1]

        # Calculate the expected results for the three cases of argument-bearing
        # opcodes: non-branching, absolute jumps and relative jumps.
        non_jump_labels = []
        abs_jump_labels = args
        rel_jump_labels = [arg + ofs for arg, ofs in zip(args, offsets)]

# scaffolding: delete for real use
        print()
        print("raw offsets   ", offsets)
        print("absolute jumps", abs_jump_labels)
        print("relative jumps", rel_jump_labels)
        for opcode, expected in ((opmap["LOAD_GLOBAL"],   non_jump_labels),
                                 (opmap["JUMP_ABSOLUTE"], abs_jump_labels),
                                 (opmap["JUMP_FORWARD"],  rel_jump_labels)):
            print()
            print("Bytecode snippets:")
            for template in templates:
                print("raw template", template)
                expanded = template % opcode.to_bytes(1, "little")
                print("expanded    ", expanded)
                dis.dis(expanded)

            print("All together now:")
            bytecode = b"".join(code % opcode.to_bytes(1, "little") for code in templates)
            dis.dis(bytecode)

            print("Check the findlabels output:")
            labels = dis.findlabels(bytecode)
            print("expected", expected)
            print("actual  ", labels)
        print()
# end scaffolding

        for opcode in range(256):
            # Ignore the COMPARE_xxx opcodes, as they are not extendable.  Any others?
            if opcode >= HAVE_ARGUMENT and not opcode in hascompare:
                bytecode = b"".join(code % opcode.to_bytes(1, "little") for code in templates)
                labels = dis.findlabels(bytecode)
                if opcode in hasjrel:
                    self.assertEqual(labels, rel_jump_labels)
                elif opcode in hasjabs:
                    self.assertEqual(labels, abs_jump_labels)
                else:
                    self.assertEqual(labels, non_jump_labels)

                    
suite = unittest.TestLoader().loadTestsFromTestCase(TestFindLabels)
unittest.TextTestRunner(verbosity=2).run(suite)