gh-87092: Make jump target label equal to the offset of the target in the instructions sequence (#102093)

1 files changed, 36 insertions, 47 deletions

diff --git a/Lib/test/support/bytecode_helper.py b/Lib/test/support/bytecode_helper.py
index 65ae7a2..190fe87 100644
--- a/Lib/test/support/bytecode_helper.py
+++ b/Lib/test/support/bytecode_helper.py
@@ -50,18 +50,13 @@ class CompilationStepTestCase(unittest.TestCase):
     HAS_TARGET = set(dis.hasjrel + dis.hasjabs + dis.hasexc)
     HAS_ARG_OR_TARGET = HAS_ARG.union(HAS_TARGET)
 
-    def setUp(self):
-        self.last_label = 0
-
-    def Label(self):
-        self.last_label += 1
-        return self.last_label
+    class Label:
+        pass
 
     def assertInstructionsMatch(self, actual_, expected_):
         # get two lists where each entry is a label or
-        # an instruction tuple. Compare them, while mapping
-        # each actual label to a corresponding expected label
-        # based on their locations.
+        # an instruction tuple. Normalize the labels to the
+        # instruction count of the target, and compare the lists.
 
         self.assertIsInstance(actual_, list)
         self.assertIsInstance(expected_, list)
@@ -82,39 +77,35 @@ class CompilationStepTestCase(unittest.TestCase):
                 act = act[:len(exp)]
             self.assertEqual(exp, act)
 
+    def resolveAndRemoveLabels(self, insts):
+        idx = 0
+        res = []
+        for item in insts:
+            assert isinstance(item, (self.Label, tuple))
+            if isinstance(item, self.Label):
+                item.value = idx
+            else:
+                idx += 1
+                res.append(item)
+
+        return res
+
     def normalize_insts(self, insts):
         """ Map labels to instruction index.
-            Remove labels which are not used as jump targets.
             Map opcodes to opnames.
         """
-        labels_map = {}
-        targets = set()
-        idx = 1
-        for item in insts:
-            assert isinstance(item, (int, tuple))
-            if isinstance(item, tuple):
-                opcode, oparg, *_ = item
-                if dis.opmap.get(opcode, opcode) in self.HAS_TARGET:
-                    targets.add(oparg)
-                idx += 1
-            elif isinstance(item, int):
-                assert item not in labels_map, "label reused"
-                labels_map[item] = idx
-
+        insts = self.resolveAndRemoveLabels(insts)
         res = []
         for item in insts:
-            if isinstance(item, int) and item in targets:
-                if not res or labels_map[item] != res[-1]:
-                    res.append(labels_map[item])
-            elif isinstance(item, tuple):
-                opcode, oparg, *loc = item
-                opcode = dis.opmap.get(opcode, opcode)
-                if opcode in self.HAS_TARGET:
-                    arg = labels_map[oparg]
-                else:
-                    arg = oparg if opcode in self.HAS_TARGET else None
-                opcode = dis.opname[opcode]
-                res.append((opcode, arg, *loc))
+            assert isinstance(item, tuple)
+            opcode, oparg, *loc = item
+            opcode = dis.opmap.get(opcode, opcode)
+            if isinstance(oparg, self.Label):
+                arg = oparg.value
+            else:
+                arg = oparg if opcode in self.HAS_ARG else None
+            opcode = dis.opname[opcode]
+            res.append((opcode, arg, *loc))
         return res
 
 
@@ -129,20 +120,18 @@ class CfgOptimizationTestCase(CompilationStepTestCase):
 
     def complete_insts_info(self, insts):
         # fill in omitted fields in location, and oparg 0 for ops with no arg.
-        instructions = []
+        res = []
         for item in insts:
-            if isinstance(item, int):
-                instructions.append(item)
-            else:
-                assert isinstance(item, tuple)
-                inst = list(reversed(item))
-                opcode = dis.opmap[inst.pop()]
-                oparg = inst.pop() if opcode in self.HAS_ARG_OR_TARGET else 0
-                loc = inst + [-1] * (4 - len(inst))
-                instructions.append((opcode, oparg, *loc))
-        return instructions
+            assert isinstance(item, tuple)
+            inst = list(reversed(item))
+            opcode = dis.opmap[inst.pop()]
+            oparg = inst.pop() if opcode in self.HAS_ARG_OR_TARGET else 0
+            loc = inst + [-1] * (4 - len(inst))
+            res.append((opcode, oparg, *loc))
+        return res
 
     def get_optimized(self, insts, consts):
+        insts = self.normalize_insts(insts)
         insts = self.complete_insts_info(insts)
         insts = optimize_cfg(insts, consts)
         return insts, consts