gh-103743: Add PyUnstable_Object_GC_NewWithExtraData (GH-103744)

Co-authored-by: Petr Viktorin <encukou@gmail.com>
Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>

6 files changed, 156 insertions, 2 deletions

diff --git a/Doc/c-api/gcsupport.rst b/Doc/c-api/gcsupport.rst
index cb5d64a..c3260a2 100644
--- a/Doc/c-api/gcsupport.rst
+++ b/Doc/c-api/gcsupport.rst
@@ -59,12 +59,31 @@ rules:
    Analogous to :c:func:`PyObject_New` but for container objects with the
    :const:`Py_TPFLAGS_HAVE_GC` flag set.
 
-
 .. c:function:: TYPE* PyObject_GC_NewVar(TYPE, PyTypeObject *type, Py_ssize_t size)
 
    Analogous to :c:func:`PyObject_NewVar` but for container objects with the
    :const:`Py_TPFLAGS_HAVE_GC` flag set.
 
+.. c:function:: PyObject* PyUnstable_Object_GC_NewWithExtraData(PyTypeObject *type, size_t extra_size)
+
+   Analogous to :c:func:`PyObject_GC_New` but allocates *extra_size*
+   bytes at the end of the object (at offset
+   :c:member:`~PyTypeObject.tp_basicsize`).
+   The allocated memory is initialized to zeros,
+   except for the :c:type:`Python object header <PyObject>`.
+
+   The extra data will be deallocated with the object, but otherwise it is
+   not managed by Python.
+
+   .. warning::
+      The function is marked as unstable because the final mechanism
+      for reserving extra data after an instance is not yet decided.
+      For allocating a variable number of fields, prefer using
+      :c:type:`PyVarObject` and :c:member:`~PyTypeObject.tp_itemsize`
+      instead.
+
+   .. versionadded:: 3.12
+
 
 .. c:function:: TYPE* PyObject_GC_Resize(TYPE, PyVarObject *op, Py_ssize_t newsize)
 
diff --git a/Include/cpython/objimpl.h b/Include/cpython/objimpl.h
index 0b038d3..5a8cdd5 100644
--- a/Include/cpython/objimpl.h
+++ b/Include/cpython/objimpl.h
@@ -90,3 +90,6 @@ PyAPI_FUNC(int) PyObject_IS_GC(PyObject *obj);
 PyAPI_FUNC(int) PyType_SUPPORTS_WEAKREFS(PyTypeObject *type);
 
 PyAPI_FUNC(PyObject **) PyObject_GET_WEAKREFS_LISTPTR(PyObject *op);
+
+PyAPI_FUNC(PyObject *) PyUnstable_Object_GC_NewWithExtraData(PyTypeObject *,
+                                                             size_t);
diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
index 9470cf1..9d5d1ca 100644
--- a/Lib/test/test_capi/test_misc.py
+++ b/Lib/test/test_capi/test_misc.py
@@ -1043,6 +1043,20 @@ class CAPITest(unittest.TestCase):
         self.assertEqual(_testcapi.function_get_kw_defaults(some), None)
         self.assertEqual(some.__kwdefaults__, None)
 
+    def test_unstable_gc_new_with_extra_data(self):
+        class Data(_testcapi.ObjExtraData):
+            __slots__ = ('x', 'y')
+
+        d = Data()
+        d.x = 10
+        d.y = 20
+        d.extra = 30
+        self.assertEqual(d.x, 10)
+        self.assertEqual(d.y, 20)
+        self.assertEqual(d.extra, 30)
+        del d.extra
+        self.assertIsNone(d.extra)
+
 
 class TestPendingCalls(unittest.TestCase):
 
diff --git a/Misc/NEWS.d/next/C API/2023-04-24-10-31-59.gh-issue-103743.2xYA1K.rst b/Misc/NEWS.d/next/C API/2023-04-24-10-31-59.gh-issue-103743.2xYA1K.rst
new file mode 100644
index 0000000..d074350
--- /dev/null
+++ b/Misc/NEWS.d/next/C API/2023-04-24-10-31-59.gh-issue-103743.2xYA1K.rst
@@ -0,0 +1,2 @@
+Add :c:func:`PyUnstable_Object_GC_NewWithExtraData` function that can be used to
+allocate additional memory after an object for data not managed by Python.
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index c1892f6..a5d23b1 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -3363,7 +3363,7 @@ test_gc_visit_objects_basic(PyObject *Py_UNUSED(self),
     }
     state.target = obj;
     state.found = 0;
-    
+
     PyUnstable_GC_VisitObjects(gc_visit_callback_basic, &state);
     Py_DECREF(obj);
     if (!state.found) {
@@ -3400,6 +3400,98 @@ test_gc_visit_objects_exit_early(PyObject *Py_UNUSED(self),
     Py_RETURN_NONE;
 }
 
+typedef struct {
+    PyObject_HEAD
+} ObjExtraData;
+
+static PyObject *
+obj_extra_data_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
+{
+    size_t extra_size = sizeof(PyObject *);
+    PyObject *obj = PyUnstable_Object_GC_NewWithExtraData(type, extra_size);
+    if (obj == NULL) {
+        return PyErr_NoMemory();
+    }
+    PyObject_GC_Track(obj);
+    return obj;
+}
+
+static PyObject **
+obj_extra_data_get_extra_storage(PyObject *self)
+{
+    return (PyObject **)((char *)self + Py_TYPE(self)->tp_basicsize);
+}
+
+static PyObject *
+obj_extra_data_get(PyObject *self, void *Py_UNUSED(ignored))
+{
+    PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
+    PyObject *value = *extra_storage;
+    if (!value) {
+        Py_RETURN_NONE;
+    }
+    return Py_NewRef(value);
+}
+
+static int
+obj_extra_data_set(PyObject *self, PyObject *newval, void *Py_UNUSED(ignored))
+{
+    PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
+    Py_CLEAR(*extra_storage);
+    if (newval) {
+        *extra_storage = Py_NewRef(newval);
+    }
+    return 0;
+}
+
+static PyGetSetDef obj_extra_data_getset[] = {
+    {"extra", (getter)obj_extra_data_get, (setter)obj_extra_data_set, NULL},
+    {NULL}
+};
+
+static int
+obj_extra_data_traverse(PyObject *self, visitproc visit, void *arg)
+{
+    PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
+    PyObject *value = *extra_storage;
+    Py_VISIT(value);
+    return 0;
+}
+
+static int
+obj_extra_data_clear(PyObject *self)
+{
+    PyObject **extra_storage = obj_extra_data_get_extra_storage(self);
+    Py_CLEAR(*extra_storage);
+    return 0;
+}
+
+static void
+obj_extra_data_dealloc(PyObject *self)
+{
+    PyTypeObject *tp = Py_TYPE(self);
+    PyObject_GC_UnTrack(self);
+    obj_extra_data_clear(self);
+    tp->tp_free(self);
+    Py_DECREF(tp);
+}
+
+static PyType_Slot ObjExtraData_Slots[] = {
+    {Py_tp_getset, obj_extra_data_getset},
+    {Py_tp_dealloc, obj_extra_data_dealloc},
+    {Py_tp_traverse, obj_extra_data_traverse},
+    {Py_tp_clear, obj_extra_data_clear},
+    {Py_tp_new, obj_extra_data_new},
+    {Py_tp_free, PyObject_GC_Del},
+    {0, NULL},
+};
+
+static PyType_Spec ObjExtraData_TypeSpec = {
+    .name = "_testcapi.ObjExtraData",
+    .basicsize = sizeof(ObjExtraData),
+    .flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
+    .slots = ObjExtraData_Slots,
+};
 
 struct atexit_data {
     int called;
@@ -4124,6 +4216,17 @@ PyInit__testcapi(void)
     Py_INCREF(&MethStatic_Type);
     PyModule_AddObject(m, "MethStatic", (PyObject *)&MethStatic_Type);
 
+    PyObject *ObjExtraData_Type = PyType_FromModuleAndSpec(
+        m, &ObjExtraData_TypeSpec, NULL);
+    if (ObjExtraData_Type == 0) {
+        return NULL;
+    }
+    int ret = PyModule_AddType(m, (PyTypeObject*)ObjExtraData_Type);
+    Py_DECREF(&ObjExtraData_Type);
+    if (ret < 0) {
+        return NULL;
+    }
+
     PyModule_AddObject(m, "CHAR_MAX", PyLong_FromLong(CHAR_MAX));
     PyModule_AddObject(m, "CHAR_MIN", PyLong_FromLong(CHAR_MIN));
     PyModule_AddObject(m, "UCHAR_MAX", PyLong_FromLong(UCHAR_MAX));
diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index 3fd5f4c..8a4d1a4 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -2367,6 +2367,19 @@ _PyObject_GC_NewVar(PyTypeObject *tp, Py_ssize_t nitems)
     return op;
 }
 
+PyObject *
+PyUnstable_Object_GC_NewWithExtraData(PyTypeObject *tp, size_t extra_size)
+{
+    size_t presize = _PyType_PreHeaderSize(tp);
+    PyObject *op = gc_alloc(_PyObject_SIZE(tp) + extra_size, presize);
+    if (op == NULL) {
+        return NULL;
+    }
+    memset(op, 0, _PyObject_SIZE(tp) + extra_size);
+    _PyObject_Init(op, tp);
+    return op;
+}
+
 PyVarObject *
 _PyObject_GC_Resize(PyVarObject *op, Py_ssize_t nitems)
 {