gh-93678: add _testinternalcapi.optimize_cfg() and test utils for compiler optimization unit tests (GH-96007)

1 files changed, 93 insertions, 0 deletions

diff --git a/Lib/test/support/bytecode_helper.py b/Lib/test/support/bytecode_helper.py
index 471d4a6..05b5491 100644
--- a/Lib/test/support/bytecode_helper.py
+++ b/Lib/test/support/bytecode_helper.py
@@ -3,6 +3,7 @@
 import unittest
 import dis
 import io
+from _testinternalcapi import optimize_cfg
 
 _UNSPECIFIED = object()
 
@@ -40,3 +41,95 @@ class BytecodeTestCase(unittest.TestCase):
                     msg = '(%s,%r) occurs in bytecode:\n%s'
                     msg = msg % (opname, argval, disassembly)
                     self.fail(msg)
+
+
+class CfgOptimizationTestCase(unittest.TestCase):
+
+    HAS_ARG = set(dis.hasarg)
+    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
+
+    def complete_insts_info(self, insts):
+        # fill in omitted fields in location, and oparg 0 for ops with no arg.
+        instructions = []
+        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
+
+    def normalize_insts(self, insts):
+        """ Map labels to instruction index.
+            Remove labels which are not used as jump targets.
+        """
+        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
+
+        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))
+        return res
+
+    def get_optimized(self, insts, consts):
+        insts = self.complete_insts_info(insts)
+        insts = optimize_cfg(insts, consts)
+        return insts, consts
+
+    def compareInstructions(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.
+
+        self.assertIsInstance(actual_, list)
+        self.assertIsInstance(expected_, list)
+
+        actual = self.normalize_insts(actual_)
+        expected = self.normalize_insts(expected_)
+        self.assertEqual(len(actual), len(expected))
+
+        # compare instructions
+        for act, exp in zip(actual, expected):
+            if isinstance(act, int):
+                self.assertEqual(exp, act)
+                continue
+            self.assertIsInstance(exp, tuple)
+            self.assertIsInstance(act, tuple)
+            # pad exp with -1's (if location info is incomplete)
+            exp += (-1,) * (len(act) - len(exp))
+            self.assertEqual(exp, act)