gh-102013: Add PyUnstable_GC_VisitObjects (#102014)

2 files changed, 93 insertions, 0 deletions

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index ea67017..f45d031 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -3310,6 +3310,73 @@ function_set_kw_defaults(PyObject *self, PyObject *args)
     Py_RETURN_NONE;
 }
 
+struct gc_visit_state_basic {
+    PyObject *target;
+    int found;
+};
+
+static int
+gc_visit_callback_basic(PyObject *obj, void *arg)
+{
+    struct gc_visit_state_basic *state = (struct gc_visit_state_basic *)arg;
+    if (obj == state->target) {
+        state->found = 1;
+        return 0;
+    }
+    return 1;
+}
+
+static PyObject *
+test_gc_visit_objects_basic(PyObject *Py_UNUSED(self),
+                            PyObject *Py_UNUSED(ignored))
+{
+    PyObject *obj;
+    struct gc_visit_state_basic state;
+
+    obj = PyList_New(0);
+    if (obj == NULL) {
+        return NULL;
+    }
+    state.target = obj;
+    state.found = 0;
+    
+    PyUnstable_GC_VisitObjects(gc_visit_callback_basic, &state);
+    Py_DECREF(obj);
+    if (!state.found) {
+        PyErr_SetString(
+             PyExc_AssertionError,
+             "test_gc_visit_objects_basic: Didn't find live list");
+         return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+static int
+gc_visit_callback_exit_early(PyObject *obj, void *arg)
+ {
+    int *visited_i = (int *)arg;
+    (*visited_i)++;
+    if (*visited_i == 2) {
+        return 0;
+    }
+    return 1;
+}
+
+static PyObject *
+test_gc_visit_objects_exit_early(PyObject *Py_UNUSED(self),
+                                 PyObject *Py_UNUSED(ignored))
+{
+    int visited_i = 0;
+    PyUnstable_GC_VisitObjects(gc_visit_callback_exit_early, &visited_i);
+    if (visited_i != 2) {
+        PyErr_SetString(
+            PyExc_AssertionError,
+            "test_gc_visit_objects_exit_early: did not exit when expected");
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyObject *test_buildvalue_issue38913(PyObject *, PyObject *);
 
 static PyMethodDef TestMethods[] = {
@@ -3452,6 +3519,8 @@ static PyMethodDef TestMethods[] = {
     {"function_set_defaults", function_set_defaults, METH_VARARGS, NULL},
     {"function_get_kw_defaults", function_get_kw_defaults, METH_O, NULL},
     {"function_set_kw_defaults", function_set_kw_defaults, METH_VARARGS, NULL},
+    {"test_gc_visit_objects_basic", test_gc_visit_objects_basic, METH_NOARGS, NULL},
+    {"test_gc_visit_objects_exit_early", test_gc_visit_objects_exit_early, METH_NOARGS, NULL},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index 5879c5e..4eaa549 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -2401,3 +2401,27 @@ PyObject_GC_IsFinalized(PyObject *obj)
     }
     return 0;
 }
+
+void
+PyUnstable_GC_VisitObjects(gcvisitobjects_t callback, void *arg)
+{
+    size_t i;
+    GCState *gcstate = get_gc_state();
+    int origenstate = gcstate->enabled;
+    gcstate->enabled = 0;
+    for (i = 0; i < NUM_GENERATIONS; i++) {
+        PyGC_Head *gc_list, *gc;
+        gc_list = GEN_HEAD(gcstate, i);
+        for (gc = GC_NEXT(gc_list); gc != gc_list; gc = GC_NEXT(gc)) {
+            PyObject *op = FROM_GC(gc);
+            Py_INCREF(op);
+            int res = callback(op, arg);
+            Py_DECREF(op);
+            if (!res) {
+                goto done;
+            }
+        }
+    }
+done:
+    gcstate->enabled = origenstate;
+}