gh-126835: Fold unary & binary complex constant expressions during code generation.

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)

Include/internal/pycore_global_strings.h

@@ -613,6 +613,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(opener)
         STRUCT_FOR_ID(operation)
         STRUCT_FOR_ID(optimize)
+        STRUCT_FOR_ID(optimize_ast)
         STRUCT_FOR_ID(options)
         STRUCT_FOR_ID(order)
         STRUCT_FOR_ID(origin)

Lib/test/support/bytecode_helper.py

@@ -89,6 +89,13 @@ def assertInstructionsMatch(self, actual_seq, expected):
             idx = max([p[0] for p in enumerate(exp) if p[1] != -1])
             self.assertEqual(exp[:idx], act[:idx])
 
+    def assertNotInInstructionSequence(self, seq, expected_opcode):
+        self.assertIn(expected_opcode, dis.opmap)
+        for instr in seq.get_instructions():
+            opcode, *_ = instr
+            if dis.opmap[expected_opcode] == opcode:
+                self.fail(f"{expected_opcode} appears in instructions sequence.")
+
     def resolveAndRemoveLabels(self, insts):
         idx = 0
         res = []
@@ -138,8 +145,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,57 @@ 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
+
+                case {-0: 0}:                pass  # match unary const int
+                case {-0.1: 0}:              pass  # match unary const float
+                case {-0j: 0}:               pass  # match unary const complex
+                case {1 + 2j: 0}:            pass  # match const int + const complex
+                case {1 - 2j: 0}:            pass  # match const int - const complex
+                case {1.1 + 2.1j: 0}:        pass  # match const float + const complex
+                case {1.1 - 2.1j: 0}:        pass  # match const float - const complex
+                case {-0 + 1j: 0}:           pass  # match unary const int + complex
+                case {-0 - 1j: 0}:           pass  # match unary const int - complex
+                case {-0.1 + 1.1j: 0}:       pass  # match unary const float + complex
+                case {-0.1 - 1.1j: 0}:       pass  # match unary const float - complex
+        """)
+        import ast
+        a = ast.parse(snippet, "my_file.py", "exec")
+        code = self.generate_code(a, optimize_ast=False)
+        self.assertNotInInstructionSequence(code, 'BINARY_OP')
+        self.assertNotInInstructionSequence(code, 'UNARY_NEGATIVE')

Modules/_testinternalcapi.c

@@ -711,18 +711,21 @@ _testinternalcapi.compiler_codegen -> object
   filename: object
   optimize: int
   compile_mode: int = 0
+  optimize_ast: int = 1
 
 Apply compiler code generation to an AST.
 [clinic start generated code]*/
 
 static PyObject *
 _testinternalcapi_compiler_codegen_impl(PyObject *module, PyObject *ast,
                                         PyObject *filename, int optimize,
-                                        int compile_mode)
-/*[clinic end generated code: output=40a68f6e13951cc8 input=a0e00784f1517cd7]*/
+                                        int compile_mode, int optimize_ast)
+/*[clinic end generated code: output=b4bf87f1213effd1 input=ece7e3ca206d738a]*/
 {
-    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);
 }