gh-117139: Set up the tagged evaluation stack

Include/internal/pycore_pystate.h

@@ -286,6 +286,16 @@ static inline struct _Py_object_freelists* _Py_object_freelists_GET(void)
 #endif
 }
 
+extern PyInterpreterState *_Py_interp_for_stop_the_world(struct _stoptheworld_state *stw);
+
+// Loops over threads for a stop-the-world event.
+// For global: all threads in all interpreters
+// For per-interpreter: all threads in the interpreter
+#define _Py_FOR_EACH_THREAD(stw, i, t)                                      \
+    for (i = _Py_interp_for_stop_the_world((stw));                              \
+            i != NULL; i = ((stw->is_global) ? i->next : NULL))             \
+        for (t = i->threads.head; t; t = t->next)
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_tagged.h

@@ -50,6 +50,7 @@ typedef union {
     #define Py_STACK_TAG(obj) ((_PyStackRef){.bits = ((uintptr_t)(obj))})
 #endif
 
+#define Py_STACK_TAG_UNSAFE(obj) ((_PyStackRef){.bits = ((uintptr_t)(obj))})
 
 #if defined(Py_TAG_TEST)
     #define Py_STACK_UNTAG_OWNED(tagged) Py_STACK_UNTAG_BORROWED(tagged)
@@ -60,6 +61,7 @@ typedef union {
             assert(_PyObject_HasDeferredRefcount(Py_STACK_UNTAG_BORROWED(tagged)));
             return Py_NewRef(Py_STACK_UNTAG_BORROWED(tagged));
         }
+        assert(!_PyObject_HasDeferredRefcount(Py_STACK_UNTAG_BORROWED(tagged)));
         return Py_STACK_UNTAG_BORROWED(tagged);
     }
     #define Py_STACK_UNTAG_OWNED(tagged) _Py_Stack_Untag_Owned(tagged)
@@ -128,11 +130,11 @@ _Py_untag_stack_owned(PyObject **dst, const _PyStackRef *src, size_t length) {
 #if defined(Py_GIL_DISABLED)
     static inline void
     _Py_IncRef_StackRef(_PyStackRef tagged) {
-        if ((tagged.bits & Py_TAG_DEFERRED) == Py_TAG_DEFERRED) {
-            assert(_PyObject_HasDeferredRefcount(Py_STACK_UNTAG_BORROWED(tagged)));
-            return;
-        }
-        assert(!_PyObject_HasDeferredRefcount(Py_STACK_UNTAG_BORROWED(tagged)));
+//        if ((tagged.bits & Py_TAG_DEFERRED) == Py_TAG_DEFERRED) {
+//            assert(_PyObject_HasDeferredRefcount(Py_STACK_UNTAG_BORROWED(tagged)));
+//            return;
+//        }
+//        assert(!_PyObject_HasDeferredRefcount(Py_STACK_UNTAG_BORROWED(tagged)));
         return Py_INCREF(Py_STACK_UNTAG_BORROWED(tagged));
     }
     #define Py_INCREF_STACKREF(op) _Py_IncRef_StackRef(op)

Python/ceval.c

@@ -252,6 +252,10 @@ static _PyInterpreterFrame *
 _PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
                         PyObject *locals, _PyStackRef const* args,
                         size_t argcount, PyObject *kwnames);
+static _PyInterpreterFrame *
+_PyEvalFramePushAndInit_UnTagged(PyThreadState *tstate, PyFunctionObject *func,
+                        PyObject *locals, PyObject *const* args,
+                        size_t argcount, PyObject *kwnames);
 static  _PyInterpreterFrame *
 _PyEvalFramePushAndInit_Ex(PyThreadState *tstate, PyFunctionObject *func,
     PyObject *locals, Py_ssize_t nargs, PyObject *callargs, PyObject *kwargs);
@@ -1481,7 +1485,7 @@ initialize_locals(PyThreadState *tstate, PyFunctionObject *func,
         for (i = 0; i < kwcount; i++) {
             PyObject **co_varnames;
             PyObject *keyword = PyTuple_GET_ITEM(kwnames, i);
-            PyObject *value = Py_STACK_UNTAG_BORROWED(args[i+argcount]);
+            _PyStackRef value_tagged = args[i+argcount];
             Py_ssize_t j;
 
             if (keyword == NULL || !PyUnicode_Check(keyword)) {
@@ -1554,16 +1558,15 @@ initialize_locals(PyThreadState *tstate, PyFunctionObject *func,
                 goto kw_fail;
             }
 
-            if (PyDict_SetItem(kwdict, keyword, value) == -1) {
+            if (PyDict_SetItem(kwdict, keyword, Py_STACK_UNTAG_OWNED(value_tagged)) == -1) {
                 goto kw_fail;
             }
-            Py_DECREF(value);
+            Py_DECREF_STACKREF(value_tagged);
             continue;
 
         kw_fail:
             for (;i < kwcount; i++) {
-                PyObject *value = Py_STACK_UNTAG_BORROWED(args[i+argcount]);
-                Py_DECREF(value);
+                Py_DECREF_STACKREF(args[i+argcount]);
             }
             goto fail_post_args;
 
@@ -1574,7 +1577,7 @@ initialize_locals(PyThreadState *tstate, PyFunctionObject *func,
                           func->func_qualname, keyword);
                 goto kw_fail;
             }
-            localsplus[j] = Py_STACK_TAG(value);
+            localsplus[j] = value_tagged;
         }
     }
 
@@ -1737,6 +1740,27 @@ _PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
     return NULL;
 }
 
+static _PyInterpreterFrame *
+_PyEvalFramePushAndInit_UnTagged(PyThreadState *tstate, PyFunctionObject *func,
+                        PyObject *locals, PyObject *const* args,
+                        size_t argcount, PyObject *kwnames)
+{
+#if defined(Py_GIL_DISABLED) || defined(Py_TAG_TEST)
+    _PyStackRef *tagged_args_buffer = PyMem_Malloc(sizeof(_PyStackRef) * argcount);
+    if (tagged_args_buffer == NULL) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+    for (size_t i = 0; i < argcount; i++) {
+        tagged_args_buffer[i] = Py_STACK_TAG(args[i]);
+    }
+    return _PyEvalFramePushAndInit(tstate, func, locals, (_PyStackRef const *)tagged_args_buffer, argcount, kwnames);
+#else
+    assert(Py_TAG == 0);
+    return _PyEvalFramePushAndInit(tstate, func, locals, (_PyStackRef const *)args, argcount, kwnames);
+#endif
+}
+
 /* Same as _PyEvalFramePushAndInit but takes an args tuple and kwargs dict.
    Steals references to func, callargs and kwargs.
 */
@@ -1761,9 +1785,9 @@ _PyEvalFramePushAndInit_Ex(PyThreadState *tstate, PyFunctionObject *func,
             Py_INCREF(PyTuple_GET_ITEM(callargs, i));
         }
     }
-    _PyInterpreterFrame *new_frame = _PyEvalFramePushAndInit(
+    _PyInterpreterFrame *new_frame = _PyEvalFramePushAndInit_UnTagged(
         tstate, (PyFunctionObject *)func, locals,
-        (_PyStackRef const *)newargs, nargs, kwnames
+        newargs, nargs, kwnames
     );
     if (has_dict) {
         _PyStack_UnpackDict_FreeNoDecRef(newargs, kwnames);
@@ -1799,8 +1823,8 @@ _PyEval_Vector(PyThreadState *tstate, PyFunctionObject *func,
             Py_INCREF(args[i+argcount]);
         }
     }
-    _PyInterpreterFrame *frame = _PyEvalFramePushAndInit(
-        tstate, func, locals, (_PyStackRef * const)args, argcount, kwnames);
+    _PyInterpreterFrame *frame = _PyEvalFramePushAndInit_UnTagged(
+        tstate, func, locals, args, argcount, kwnames);
     if (frame == NULL) {
         return NULL;
     }

Python/gc_free_threading.c

@@ -15,6 +15,7 @@
 #include "pycore_tstate.h"        // _PyThreadStateImpl
 #include "pycore_weakref.h"       // _PyWeakref_ClearRef()
 #include "pydtrace.h"
+#include "pycore_tagged.h"
 
 #ifdef Py_GIL_DISABLED
 
@@ -308,6 +309,27 @@ gc_visit_heaps(PyInterpreterState *interp, mi_block_visit_fun *visitor,
     return err;
 }
 
+static void
+gc_visit_thread_stacks(struct _stoptheworld_state *stw)
+{
+    HEAD_LOCK(&_PyRuntime);
+    PyInterpreterState *interp;
+    PyThreadState *tstate;
+    _Py_FOR_EACH_THREAD(stw, interp, tstate) {
+        _PyStackChunk *curr_chunk = tstate->datastack_chunk;
+        while (curr_chunk != NULL) {
+            for (size_t curr_i = 0; curr_i < curr_chunk->top; curr_i++) {
+                _PyStackRef curr_o = Py_STACK_TAG_UNSAFE(curr_chunk->data[curr_i]);
+                if ((curr_o.bits & Py_TAG_DEFERRED) == Py_TAG_DEFERRED) {
+                    gc_add_refs(Py_STACK_UNTAG_OWNED(curr_o), 1);
+                }
+            }
+            curr_chunk = curr_chunk->previous;
+        }
+    }
+    HEAD_UNLOCK(&_PyRuntime);
+}
+
 static void
 merge_queued_objects(_PyThreadStateImpl *tstate, struct collection_state *state)
 {
@@ -566,6 +588,8 @@ deduce_unreachable_heap(PyInterpreterState *interp,
     gc_visit_heaps(interp, &validate_gc_objects, &state->base);
 #endif
 
+    gc_visit_thread_stacks(&interp->stoptheworld);
+
     // Transitively mark reachable objects by clearing the
     // _PyGC_BITS_UNREACHABLE flag.
     if (gc_visit_heaps(interp, &mark_heap_visitor, &state->base) < 0) {

Python/pystate.c

@@ -2079,23 +2079,14 @@ decrement_stoptheworld_countdown(struct _stoptheworld_state *stw)
 // we start with the first interpreter and then iterate over all interpreters.
 // For per-interpreter stop-the-world events, we only operate on the one
 // interpreter.
-static PyInterpreterState *
-interp_for_stop_the_world(struct _stoptheworld_state *stw)
+PyInterpreterState *
+_Py_interp_for_stop_the_world(struct _stoptheworld_state *stw)
 {
     return (stw->is_global
         ? PyInterpreterState_Head()
         : _Py_CONTAINER_OF(stw, PyInterpreterState, stoptheworld));
 }
 
-// Loops over threads for a stop-the-world event.
-// For global: all threads in all interpreters
-// For per-interpreter: all threads in the interpreter
-#define _Py_FOR_EACH_THREAD(stw, i, t)                                      \
-    for (i = interp_for_stop_the_world((stw));                              \
-            i != NULL; i = ((stw->is_global) ? i->next : NULL))             \
-        for (t = i->threads.head; t; t = t->next)
-
-
 // Try to transition threads atomically from the "detached" state to the
 // "gc stopped" state. Returns true if all threads are in the "gc stopped"
 static bool