Add support for asyn generators on fast path. Simplify counting

python · markshannon · Dec 6, 2024 · Nov 9, 2024 · Nov 10, 2024 · Nov 11, 2024
commit e8497ae4285efee06ee97bbb02bf851e9d7bfba3
diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
@@ -409,6 +409,9 @@ _PyEvalFramePushAndInit(PyThreadState *tstate, _PyStackRef func,
                         size_t argcount, PyObject *kwnames,
                         _PyInterpreterFrame *previous);
 
+void
+_PyFrame_MoveToReachable(_PyInterpreterFrame *frame, PyGC_Head *reachable, int visited_space);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Include/internal/pycore_gc.h b/Include/internal/pycore_gc.h
@@ -401,6 +401,8 @@ extern void _PyGC_VisitObjectsWorldStopped(PyInterpreterState *interp,
                                            gcvisitobjects_t callback, void *arg);
 #endif
 
+void _PyGC_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
@@ -29,8 +29,10 @@ extern "C" {
 #define _Py_TYPE_VERSION_COROUTINE 13
 #define _Py_TYPE_VERSION_MODULE 14
 #define _Py_TYPE_VERSION_TYPE 15
+#define _Py_TYPE_VERSION_ASYNC_GENERATOR 16
+#define _Py_TYPE_VERSION_ASYNC_ASEND 17
 
-#define _Py_TYPE_VERSION_NEXT 16
+#define _Py_TYPE_VERSION_NEXT 20
 
 
 #define _Py_TYPE_BASE_VERSION_TAG (2<<16)

@@ -4475,6 +4475,42 @@ dict_popitem_impl(PyDictObject *self)
     return res;
 }
 
+void
+_PyDict_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyDictObject *mp = (PyDictObject *)op;
+    PyDictKeysObject *keys = mp->ma_keys;
+    Py_ssize_t i, n = keys->dk_nentries;
+    if (DK_IS_UNICODE(keys)) {
+        if (_PyDict_HasSplitTable(mp)) {
+            if (!mp->ma_values->embedded) {
+                for (i = 0; i < n; i++) {
+                    PyObject *value = mp->ma_values->values[i];
+                    _PyGC_MoveToReachable(value, reachable, visited_space);
+                }
+            }
+        }
+        else {
+            PyDictUnicodeEntry *entries = DK_UNICODE_ENTRIES(keys);
+            for (i = 0; i < n; i++) {
+                PyObject *value = entries[i].me_value;
+                _PyGC_MoveToReachable(value, reachable, visited_space);
+            }
+        }
+    }
+    else {
+        PyDictKeyEntry *entries = DK_ENTRIES(keys);
+        for (i = 0; i < n; i++) {
+            if (entries[i].me_value != NULL) {
+                PyObject *key = entries[i].me_key;
+                _PyGC_MoveToReachable(key, reachable, visited_space);
+                PyObject *value = entries[i].me_value;
+                _PyGC_MoveToReachable(value, reachable, visited_space);
+            }
+        }
+    }
+}
+
 static int
 dict_traverse(PyObject *op, visitproc visit, void *arg)
 {

@@ -78,6 +78,22 @@ gen_traverse(PyObject *self, visitproc visit, void *arg)
     return 0;
 }
 
+void
+_PyGen_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyGenObject *gen = (PyGenObject *)op;
+    if (gen->gi_frame_state == FRAME_CLEARED) {
+        return;
+    }
+    _PyGC_MoveToReachable(gen->gi_exc_state.exc_value, reachable, visited_space);
+    if (gen->gi_frame_state == FRAME_EXECUTING) {
+        /* if executing we already traversed it on the stack */
+        return;
+    }
+    _PyInterpreterFrame *frame = &gen->gi_iframe;
+   _PyFrame_MoveToReachable(frame, reachable, visited_space);
+}
+
 void
 _PyGen_Finalize(PyObject *self)
 {
@@ -1426,6 +1442,14 @@ typedef struct _PyAsyncGenWrappedValue {
      _Py_CAST(_PyAsyncGenWrappedValue*, (op)))
 
 
+void
+_PyAsyncGen_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyAsyncGenObject *ag = _PyAsyncGenObject_CAST(op);
+    _PyGC_MoveToReachable(ag->ag_origin_or_finalizer, reachable, visited_space);
+    _PyGen_MoveToReachable(op, reachable, visited_space);
+}
+
 static int
 async_gen_traverse(PyObject *self, visitproc visit, void *arg)
 {
@@ -1634,24 +1658,8 @@ PyTypeObject PyAsyncGen_Type = {
     async_gen_methods,                          /* tp_methods */
     async_gen_memberlist,                       /* tp_members */
     async_gen_getsetlist,                       /* tp_getset */
-    0,                                          /* tp_base */
-    0,                                          /* tp_dict */
-    0,                                          /* tp_descr_get */
-    0,                                          /* tp_descr_set */
-    0,                                          /* tp_dictoffset */
-    0,                                          /* tp_init */
-    0,                                          /* tp_alloc */
-    0,                                          /* tp_new */
-    0,                                          /* tp_free */
-    0,                                          /* tp_is_gc */
-    0,                                          /* tp_bases */
-    0,                                          /* tp_mro */
-    0,                                          /* tp_cache */
-    0,                                          /* tp_subclasses */
-    0,                                          /* tp_weaklist */
-    0,                                          /* tp_del */
-    0,                                          /* tp_version_tag */
-    _PyGen_Finalize,                            /* tp_finalize */
+    .tp_finalize = _PyGen_Finalize,
+    .tp_version_tag = _Py_TYPE_VERSION_ASYNC_GENERATOR,
 };
 
 
@@ -1724,6 +1732,14 @@ async_gen_asend_dealloc(PyObject *self)
     _Py_FREELIST_FREE(async_gen_asends, self, PyObject_GC_Del);
 }
 
+void
+_PyAsyncAsend_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyAsyncGenASend *ags = _PyAsyncGenASend_CAST(op);
+    _PyGC_MoveToReachable((PyObject *)ags->ags_sendval, reachable, visited_space);
+    _PyAsyncGen_MoveToReachable((PyObject *)ags->ags_gen, reachable, visited_space);
+}
+
 static int
 async_gen_asend_traverse(PyObject *self, visitproc visit, void *arg)
 {
@@ -1896,17 +1912,8 @@ PyTypeObject _PyAsyncGenASend_Type = {
     PyObject_SelfIter,                          /* tp_iter */
     async_gen_asend_iternext,                   /* tp_iternext */
     async_gen_asend_methods,                    /* tp_methods */
-    0,                                          /* tp_members */
-    0,                                          /* tp_getset */
-    0,                                          /* tp_base */
-    0,                                          /* tp_dict */
-    0,                                          /* tp_descr_get */
-    0,                                          /* tp_descr_set */
-    0,                                          /* tp_dictoffset */
-    0,                                          /* tp_init */
-    0,                                          /* tp_alloc */
-    0,                                          /* tp_new */
     .tp_finalize = async_gen_asend_finalize,
+    .tp_version_tag = _Py_TYPE_VERSION_ASYNC_ASEND,
 };
 
 

diff --git a/Objects/listobject.c b/Objects/listobject.c
@@ -3276,6 +3276,17 @@ list_remove_impl(PyListObject *self, PyObject *value)
     return NULL;
 }
 
+void
+_PyList_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyListObject *o = (PyListObject *)op;
+    Py_ssize_t i;
+    for (i = Py_SIZE(o); --i >= 0; ) {
+        PyObject *item = o->ob_item[i];
+        _PyGC_MoveToReachable(item, reachable, visited_space);
+    }
+}
+
 static int
 list_traverse(PyObject *self, visitproc visit, void *arg)
 {

diff --git a/Objects/setobject.c b/Objects/setobject.c
@@ -492,11 +492,6 @@ set_next(PySetObject *so, Py_ssize_t *pos_ptr, setentry **entry_ptr)
     return 1;
 }
 
-int _PyTemporary_SetNext(PySetObject *so, Py_ssize_t *pos_ptr, setentry **entry_ptr)
-{
-    return set_next(so, pos_ptr, entry_ptr);
-}
-
 static void
 set_dealloc(PyObject *self)
 {
@@ -705,6 +700,17 @@ set_traverse(PyObject *self, visitproc visit, void *arg)
     return 0;
 }
 
+void
+_PySet_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PySetObject *so = (PySetObject *)op;
+    Py_ssize_t pos = 0;
+    setentry *entry;
+    while (set_next(so, &pos, &entry)) {
+        _PyGC_MoveToReachable(entry->key, reachable, visited_space);
+    }
+}
+
 /* Work to increase the bit dispersion for closely spaced hash values.
    This is important because some use cases have many combinations of a
    small number of elements with nearby hashes so that many distinct

diff --git a/Objects/tupleobject.c b/Objects/tupleobject.c
@@ -623,6 +623,16 @@ tuple_count(PyTupleObject *self, PyObject *value)
     return PyLong_FromSsize_t(count);
 }
 
+void
+_PyTuple_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyTupleObject *o = (PyTupleObject *)op;
+    for (Py_ssize_t i = Py_SIZE(o); --i >= 0; ) {
+        PyObject *item = o->ob_item[i];
+        _PyGC_MoveToReachable(item, reachable, visited_space);
+    }
+}
+
 static int
 tuple_traverse(PyObject *self, visitproc visit, void *arg)
 {

@@ -6214,6 +6214,19 @@ PyDoc_STRVAR(type_doc,
 "type(object) -> the object's type\n"
 "type(name, bases, dict, **kwds) -> a new type");
 
+
+void
+_PyType_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space)
+{
+    PyTypeObject *type = (PyTypeObject *)op;
+    _PyGC_MoveToReachable(type->tp_dict, reachable, visited_space);
+    _PyGC_MoveToReachable(type->tp_cache, reachable, visited_space);
+    _PyGC_MoveToReachable(type->tp_mro, reachable, visited_space);
+    _PyGC_MoveToReachable(type->tp_bases, reachable, visited_space);
+    _PyGC_MoveToReachable((PyObject *)type->tp_base, reachable, visited_space);
+    _PyGC_MoveToReachable(((PyHeapTypeObject *)type)->ht_module, reachable, visited_space);
+}
+
 static int
 type_traverse(PyObject *self, visitproc visit, void *arg)
 {
-Original file line number
+Diff line change
@@ Expand Up @@
                                                gcvisitobjects_t callback, void *arg);
     #endif
+    void _PyGC_MoveToReachable(PyObject *op, PyGC_Head *reachable, int visited_space);
     #ifdef __cplusplus
     }
     #endif
@@ Expand Down @@