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

iritkatriel · web-flow · commit 420f39f457a9 · 2022-08-24T11:02:53.000+01:00
diff --git a/Include/internal/pycore_compile.h b/Include/internal/pycore_compile.h
@@ -38,6 +38,11 @@ extern int _PyAST_Optimize(
     struct _arena *arena,
     _PyASTOptimizeState *state);
 
+/* Access compiler internals for unit testing */
+PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
+        PyObject *instructions,
+        PyObject *consts);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -298,6 +298,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(code)
         STRUCT_FOR_ID(command)
         STRUCT_FOR_ID(comment_factory)
+        STRUCT_FOR_ID(consts)
         STRUCT_FOR_ID(context)
         STRUCT_FOR_ID(cookie)
         STRUCT_FOR_ID(copy)
@@ -407,6 +408,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(input)
         STRUCT_FOR_ID(insert_comments)
         STRUCT_FOR_ID(insert_pis)
+        STRUCT_FOR_ID(instructions)
         STRUCT_FOR_ID(intern)
         STRUCT_FOR_ID(intersection)
         STRUCT_FOR_ID(isatty)
diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
diff --git 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 @@ def assertNotInBytecode(self, x, opname, argval=_UNSPECIFIED):
                     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)
diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
@@ -4,7 +4,7 @@
 import textwrap
 import unittest
 
-from test.support.bytecode_helper import BytecodeTestCase
+from test.support.bytecode_helper import BytecodeTestCase, CfgOptimizationTestCase
 
 
 def compile_pattern_with_fast_locals(pattern):
@@ -864,5 +864,81 @@ def trace(frame, event, arg):
         self.assertNotInBytecode(f, "LOAD_FAST_CHECK")
 
 
+class DirectiCfgOptimizerTests(CfgOptimizationTestCase):
+
+    def cfg_optimization_test(self, insts, expected_insts,
+                              consts=None, expected_consts=None):
+        if expected_consts is None:
+            expected_consts = consts
+        opt_insts, opt_consts = self.get_optimized(insts, consts)
+        self.compareInstructions(opt_insts, expected_insts)
+        self.assertEqual(opt_consts, expected_consts)
+
+    def test_conditional_jump_forward_non_const_condition(self):
+        insts = [
+            ('LOAD_NAME', 1, 11),
+            ('POP_JUMP_IF_TRUE', lbl := self.Label(), 12),
+            ('LOAD_CONST', 2, 13),
+            lbl,
+            ('LOAD_CONST', 3, 14),
+        ]
+        expected = [
+            ('LOAD_NAME', '1', 11),
+            ('POP_JUMP_IF_TRUE', lbl := self.Label(), 12),
+            ('LOAD_CONST', '2', 13),
+            lbl,
+            ('LOAD_CONST', '3', 14)
+        ]
+        self.cfg_optimization_test(insts, expected, consts=list(range(5)))
+
+    def test_conditional_jump_forward_const_condition(self):
+        # The unreachable branch of the jump is removed
+
+        insts = [
+            ('LOAD_CONST', 3, 11),
+            ('POP_JUMP_IF_TRUE', lbl := self.Label(), 12),
+            ('LOAD_CONST', 2, 13),
+            lbl,
+            ('LOAD_CONST', 3, 14),
+        ]
+        expected = [
+            ('NOP', None, 11),
+            ('JUMP', lbl := self.Label(), 12),
+            lbl,
+            ('LOAD_CONST', '3', 14)
+        ]
+        self.cfg_optimization_test(insts, expected, consts=list(range(5)))
+
+    def test_conditional_jump_backward_non_const_condition(self):
+        insts = [
+            lbl1 := self.Label(),
+            ('LOAD_NAME', 1, 11),
+            ('POP_JUMP_IF_TRUE', lbl1, 12),
+            ('LOAD_CONST', 2, 13),
+        ]
+        expected = [
+            lbl := self.Label(),
+            ('LOAD_NAME', '1', 11),
+            ('POP_JUMP_IF_TRUE', lbl, 12),
+            ('LOAD_CONST', '2', 13)
+        ]
+        self.cfg_optimization_test(insts, expected, consts=list(range(5)))
+
+    def test_conditional_jump_backward_const_condition(self):
+        # The unreachable branch of the jump is removed
+        insts = [
+            lbl1 := self.Label(),
+            ('LOAD_CONST', 1, 11),
+            ('POP_JUMP_IF_TRUE', lbl1, 12),
+            ('LOAD_CONST', 2, 13),
+        ]
+        expected = [
+            lbl := self.Label(),
+            ('NOP', None, 11),
+            ('JUMP', lbl, 12)
+        ]
+        self.cfg_optimization_test(insts, expected, consts=list(range(5)))
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-08-15-20-52-41.gh-issue-93678.X7GuIJ.rst b/Misc/NEWS.d/next/Core and Builtins/2022-08-15-20-52-41.gh-issue-93678.X7GuIJ.rst
@@ -0,0 +1 @@
+Added test a harness for direct unit tests of the compiler's optimization stage. The ``_testinternalcapi.optimize_cfg()`` function runs the optimiser on a sequence of instructions. The ``CfgOptimizationTestCase`` class in ``test.support`` has utilities for invoking the optimizer and checking the output.
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -14,6 +14,7 @@
 #include "Python.h"
 #include "pycore_atomic_funcs.h" // _Py_atomic_int_get()
 #include "pycore_bitutils.h"     // _Py_bswap32()
+#include "pycore_compile.h"      // _PyCompile_OptimizeCfg()
 #include "pycore_fileutils.h"    // _Py_normpath
 #include "pycore_frame.h"        // _PyInterpreterFrame
 #include "pycore_gc.h"           // PyGC_Head
@@ -25,7 +26,12 @@
 #include "pycore_pystate.h"      // _PyThreadState_GET()
 #include "osdefs.h"              // MAXPATHLEN
 
+#include "clinic/_testinternalcapi.c.h"
 
+/*[clinic input]
+module _testinternalcapi
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=7bb583d8c9eb9a78]*/
 static PyObject *
 get_configs(PyObject *self, PyObject *Py_UNUSED(args))
 {
@@ -525,6 +531,25 @@ set_eval_frame_record(PyObject *self, PyObject *list)
 }
 
 
+/*[clinic input]
+
+_testinternalcapi.optimize_cfg -> object
+
+  instructions: object
+  consts: object
+
+Apply compiler optimizations to an instruction list.
+[clinic start generated code]*/
+
+static PyObject *
+_testinternalcapi_optimize_cfg_impl(PyObject *module, PyObject *instructions,
+                                    PyObject *consts)
+/*[clinic end generated code: output=5412aeafca683c8b input=7e8a3de86ebdd0f9]*/
+{
+    return _PyCompile_OptimizeCfg(instructions, consts);
+}
+
+
 static PyMethodDef TestMethods[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
@@ -543,6 +568,7 @@ static PyMethodDef TestMethods[] = {
     {"DecodeLocaleEx", decode_locale_ex, METH_VARARGS},
     {"set_eval_frame_default", set_eval_frame_default, METH_NOARGS, NULL},
     {"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
+    _TESTINTERNALCAPI_OPTIMIZE_CFG_METHODDEF
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/clinic/_testinternalcapi.c.h b/Modules/clinic/_testinternalcapi.c.h
diff --git a/Python/compile.c b/Python/compile.c

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Added test a harness for direct unit tests of the compiler's optimization stage. The ``_testinternalcapi.optimize_cfg()`` function runs the optimiser on a sequence of instructions. The ``CfgOptimizationTestCase`` class in ``test.support`` has utilities for invoking the optimizer and checking the output.