port comprehension inlining to Python/compile.c also

Summary:
Original comprehension inlining diff from 3.8: D28940584 (03040db)

In D39216342 we ported comprehension inlining only to the py-compiler.
Since then we've discovered that py-compiler is in practice still used only for strict
and static modules, so we should port valuable optimizations to `compile.c` also.
This ports comprehension inlining to `compile.c` as well.

Ensure that we run the comprehension-inliner tests against both py-compiler
and compile.c.

Also requires moving the implementation in py-compiler from
`CinderCodeGenerator` up to the base class `CinderBaseCodeGenerator`, since
side-by-side comparison testing betweeen py-compiler and compile.c uses
`CinderBaseCodeGenerator`.

In 3.8 there was a difference in behavior between py-compiler and compile.c on
the two `nested_diff_scopes` inlining tests. This difference was not caught in 3.8
because the inlining tests only ran against `compile.c`, and the tested case did
not occur in the corpus for side-by-side testing. In porting inlining to py-compiler
in 3.10, we switched the inlining tests to run against py-compiler, and adjusted the
two `nested_diff_scopes` tests accordingly; this resulted in adding
`DELETE_DEREF` opcodes in those two tests that weren't present in those tests
in 3.8. In this diff I remove those opcodes again and adjust the py-compiler
implementation so that they aren't emitted.

In actual fact neither the presence nor absence of those `DELETE_DEREF`
opcodes is correct. It's not safe to inline a comprehension with cell vars at all,
so in the two `nested_diff_scopes` tests we should not be inlining the
comprehension with the lambda at all. I will make this change as a separate
bugfix in 3.8 and port it separately to 3.10.

Reviewed By: tekknolagi

Differential Revision: D40656285

fbshipit-source-id: 405d3fe

Author: Carl Meyer

Include/internal/pycore_symtable.h

@@ -39,6 +39,7 @@ struct symtable {
                                        the symbol table */
     int recursion_depth;            /* current recursion depth */
     int recursion_limit;            /* recursion limit */
+    int st_inline_comprehensions;   /* inline comprehensions? */
 };
 
 typedef struct _symtable_entry {
@@ -72,6 +73,7 @@ typedef struct _symtable_entry {
     int ste_end_col_offset;  /* end offset of first line of block */
     int ste_opt_lineno;      /* lineno of last exec or import * */
     int ste_opt_col_offset;  /* offset of last exec or import * */
+    unsigned int ste_inlined_comprehension; /* comprehension is inlined; symbols merged in parent scope */
     struct symtable *ste_table;
 } PySTEntryObject;
 
@@ -85,6 +87,11 @@ extern struct symtable* _PySymtable_Build(
     struct _mod *mod,
     PyObject *filename,
     PyFutureFeatures *future);
+extern struct symtable* _PySymtable_BuildEx(
+    struct _mod *mod,
+    PyObject *filename,
+    PyFutureFeatures *future,
+    int inline_comprehensions);
 PyAPI_FUNC(PySTEntryObject *) PySymtable_Lookup(struct symtable *, void *);
 
 extern void _PySymtable_Free(struct symtable *);

Lib/compiler/pycodegen.py

@@ -3018,13 +3018,86 @@ def __init__(self, kind, block, exit, unwinding_datum):
         self.unwinding_datum = unwinding_datum
 
 
-# This is identical to the base code generator, except for the fact that CO_SUPPRESS_JIT is emitted.
 class CinderBaseCodeGenerator(CodeGenerator):
+    """
+    Code generator equivalent to `Python/compile.c` in Cinder.
+
+    The base `CodeGenerator` is equivalent to upstream `Python/compile.c`.
+    """
+
     flow_graph = pyassem.PyFlowGraphCinder
+    _SymbolVisitor = symbols.CinderSymbolVisitor
+
+    # TODO(T132400505): Split into smaller methods.
+    def compile_comprehension(
+        self,
+        node: CompNode,
+        name: str,
+        elt: ast.expr,
+        val: ast.expr | None,
+        opcode: str,
+        oparg: object = 0,
+    ) -> None:
+        # fetch the scope that correspond to comprehension
+        scope = self.scopes[node]
+        if scope.inlined:
+            # for inlined comprehension process with current generator
+            gen = self
+        else:
+            gen = self.make_func_codegen(
+                node, self.conjure_arguments([ast.arg(".0", None)]), name, node.lineno
+            )
+        gen.set_lineno(node)
+
+        if opcode:
+            gen.emit(opcode, oparg)
+
+        gen.compile_comprehension_generator(
+            node.generators, 0, 0, elt, val, type(node), not scope.inlined
+        )
+
+        if scope.inlined:
+            # collect list of defs that were introduced by comprehension
+            # note that we need to exclude:
+            # - .0 parameter since it is used
+            # - non-local names (typically named expressions), they are
+            #   defined in enclosing scope and thus should not be deleted
+            to_delete = [
+                v
+                for v in scope.defs
+                if v != ".0"
+                and v not in scope.nonlocals
+                and v not in scope.parent.cells
+            ]
+            # sort names to have deterministic deletion order
+            to_delete.sort()
+            for v in to_delete:
+                self.delName(v)
+            return
+
+        if not isinstance(node, ast.GeneratorExp):
+            gen.emit("RETURN_VALUE")
+
+        gen.finishFunction()
+
+        self._makeClosure(gen, 0)
+
+        # precomputation of outmost iterable
+        self.visit(node.generators[0].iter)
+        if node.generators[0].is_async:
+            self.emit("GET_AITER")
+        else:
+            self.emit("GET_ITER")
+        self.emit("CALL_FUNCTION", 1)
+
+        if gen.scope.coroutine and type(node) is not ast.GeneratorExp:
+            self.emit("GET_AWAITABLE")
+            self.emit("LOAD_CONST", None)
+            self.emit("YIELD_FROM")
 
 
 class CinderCodeGenerator(CinderBaseCodeGenerator):
-    _SymbolVisitor = symbols.CinderSymbolVisitor
+    """Contains some optimizations not (yet) present in Python/compile.c."""
 
     def set_qual_name(self, qualname):
         self._qual_name = qualname
@@ -3106,71 +3179,6 @@ def findFutures(self, node):
                 future_flags |= consts.CO_FUTURE_EAGER_IMPORTS
         return future_flags
 
-    # TODO(T132400505): Split into smaller methods.
-    def compile_comprehension(
-        self,
-        node: CompNode,
-        name: str,
-        elt: ast.expr,
-        val: ast.expr | None,
-        opcode: str,
-        oparg: object = 0,
-    ) -> None:
-        # fetch the scope that correspond to comprehension
-        scope = self.scopes[node]
-        if scope.inlined:
-            # for inlined comprehension process with current generator
-            gen = self
-        else:
-            gen = self.make_func_codegen(
-                node, self.conjure_arguments([ast.arg(".0", None)]), name, node.lineno
-            )
-        gen.set_lineno(node)
-
-        if opcode:
-            gen.emit(opcode, oparg)
-
-        gen.compile_comprehension_generator(
-            node.generators, 0, 0, elt, val, type(node), not scope.inlined
-        )
-
-        if scope.inlined:
-            # collect list of defs that were introduced by comprehension
-            # note that we need to exclude:
-            # - .0 parameter since it is used
-            # - non-local names (typically named expressions), they are
-            #   defined in enclosing scope and thus should not be deleted
-            to_delete = [
-                v
-                for v in scope.defs
-                if v != ".0" and v not in scope.nonlocals and v not in scope.cells
-            ]
-            # sort names to have deterministic deletion order
-            to_delete.sort()
-            for v in to_delete:
-                self.delName(v)
-            return
-
-        if not isinstance(node, ast.GeneratorExp):
-            gen.emit("RETURN_VALUE")
-
-        gen.finishFunction()
-
-        self._makeClosure(gen, 0)
-
-        # precomputation of outmost iterable
-        self.visit(node.generators[0].iter)
-        if node.generators[0].is_async:
-            self.emit("GET_AITER")
-        else:
-            self.emit("GET_ITER")
-        self.emit("CALL_FUNCTION", 1)
-
-        if gen.scope.coroutine and type(node) is not ast.GeneratorExp:
-            self.emit("GET_AWAITABLE")
-            self.emit("LOAD_CONST", None)
-            self.emit("YIELD_FROM")
-
 
 def get_default_generator():
     if "cinder" in sys.version:

Lib/compiler/symbols.py

@@ -777,6 +777,11 @@ def inline_nested_comprehensions(self):
             for u in comp.uses.keys():
                 self.add_use(u)
 
+            # cell vars in comprehension become cells in current scope
+            for c in comp.cells.keys():
+                if c != ".0":
+                    self.cells[c] = 1
+
             # splice children of comprehension into current scope
             # replacing existing entry for 'comp'
             i = self.children.index(comp)

Lib/importlib/_bootstrap_external.py

@@ -357,8 +357,9 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.10b1 3442 (Port Cinder-specific PyCodeObject.co_qualname from 3.8)
 #     Python 3.10b1 3443 (remove primitive enums, migrate PRIMITIVE_(UN)BOX back to opargs)
 #     Python 3.10b1 3444 (optimizations of LOAD_METHOD_SUPER and LOAD_ATTR_SUPER)
-#     Python 3.10b1 3445 (comprehension inliner)
+#     Python 3.10b1 3445 (comprehension inliner in Lib/compiler)
 #     Python 3.10b1 3446 (Set default PyCodeObject.co_qualname missed in 3442)
+#     Python 3.10b1 3447 (comprehension inliner in Python/compile.c)
 
 
 #
@@ -369,7 +370,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3446).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3447).to_bytes(2, 'little') + b'\r\n'
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'

Lib/test/test_compiler/test_cinder.py

@@ -291,6 +291,11 @@ def f(self):
 
 
 class ComprehensionInlinerTests(DisTests):
+    compiler = staticmethod(py_compile)
+
+    def compile(self, code_str):
+        return self.compiler(dedent(code_str), "<string>", "exec")
+
     def test_sync_comp_top(self):
         # ensure module level comprehensions are not inlined
         src = """
@@ -307,7 +312,7 @@ def test_sync_comp_top(self):
              14 LOAD_CONST               2 (None)
              16 RETURN_VALUE
 """
-        co = py_compile(dedent(src), "<string>", "exec")
+        co = self.compile(src)
         self.do_disassembly_test(co, expected)
 
     def test_inline_sync_comp_nested_diff_scopes_1(self):
@@ -325,27 +330,26 @@ def f():
              10 LOAD_DEREF               0 (x)
              12 LIST_APPEND              2
              14 JUMP_ABSOLUTE            3 (to 6)
-        >>   16 DELETE_DEREF             0 (x)
-             18 POP_TOP
-
-  4          20 BUILD_LIST               0
-             22 LOAD_GLOBAL              0 (lst)
-             24 GET_ITER
-        >>   26 FOR_ITER                 8 (to 44)
-             28 STORE_DEREF              0 (x)
-             30 LOAD_CLOSURE             0 (x)
-             32 BUILD_TUPLE              1
-             34 LOAD_CONST               1 (<code object <lambda> at 0x..., file "<string>", line 4>)
-             36 LOAD_CONST               2 ('f.<locals>.<lambda>')
-             38 MAKE_FUNCTION            8 (closure)
-             40 LIST_APPEND              2
-             42 JUMP_ABSOLUTE           13 (to 26)
-        >>   44 POP_TOP
-             46 LOAD_CONST               0 (None)
-             48 RETURN_VALUE
+        >>   16 POP_TOP
+
+  4          18 BUILD_LIST               0
+             20 LOAD_GLOBAL              0 (lst)
+             22 GET_ITER
+        >>   24 FOR_ITER                 8 (to 42)
+             26 STORE_DEREF              0 (x)
+             28 LOAD_CLOSURE             0 (x)
+             30 BUILD_TUPLE              1
+             32 LOAD_CONST               1 (<code object <lambda> at 0x..., file "<string>", line 4>)
+             34 LOAD_CONST               2 ('f.<locals>.<lambda>')
+             36 MAKE_FUNCTION            8 (closure)
+             38 LIST_APPEND              2
+             40 JUMP_ABSOLUTE           12 (to 24)
+        >>   42 POP_TOP
+             44 LOAD_CONST               0 (None)
+             46 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_inline_sync_comp_nested_diff_scopes_2(self):
@@ -377,13 +381,12 @@ def f():
              36 LOAD_DEREF               0 (x)
              38 LIST_APPEND              2
              40 JUMP_ABSOLUTE           16 (to 32)
-        >>   42 DELETE_DEREF             0 (x)
-             44 POP_TOP
-             46 LOAD_CONST               0 (None)
-             48 RETURN_VALUE
+        >>   42 POP_TOP
+             44 LOAD_CONST               0 (None)
+             46 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_inline_sync_comp_nested_comprehensions(self):
@@ -414,7 +417,7 @@ def f():
              38 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_inline_sync_comp_named_expr_1(self):
@@ -441,7 +444,7 @@ def f():
              30 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_inline_async_comp_free_var1(self):
@@ -490,7 +493,7 @@ async def f(lst):
              68 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_comprehension_inlining_name_conflict_with_implicit_global(self):
@@ -520,7 +523,7 @@ def g():
 """
 
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_use_param_1(self):
@@ -554,7 +557,7 @@ def f(self, name, data, files=(), dirs=()):
              44 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
 
     def test_inline_comp_global1(self):
@@ -586,5 +589,9 @@ def f():
              34 RETURN_VALUE
 """
         g = {}
-        exec(py_compile(dedent(src), "<string>", "exec"), g)
+        exec(self.compile(src), g)
         self.do_disassembly_test(g["f"], expected)
+
+
+class ComprehensionInlinerBuiltinCompilerTests(ComprehensionInlinerTests):
+    compiler = staticmethod(compile)

Lib/test/test_compiler/test_py310.py

@@ -58,20 +58,3 @@ def f():
                 return x
         """
         self._check(codestr)
-
-    def test_no_nested_async_comprehension(self):
-        codestr = """
-            async def foo(a):
-                return {k: [y for y in k if await bar(y)] for k in a}
-        """
-
-        # The base code generator matches 3.10 upstream and thus has this
-        # restriction, but the restriction has been lifted in 3.11
-        # (see https://github.com/python/cpython/pull/6766), so we also lift
-        # it in CinderCodeGenerator.
-        self._check_error(
-            codestr,
-            "asynchronous comprehension outside of an asynchronous function",
-            generator=BaseCodeGenerator,
-        )
-        self.compile(codestr, generator=CinderCodeGenerator)