Add more tests

Author: mpage

Include/internal/pycore_flowgraph.h

@@ -41,6 +41,9 @@ PyAPI_FUNC(PyObject*) _PyCompile_OptimizeCfg(
         PyObject *consts,
         int nlocals);
 
+// Export for '_testinternalcapi' shared extension
+PyAPI_FUNC(PyObject*) _PyCompile_OptimizeLoadFast(PyObject *instructions);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_peepholer.py

@@ -4,9 +4,14 @@
 import sys
 import textwrap
 import unittest
+try:
+    import _testinternalcapi
+except ImportError:
+    _testinternalcapi = None
 
 from test import support
-from test.support.bytecode_helper import BytecodeTestCase, CfgOptimizationTestCase
+from test.support.bytecode_helper import (
+    BytecodeTestCase, CfgOptimizationTestCase, CompilationStepTestCase)
 
 
 def compile_pattern_with_fast_locals(pattern):
@@ -2353,5 +2358,97 @@ def test_list_to_tuple_get_iter_is_safe(self):
         self.assertEqual(items, [])
 
 
+@unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
+class OptimizeLoadFastTestCase(CompilationStepTestCase):
+    def check(self, insts, expected_insts):
+        self.check_instructions(insts)
+        self.check_instructions(expected_insts)
+        seq = self.seq_from_insts(insts)
+        opt_insts = _testinternalcapi.optimize_load_fast(seq)
+        expected_insts = self.seq_from_insts(expected_insts).get_instructions()
+        self.assertInstructionsMatch(opt_insts, expected_insts)
+
+    def test_optimized(self):
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_FAST", 1, 2),
+            ("BINARY_OP", 2, 3),
+        ]
+        expected = [
+            ("LOAD_FAST_BORROW", 0, 1),
+            ("LOAD_FAST_BORROW", 1, 2),
+            ("BINARY_OP", 2, 3),
+        ]
+        self.check(insts, expected)
+
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_CONST", 1, 2),
+            ("SWAP", 2, 3),
+            ("POP_TOP", None, 4),
+        ]
+        expected = [
+            ("LOAD_FAST_BORROW", 0, 1),
+            ("LOAD_CONST", 1, 2),
+            ("SWAP", 2, 3),
+            ("POP_TOP", None, 4),
+        ]
+        self.check(insts, expected)
+
+    def test_unoptimized_if_unconsumed(self):
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_FAST", 1, 2),
+            ("POP_TOP", None, 3),
+        ]
+        expected = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_FAST_BORROW", 1, 2),
+            ("POP_TOP", None, 3),
+        ]
+        self.check(insts, expected)
+
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("COPY", 1, 2),
+            ("POP_TOP", None, 3),
+        ]
+        self.check(insts, insts)
+
+    def test_unoptimized_if_support_killed(self):
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_CONST", 0, 2),
+            ("STORE_FAST", 0, 3),
+            ("POP_TOP", None, 4),
+        ]
+        self.check(insts, insts)
+
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_CONST", 0, 2),
+            ("LOAD_CONST", 0, 3),
+            ("STORE_FAST_STORE_FAST", 0 << 4 | 1, 4),
+            ("POP_TOP", None, 5),
+        ]
+        self.check(insts, insts)
+
+    def test_unoptimized_if_aliased(self):
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("STORE_FAST", 1, 2),
+        ]
+        self.check(insts, insts)
+
+        insts = [
+            ("LOAD_FAST", 0, 1),
+            ("LOAD_CONST", 0, 3),
+            ("STORE_FAST_STORE_FAST", 0 << 4 | 1, 4),
+        ]
+        self.check(insts, insts)
+
+
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_testinternalcapi.c

@@ -748,6 +748,23 @@ _testinternalcapi_optimize_cfg_impl(PyObject *module, PyObject *instructions,
     return _PyCompile_OptimizeCfg(instructions, consts, nlocals);
 }
 
+/*[clinic input]
+
+_testinternalcapi.optimize_load_fast -> object
+
+  instructions: object
+
+Optimize LOAD_FAST{_LOAD_FAST} instructions.
+[clinic start generated code]*/
+
+static PyObject *
+_testinternalcapi_optimize_load_fast_impl(PyObject *module,
+                                          PyObject *instructions)
+/*[clinic end generated code: output=6f975349c976d017 input=c59f3eac68308c01]*/
+{
+    return _PyCompile_OptimizeLoadFast(instructions);
+}
+
 static int
 get_nonnegative_int_from_dict(PyObject *dict, const char *key) {
     PyObject *obj = PyDict_GetItemString(dict, key);
@@ -2023,6 +2040,7 @@ static PyMethodDef module_functions[] = {
     _TESTINTERNALCAPI_NEW_INSTRUCTION_SEQUENCE_METHODDEF
     _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF
     _TESTINTERNALCAPI_OPTIMIZE_CFG_METHODDEF
+    _TESTINTERNALCAPI_OPTIMIZE_LOAD_FAST_METHODDEF
     _TESTINTERNALCAPI_ASSEMBLE_CODE_OBJECT_METHODDEF
     {"get_interp_settings", get_interp_settings, METH_VARARGS, NULL},
     {"clear_extension", clear_extension, METH_VARARGS, NULL},

Modules/clinic/_testinternalcapi.c.h

@@ -2572,12 +2572,15 @@ load_fast_push_block(basicblock ***sp, basicblock *target, int start_depth)
  * non-violating LOAD_FAST{_LOAD_FAST} can be optimized.
  */
 static int
-optimize_load_fast(cfg_builder *g)
+optimize_load_fast(cfg_builder *g, bool compute_stackdepth)
 {
     int status;
     ref_stack refs = {0};
     int max_instrs = 0;
     basicblock *entryblock = g->g_entryblock;
+    if (compute_stackdepth) {
+        calculate_stackdepth(g);
+    }
     for (basicblock *b = entryblock; b != NULL; b = b->b_next) {
         max_instrs = Py_MAX(max_instrs, b->b_iused);
     }
@@ -2620,6 +2623,7 @@ optimize_load_fast(cfg_builder *g)
             assert(opcode != EXTENDED_ARG);
             switch (opcode) {
                 case COPY: {
+                    assert(oparg > 0);
                     Py_ssize_t idx = refs.size - oparg;
                     ref r = ref_stack_at(&refs, idx);
                     if (ref_stack_push(&refs, r) < 0) {
@@ -2687,6 +2691,7 @@ optimize_load_fast(cfg_builder *g)
                 }
 
                 case SWAP: {
+                    assert(oparg >= 2);
                     ref_stack_swap_top(&refs, oparg);
                     break;
                 }
@@ -3711,7 +3716,7 @@ _PyCfg_OptimizedCfgToInstructionSequence(cfg_builder *g,
     /* Can't modify the bytecode after inserting instructions that produce
      * borrowed references.
      */
-    RETURN_IF_ERROR(optimize_load_fast(g));
+    RETURN_IF_ERROR(optimize_load_fast(g, /* compute_stackdepth */ false));
 
     /* Can't modify the bytecode after computing jump offsets. */
     if (_PyCfg_ToInstructionSequence(g, seq) < 0) {
@@ -3802,3 +3807,26 @@ _PyCompile_OptimizeCfg(PyObject *seq, PyObject *consts, int nlocals)
     _PyCfgBuilder_Free(g);
     return res;
 }
+
+PyObject *
+_PyCompile_OptimizeLoadFast(PyObject *seq)
+{
+    if (!_PyInstructionSequence_Check(seq)) {
+        PyErr_SetString(PyExc_ValueError, "expected an instruction sequence");
+        return NULL;
+    }
+
+    cfg_builder *g = _PyCfg_FromInstructionSequence((_PyInstructionSequence*)seq);
+    if (g == NULL) {
+        return NULL;
+    }
+
+    if (optimize_load_fast(g, /* compute_stackdepth */ true) != SUCCESS) {
+        _PyCfgBuilder_Free(g);
+        return NULL;
+    }
+
+    PyObject *res = cfg_to_instruction_sequence(g);
+    _PyCfgBuilder_Free(g);
+    return res;
+}