fold unary & binary complex constant expressions in codegen

Author: WolframAlph

Include/cpython/compile.h

@@ -20,6 +20,7 @@
 #define PyCF_ALLOW_TOP_LEVEL_AWAIT 0x2000
 #define PyCF_ALLOW_INCOMPLETE_INPUT 0x4000
 #define PyCF_OPTIMIZED_AST (0x8000 | PyCF_ONLY_AST)
+#define PyCF_DONT_OPTIMIZE_AST 0x10000
 #define PyCF_COMPILE_MASK (PyCF_ONLY_AST | PyCF_ALLOW_TOP_LEVEL_AWAIT | \
                            PyCF_TYPE_COMMENTS | PyCF_DONT_IMPLY_DEDENT | \
                            PyCF_ALLOW_INCOMPLETE_INPUT | PyCF_OPTIMIZED_AST)

Lib/test/support/bytecode_helper.py

@@ -138,8 +138,8 @@ def check_instructions(self, insts):
 @unittest.skipIf(_testinternalcapi is None, "requires _testinternalcapi")
 class CodegenTestCase(CompilationStepTestCase):
 
-    def generate_code(self, ast):
-        insts, _ = _testinternalcapi.compiler_codegen(ast, "my_file.py", 0)
+    def generate_code(self, ast, optimize_ast=True):
+        insts, _ = _testinternalcapi.compiler_codegen(ast, "my_file.py", 0, 0, optimize_ast)
         return insts
 
 

Lib/test/test_compiler_codegen.py

@@ -15,10 +15,10 @@ def assertInstructionsMatch_recursive(self, insts, expected_insts):
         for n_insts, n_expected in zip(insts.get_nested(), expected_nested):
             self.assertInstructionsMatch_recursive(n_insts, n_expected)
 
-    def codegen_test(self, snippet, expected_insts):
+    def codegen_test(self, snippet, expected_insts, optimize_ast=True):
         import ast
         a = ast.parse(snippet, "my_file.py", "exec")
-        insts = self.generate_code(a)
+        insts = self.generate_code(a, optimize_ast=optimize_ast)
         self.assertInstructionsMatch_recursive(insts, expected_insts)
 
     def test_if_expression(self):
@@ -157,3 +157,41 @@ def test_syntax_error__return_not_in_function(self):
         self.assertIsNone(cm.exception.text)
         self.assertEqual(cm.exception.offset, 1)
         self.assertEqual(cm.exception.end_offset, 10)
+
+    def test_dont_optimize_ast_before_codegen(self):
+        snippet = "1+2"
+        unoptimized = [
+            ('RESUME', 0, 0),
+            ('LOAD_SMALL_INT', 1, 0),
+            ('LOAD_SMALL_INT', 2, 0),
+            ('BINARY_OP', 0, 0),
+            ('POP_TOP', None, 0),
+            ('LOAD_CONST', 0, 0),
+            ('RETURN_VALUE', None, 0),
+        ]
+        self.codegen_test(snippet, unoptimized, optimize_ast=False)
+
+        optimized = [
+            ('RESUME', 0, 0),
+            ('NOP', None, 0),
+            ('LOAD_CONST', 0, 0),
+            ('RETURN_VALUE', None, 0),
+        ]
+        self.codegen_test(snippet, optimized, optimize_ast=True)
+
+    def test_match_case_fold_codegen(self):
+        snippet = textwrap.dedent("""
+            match 0:
+                case -0:                pass  # match unary const int
+                case -0.1:              pass  # match unary const float
+                case -0j:               pass  # match unary const complex
+                case 1 + 2j:            pass  # match const int + const complex
+                case 1 - 2j:            pass  # match const int - const complex
+                case 1.1 + 2.1j:        pass  # match const float + const complex
+                case 1.1 - 2.1j:        pass  # match const float - const complex
+                case -0 + 1j:           pass  # match unary const int + complex
+                case -0 - 1j:           pass  # match unary const int - complex
+                case -0.1 + 1.1j:       pass  # match unary const float + complex
+                case -0.1 - 1.1j:       pass  # match unary const float - complex
+        """)
+        self.codegen_test(snippet, [], optimize_ast=False)

Modules/_testinternalcapi.c

@@ -718,11 +718,13 @@ Apply compiler code generation to an AST.
 static PyObject *
 _testinternalcapi_compiler_codegen_impl(PyObject *module, PyObject *ast,
                                         PyObject *filename, int optimize,
-                                        int compile_mode)
+                                        int compile_mode, int optimize_ast)
 /*[clinic end generated code: output=40a68f6e13951cc8 input=a0e00784f1517cd7]*/
 {
-    PyCompilerFlags *flags = NULL;
-    return _PyCompile_CodeGen(ast, filename, flags, optimize, compile_mode);
+    PyCompilerFlags flags = _PyCompilerFlags_INIT;
+    if (!optimize_ast)
+        flags.cf_flags = PyCF_DONT_OPTIMIZE_AST;
+    return _PyCompile_CodeGen(ast, filename, &flags, optimize, compile_mode);
 }
 
 

Modules/clinic/_testinternalcapi.c.h

@@ -5355,9 +5355,114 @@ codegen_slice(compiler *c, expr_ty s)
 #define WILDCARD_STAR_CHECK(N) \
     ((N)->kind == MatchStar_kind && !(N)->v.MatchStar.name)
 
-// Limit permitted subexpressions, even if the parser & AST validator let them through
-#define MATCH_VALUE_EXPR(N) \
-    ((N)->kind == Constant_kind || (N)->kind == Attribute_kind)
+static bool
+is_unary_or_complex_expr(expr_ty e)
+{
+    if (e->kind != UnaryOp_kind) {
+        return false;
+    }
+    if (e->v.UnaryOp.op != USub) {
+        return false;
+    }
+    if (e->v.UnaryOp.operand->kind != Constant_kind) {
+        return false;
+    }
+    PyObject *constant = e->v.UnaryOp.operand->v.Constant.value;
+    return PyLong_CheckExact(constant) || PyFloat_CheckExact(constant) || PyComplex_CheckExact(constant);
+}
+
+static bool
+is_complex_binop_expr(expr_ty e)
+{
+    if (e->kind != BinOp_kind) {
+        return false;
+    }
+    if (e->v.BinOp.op != Add && e->v.BinOp.op != Sub) {
+        return false;
+    }
+    if (e->v.BinOp.right->kind != Constant_kind) {
+        return false;
+    }
+    if (e->v.BinOp.left->kind != Constant_kind && e->v.BinOp.left->kind != UnaryOp_kind) {
+        return false;
+    }
+    PyObject *leftconst;
+    if (e->v.BinOp.left->kind == UnaryOp_kind) {
+        if (e->v.BinOp.left->v.UnaryOp.operand->kind != Constant_kind) {
+            return false;
+        }
+        if (e->v.BinOp.left->v.UnaryOp.op != USub) {
+            return false;
+        }
+        leftconst = e->v.BinOp.left->v.UnaryOp.operand->v.Constant.value;
+    }
+    else {
+        leftconst = e->v.BinOp.left->v.Constant.value;
+    }
+    PyObject *rightconst = e->v.BinOp.right->v.Constant.value;
+    return (PyLong_CheckExact(leftconst) || PyFloat_CheckExact(leftconst)) && PyComplex_CheckExact(rightconst);
+}
+
+static void
+fold_node(expr_ty node, PyObject *folded)
+{
+    assert(node->kind != Constant_kind);
+    node->kind = Constant_kind;
+    node->v.Constant.kind = NULL;
+    node->v.Constant.value = folded;
+}
+
+static int
+fold_unary_or_complex_expr(expr_ty e)
+{
+    assert(e->kind == UnaryOp_kind);
+    assert(e->v.UnaryOp.op == USub);
+    assert(e->v.UnaryOp.operand->kind == Constant_kind);
+    PyObject *operand = e->v.UnaryOp.operand->v.Constant.value;
+    assert(PyLong_CheckExact(operand) || PyFloat_CheckExact(operand) || PyComplex_CheckExact(operand));
+    PyObject* folded = PyNumber_Negative(operand);
+    if (folded == NULL) {
+        return ERROR;
+    }
+    fold_node(e, folded);
+    return SUCCESS;
+}
+
+static int
+fold_binary_complex_expr(expr_ty e)
+{
+    assert(e->kind == BinOp_kind);
+    assert(e->v.BinOp.right->kind == Constant_kind);
+    assert(e->v.BinOp.left->kind == UnaryOp_kind || e->v.BinOp.left->kind == Constant_kind);
+    if (e->v.BinOp.left->kind == UnaryOp_kind) {
+        RETURN_IF_ERROR(fold_unary_or_complex_expr(e->v.BinOp.left));
+    }
+    assert(e->v.BinOp.left->kind == Constant_kind);
+    operator_ty op = e->v.BinOp.op;
+    PyObject *left = e->v.BinOp.left->v.Constant.value;
+    PyObject *right = e->v.BinOp.right->v.Constant.value;
+    assert(op == Add || op == Sub);
+    assert(PyLong_CheckExact(left) || PyFloat_CheckExact(left));
+    assert(PyComplex_CheckExact(right));
+    PyObject *folded = op == Add ? PyNumber_Add(left, right) : PyNumber_Subtract(left, right);
+    if (folded == NULL) {
+        return ERROR;
+    }
+    fold_node(e, folded);
+    return SUCCESS;
+}
+
+static int
+try_fold_unary_or_binary_complex_const_expr(expr_ty key)
+{
+    if (is_unary_or_complex_expr(key)) {
+        return fold_unary_or_complex_expr(key);
+    }
+    if (is_complex_binop_expr(key)) {
+        return fold_binary_complex_expr(key);
+    }
+    return SUCCESS;
+}
 
 // Allocate or resize pc->fail_pop to allow for n items to be popped on failure.
 static int
@@ -5688,7 +5793,7 @@ codegen_pattern_mapping_key(compiler *c, PyObject *seen, pattern_ty p, Py_ssize_
         location loc = LOC((pattern_ty) asdl_seq_GET(patterns, i));
         return _PyCompile_Error(c, loc, e);
     }
-
+    RETURN_IF_ERROR(try_fold_unary_or_binary_complex_const_expr(key));
     if (key->kind == Constant_kind) {
         int in_seen = PySet_Contains(seen, key->v.Constant.value);
         RETURN_IF_ERROR(in_seen);
@@ -6022,7 +6127,8 @@ codegen_pattern_value(compiler *c, pattern_ty p, pattern_context *pc)
 {
     assert(p->kind == MatchValue_kind);
     expr_ty value = p->v.MatchValue.value;
-    if (!MATCH_VALUE_EXPR(value)) {
+    RETURN_IF_ERROR(try_fold_unary_or_binary_complex_const_expr(value));
+    if (value->kind != Constant_kind && value->kind != Attribute_kind) {
         const char *e = "patterns may only match literals and attribute lookups";
         return _PyCompile_Error(c, LOC(p), e);
     }

Python/codegen.c

@@ -126,7 +126,8 @@ compiler_setup(compiler *c, mod_ty mod, PyObject *filename,
     c->c_optimize = (optimize == -1) ? _Py_GetConfig()->optimization_level : optimize;
     c->c_save_nested_seqs = false;
 
-    if (!_PyAST_Optimize(mod, arena, c->c_optimize, merged)) {
+    int ast_opt = !(flags->cf_flags & PyCF_DONT_OPTIMIZE_AST);
+    if (ast_opt && !_PyAST_Optimize(mod, arena, c->c_optimize, merged)) {
         return ERROR;
     }
     c->c_st = _PySymtable_Build(mod, filename, &c->c_future);