add opcode for more efficient comprehension execution

Author: carljm

Doc/library/dis.rst

@@ -1304,6 +1304,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

@@ -52,6 +52,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 */
     PyCodeObject *f_code; /* Strong reference */
     PyFrameObject *frame_obj; /* Strong reference, may be NULL. Only valid if not on C stack */
     /* Linkage section */
@@ -112,12 +115,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;
@@ -250,7 +255,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;
 }
 

Include/internal/pycore_opcode.h

@@ -431,6 +431,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a1 3515 (Embed jump mask in COMPARE_OP oparg)
 #     Python 3.12a1 3516 (Add COMPARE_AND_BRANCH instruction)
 #     Python 3.12a1 3517 (Change YIELD_VALUE oparg to exception block depth)
+#     Python 3.12a1 3518 (Add COMPREHENSION instruction)
 
 #     Python 3.13 will start with 3550
 
@@ -443,7 +444,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 = (3517).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3518).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Include/opcode.h

@@ -218,6 +218,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/importlib/_bootstrap_external.py

@@ -170,10 +170,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)
 
@@ -675,10 +674,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/opcode.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

Lib/test/test_dis.py

@@ -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.
 

Lib/test/test_sys.py

@@ -882,6 +882,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]);
@@ -1020,7 +1021,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;
 }
 
@@ -1123,7 +1124,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);

Objects/frame_layout.md

@@ -1265,8 +1265,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) {
@@ -2515,6 +2514,23 @@ dummy_func(
             kwnames = GETITEM(consts, oparg);
         }
 
+        // stack effect: (__0, __1, __array[oparg] --)
+        inst(COMPREHENSION) {
+            PyObject *code = PEEK(2);
+            PyObject *closure = oparg ? PEEK(3) : NULL;
+            _PyInterpreterFrame *new_frame = _PyEvalFramePushAndInit(
+                tstate, 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);
+        }
+
         // stack effect: (__0, __array[oparg] -- )
         inst(CALL) {
             #if ENABLE_SPECIALIZATION
@@ -2553,8 +2569,9 @@ dummy_func(
                 int code_flags = ((PyCodeObject*)PyFunction_GET_CODE(function))->co_flags;
                 PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : Py_NewRef(PyFunction_GET_GLOBALS(function));
                 STACK_SHRINK(total_args);
+                PyFunctionObject *func = (PyFunctionObject *)function;
                 _PyInterpreterFrame *new_frame = _PyEvalFramePushAndInit(
-                    tstate, (PyFunctionObject *)function, locals,
+                    tstate, func, func->func_code, func->func_closure, locals,
                     stack_pointer, positional_args, kwnames
                 );
                 kwnames = NULL;

Objects/frameobject.c

@@ -210,8 +210,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);
 
@@ -752,6 +752,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 =
@@ -1388,10 +1389,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) */
@@ -1613,18 +1613,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;
@@ -1708,7 +1709,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;
     }