gh-97933: add opcode for more efficient comprehension execution

Doc/library/dis.rst

@@ -1300,6 +1300,20 @@ iterations of the loop.
    * 2: ``raise STACK[-2] from STACK[-1]`` (raise exception instance or type at
      ``STACK[-2]`` with ``__cause__`` set to ``STACK[-1]``)
 
+.. opcode:: COMPREHENSION (flag)
+
+   Calls a comprehension code object, without creating and throwing away a
+   single-use function object. ``flag`` must be either ``0`` or ``1``, the
+   latter indicating the comprehension has free variables and a closure tuple
+   will be on the stack.
+
+   The stack should contain, from bottom to top:
+
+   * a tuple containing cells for free variables, if ``flag`` is set
+   * the code object for the comprehension
+   * the single "argument" to the comprehension (the iterated object)
+
+   .. versionadded:: 3.12
 
 .. opcode:: CALL (argc)
 

Include/internal/pycore_frame.h

@@ -53,6 +53,9 @@ typedef struct _PyInterpreterFrame {
     PyObject *f_globals; /* Borrowed reference. Only valid if not on C stack */
     PyObject *f_builtins; /* Borrowed reference. Only valid if not on C stack */
     PyObject *f_locals; /* Strong reference, may be NULL. Only valid if not on C stack */
+    // For comprehensions, f_closure and f_code may not match func_closure and
+    // func_code from f_funcobj above; f_funcobj will be the calling function.
+    PyObject *f_closure; /* Strong reference, may be NULL. Only valid if not on C stack */
     PyFrameObject *frame_obj; /* Strong reference, may be NULL. Only valid if not on C stack */
     // NOTE: This is not necessarily the last instruction started in the given
     // frame. Rather, it is the code unit *prior to* the *next* instruction. For
@@ -110,12 +113,14 @@ void _PyFrame_Copy(_PyInterpreterFrame *src, _PyInterpreterFrame *dest);
 static inline void
 _PyFrame_Initialize(
     _PyInterpreterFrame *frame, PyFunctionObject *func,
-    PyObject *locals, PyCodeObject *code, int null_locals_from)
+    PyObject *locals, PyCodeObject *code, PyObject *closure,
+    int null_locals_from)
 {
     frame->f_funcobj = (PyObject *)func;
     frame->f_code = (PyCodeObject *)Py_NewRef(code);
     frame->f_builtins = func->func_builtins;
     frame->f_globals = func->func_globals;
+    frame->f_closure = Py_XNewRef(closure);
     frame->f_locals = locals;
     frame->stacktop = code->co_nlocalsplus;
     frame->frame_obj = NULL;
@@ -248,7 +253,7 @@ _PyFrame_PushUnchecked(PyThreadState *tstate, PyFunctionObject *func, int null_l
     _PyInterpreterFrame *new_frame = (_PyInterpreterFrame *)tstate->datastack_top;
     tstate->datastack_top += code->co_framesize;
     assert(tstate->datastack_top < tstate->datastack_limit);
-    _PyFrame_Initialize(new_frame, func, NULL, code, null_locals_from);
+    _PyFrame_Initialize(new_frame, func, NULL, code, func->func_closure, null_locals_from);
     return new_frame;
 }
 

Lib/importlib/_bootstrap_external.py

@@ -434,6 +434,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a5 3518 (Add RETURN_CONST instruction)
 #     Python 3.12a5 3519 (Modify SEND instruction)
 #     Python 3.12a5 3520 (Remove PREP_RERAISE_STAR, add CALL_INTRINSIC_2)
+#     Python 3.12a5 3521 (Add COMPREHENSION instruction)
 
 #     Python 3.13 will start with 3550
 
@@ -446,7 +447,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 = (3520).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3521).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Lib/opcode.py

@@ -219,6 +219,7 @@ def pseudo_op(name, op, real_ops):
 def_op('DICT_MERGE', 164)
 def_op('DICT_UPDATE', 165)
 
+def_op('COMPREHENSION', 170)
 def_op('CALL', 171)
 def_op('KW_NAMES', 172)
 hasconst.append(172)

Lib/test/test_dis.py

@@ -163,10 +163,9 @@ def bug1333982(x=[]):
 
 %3d        LOAD_ASSERTION_ERROR
            LOAD_CONST               1 (<code object <listcomp> at 0x..., file "%s", line %d>)
-           MAKE_FUNCTION            0
            LOAD_FAST                0 (x)
            GET_ITER
-           CALL                     0
+           COMPREHENSION            0
 
 %3d        LOAD_CONST               2 (1)
 
@@ -661,10 +660,9 @@ def foo(x):
 %3d        LOAD_CLOSURE             0 (x)
            BUILD_TUPLE              1
            LOAD_CONST               1 (<code object <listcomp> at 0x..., file "%s", line %d>)
-           MAKE_FUNCTION            8 (closure)
            LOAD_DEREF               1 (y)
            GET_ITER
-           CALL                     0
+           COMPREHENSION            1
            RETURN_VALUE
 """ % (dis_nested_0,
        __file__,

Lib/test/test_sys.py

@@ -1443,7 +1443,7 @@ class C(object): pass
         def func():
             return sys._getframe()
         x = func()
-        check(x, size('3Pi3c7P2ic??2P'))
+        check(x, size('3Pi3c8P2ic??2P'))
         # function
         def func(): pass
         check(func, size('14Pi'))
@@ -1460,7 +1460,7 @@ def bar(cls):
             check(bar, size('PP'))
         # generator
         def get_gen(): yield 1
-        check(get_gen(), size('P2P4P4c7P2ic??2P'))
+        check(get_gen(), size('P2P4P4c8P2ic??2P'))
         # iterator
         check(iter('abc'), size('lP'))
         # callable-iterator

Misc/NEWS.d/next/Core and Builtins/2023-01-25-00-55-33.gh-issue-97933.1ET5Ig.rst

@@ -0,0 +1 @@
+New COMPREHENSION bytecode instruction executes a comprehension more efficiently, without allocating a single-use function object.

Objects/frame_layout.md

@@ -75,6 +75,7 @@ The specials sections contains the following pointers:
 * Builtins dict
 * Locals dict (not the "fast" locals, but the locals for eval and class creation)
 * Code object
+* Closure tuple of cells for free variables, if any.
 * Heap allocated `PyFrameObject` for this activation record, if any.
 * The function.
 

Objects/frameobject.c

@@ -884,6 +884,7 @@ frame_dealloc(PyFrameObject *f)
         frame->f_code = NULL;
         Py_CLEAR(frame->f_funcobj);
         Py_CLEAR(frame->f_locals);
+        Py_CLEAR(frame->f_closure);
         PyObject **locals = _PyFrame_GetLocalsArray(frame);
         for (int i = 0; i < frame->stacktop; i++) {
             Py_CLEAR(locals[i]);
@@ -1022,7 +1023,7 @@ init_frame(_PyInterpreterFrame *frame, PyFunctionObject *func, PyObject *locals)
 {
     PyCodeObject *code = (PyCodeObject *)func->func_code;
     _PyFrame_Initialize(frame, (PyFunctionObject*)Py_NewRef(func),
-                        Py_XNewRef(locals), code, 0);
+                        Py_XNewRef(locals), code, func->func_closure, 0);
     frame->previous = NULL;
 }
 
@@ -1125,7 +1126,7 @@ frame_init_get_vars(_PyInterpreterFrame *frame)
     }
 
     /* Free vars have not been initialized -- Do that */
-    PyObject *closure = ((PyFunctionObject *)frame->f_funcobj)->func_closure;
+    PyObject *closure = frame->f_closure;
     int offset = PyCode_GetFirstFree(co);
     for (int i = 0; i < co->co_nfreevars; ++i) {
         PyObject *o = PyTuple_GET_ITEM(closure, i);

Python/bytecodes.c

@@ -1248,8 +1248,7 @@ dummy_func(
         inst(COPY_FREE_VARS, (--)) {
             /* Copy closure variables to free variables */
             PyCodeObject *co = frame->f_code;
-            assert(PyFunction_Check(frame->f_funcobj));
-            PyObject *closure = ((PyFunctionObject *)frame->f_funcobj)->func_closure;
+            PyObject *closure = frame->f_closure;
             assert(oparg == co->co_nfreevars);
             int offset = co->co_nlocalsplus - oparg;
             for (int i = 0; i < oparg; ++i) {
@@ -2383,6 +2382,20 @@ dummy_func(
             kwnames = GETITEM(consts, oparg);
         }
 
+        inst(COMPREHENSION, (closure if (oparg), code, unused -- res)) {
+            _PyInterpreterFrame *new_frame = _PyEvalFramePushAndInit(
+                tstate, (PyFunctionObject *)Py_NewRef(frame->f_funcobj), code,
+                closure, NULL, stack_pointer - 1, 1, NULL
+            );
+            Py_XDECREF(code);
+            Py_XDECREF(closure);
+            STACK_SHRINK(oparg + 2);
+            if (new_frame == NULL) {
+                goto error;
+            }
+            DISPATCH_INLINED(new_frame);
+        }
+
         // Cache layout: counter/1, func_version/2, min_args/1
         // Neither CALL_INTRINSIC_1/2 nor CALL_FUNCTION_EX are members!
         family(call, INLINE_CACHE_ENTRIES_CALL) = {
@@ -2451,8 +2464,9 @@ dummy_func(
             {
                 int code_flags = ((PyCodeObject*)PyFunction_GET_CODE(callable))->co_flags;
                 PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : Py_NewRef(PyFunction_GET_GLOBALS(callable));
+                PyFunctionObject *func = (PyFunctionObject *)callable;
                 _PyInterpreterFrame *new_frame = _PyEvalFramePushAndInit(
-                    tstate, (PyFunctionObject *)callable, locals,
+                    tstate, func, func->func_code, func->func_closure, locals,
                     args, positional_args, kwnames
                 );
                 kwnames = NULL;

Python/ceval.c

@@ -208,8 +208,8 @@ static void format_awaitable_error(PyThreadState *, PyTypeObject *, int);
 static int get_exception_handler(PyCodeObject *, int, int*, int*, int*);
 static _PyInterpreterFrame *
 _PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
-                        PyObject *locals, PyObject* const* args,
-                        size_t argcount, PyObject *kwnames);
+                        PyObject *code, PyObject *closure, PyObject *locals,
+                        PyObject* const* args, size_t argcount, PyObject *kwnames);
 static void
 _PyEvalFrameClearAndPop(PyThreadState *tstate, _PyInterpreterFrame *frame);
 
@@ -744,6 +744,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     entry_frame.frame_obj = (PyFrameObject*)0xaaa2;
     entry_frame.f_globals = (PyObject*)0xaaa3;
     entry_frame.f_builtins = (PyObject*)0xaaa4;
+    entry_frame.f_closure = (PyObject*)0xaaa5;
 #endif
     entry_frame.f_code = tstate->interp->interpreter_trampoline;
     entry_frame.prev_instr =
@@ -1382,10 +1383,9 @@ get_exception_handler(PyCodeObject *code, int index, int *level, int *handler, i
 
 static int
 initialize_locals(PyThreadState *tstate, PyFunctionObject *func,
-    PyObject **localsplus, PyObject *const *args,
+    PyCodeObject *co, PyObject **localsplus, PyObject *const *args,
     Py_ssize_t argcount, PyObject *kwnames)
 {
-    PyCodeObject *co = (PyCodeObject*)func->func_code;
     const Py_ssize_t total_args = co->co_argcount + co->co_kwonlyargcount;
 
     /* Create a dictionary for keyword parameters (**kwags) */
@@ -1607,18 +1607,19 @@ initialize_locals(PyThreadState *tstate, PyFunctionObject *func,
 /* Consumes references to func, locals and all the args */
 static _PyInterpreterFrame *
 _PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
-                        PyObject *locals, PyObject* const* args,
-                        size_t argcount, PyObject *kwnames)
+                        PyObject *codeobj, PyObject *closure, PyObject *locals,
+                        PyObject* const* args, size_t argcount, PyObject *kwnames)
 {
-    PyCodeObject * code = (PyCodeObject *)func->func_code;
     CALL_STAT_INC(frames_pushed);
+    assert(PyCode_Check(codeobj));
+    PyCodeObject *code = (PyCodeObject *)codeobj;
     _PyInterpreterFrame *frame = _PyThreadState_PushFrame(tstate, code->co_framesize);
     if (frame == NULL) {
         goto fail;
     }
-    _PyFrame_Initialize(frame, func, locals, code, 0);
+    _PyFrame_Initialize(frame, func, locals, code, closure, 0);
     PyObject **localsarray = &frame->localsplus[0];
-    if (initialize_locals(tstate, func, localsarray, args, argcount, kwnames)) {
+    if (initialize_locals(tstate, func, code, localsarray, args, argcount, kwnames)) {
         assert(frame->owner != FRAME_OWNED_BY_GENERATOR);
         _PyEvalFrameClearAndPop(tstate, frame);
         return NULL;
@@ -1702,7 +1703,8 @@ _PyEval_Vector(PyThreadState *tstate, PyFunctionObject *func,
         }
     }
     _PyInterpreterFrame *frame = _PyEvalFramePushAndInit(
-        tstate, func, locals, args, argcount, kwnames);
+        tstate, func, func->func_code, func->func_closure, locals,
+        args, argcount, kwnames);
     if (frame == NULL) {
         return NULL;
     }