bpo-34784: Fix PyStructSequence_NewType with heap-allocated StructSequence (GH-9665)

1 files changed, 133 insertions, 92 deletions

diff --git a/Objects/structseq.c b/Objects/structseq.c
index 1705837..d513710 100644
--- a/Objects/structseq.c
+++ b/Objects/structseq.c
@@ -287,115 +287,118 @@ static PyMethodDef structseq_methods[] = {
     {NULL, NULL}
 };
 
-static PyTypeObject _struct_sequence_template = {
-    PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    NULL,                                       /* tp_name */
-    sizeof(PyStructSequence) - sizeof(PyObject *), /* tp_basicsize */
-    sizeof(PyObject *),                         /* tp_itemsize */
-    (destructor)structseq_dealloc,              /* tp_dealloc */
-    0,                                          /* tp_print */
-    0,                                          /* tp_getattr */
-    0,                                          /* tp_setattr */
-    0,                                          /* tp_reserved */
-    (reprfunc)structseq_repr,                   /* tp_repr */
-    0,                                          /* tp_as_number */
-    0,                                          /* tp_as_sequence */
-    0,                                          /* tp_as_mapping */
-    0,                                          /* tp_hash */
-    0,                                          /* tp_call */
-    0,                                          /* tp_str */
-    0,                                          /* tp_getattro */
-    0,                                          /* tp_setattro */
-    0,                                          /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT,                         /* tp_flags */
-    NULL,                                       /* tp_doc */
-    0,                                          /* tp_traverse */
-    0,                                          /* tp_clear */
-    0,                                          /* tp_richcompare */
-    0,                                          /* tp_weaklistoffset */
-    0,                                          /* tp_iter */
-    0,                                          /* tp_iternext */
-    structseq_methods,                          /* tp_methods */
-    NULL,                                       /* 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 */
-    structseq_new,                              /* tp_new */
-};
+static Py_ssize_t
+count_members(PyStructSequence_Desc *desc, Py_ssize_t *n_unnamed_members) {
+    Py_ssize_t i;
+
+    *n_unnamed_members = 0;
+    for (i = 0; desc->fields[i].name != NULL; ++i) {
+        if (desc->fields[i].name == PyStructSequence_UnnamedField) {
+            (*n_unnamed_members)++;
+        }
+    }
+    return i;
+}
+
+static int
+initialize_structseq_dict(PyStructSequence_Desc *desc, PyObject* dict,
+                          Py_ssize_t n_members, Py_ssize_t n_unnamed_members) {
+    PyObject *v;
+
+#define SET_DICT_FROM_SIZE(key, value)                                         \
+    do {                                                                       \
+        v = PyLong_FromSsize_t(value);                                         \
+        if (v == NULL) {                                                       \
+            return -1;                                                         \
+        }                                                                      \
+        if (PyDict_SetItemString(dict, key, v) < 0) {                          \
+            Py_DECREF(v);                                                      \
+            return -1;                                                         \
+        }                                                                      \
+        Py_DECREF(v);                                                          \
+    } while (0)
+
+    SET_DICT_FROM_SIZE(visible_length_key, desc->n_in_sequence);
+    SET_DICT_FROM_SIZE(real_length_key, n_members);
+    SET_DICT_FROM_SIZE(unnamed_fields_key, n_unnamed_members);
+    return 0;
+}
+
+static void
+initialize_members(PyStructSequence_Desc *desc, PyMemberDef* members,
+                   Py_ssize_t n_members) {
+    Py_ssize_t i, k;
+
+    for (i = k = 0; i < n_members; ++i) {
+        if (desc->fields[i].name == PyStructSequence_UnnamedField) {
+            continue;
+        }
+
+        /* The names and docstrings in these MemberDefs are statically */
+        /* allocated so it is expected that they'll outlive the MemberDef */
+        members[k].name = desc->fields[i].name;
+        members[k].type = T_OBJECT;
+        members[k].offset = offsetof(PyStructSequence, ob_item)
+          + i * sizeof(PyObject*);
+        members[k].flags = READONLY;
+        members[k].doc = desc->fields[i].doc;
+        k++;
+    }
+    members[k].name = NULL;
+}
 
 int
 PyStructSequence_InitType2(PyTypeObject *type, PyStructSequence_Desc *desc)
 {
-    PyObject *dict;
-    PyMemberDef* members;
-    Py_ssize_t n_members, n_unnamed_members, i, k;
-    PyObject *v;
+    PyMemberDef *members;
+    Py_ssize_t n_members, n_unnamed_members;
 
 #ifdef Py_TRACE_REFS
     /* if the type object was chained, unchain it first
        before overwriting its storage */
     if (type->ob_base.ob_base._ob_next) {
-        _Py_ForgetReference((PyObject*)type);
+        _Py_ForgetReference((PyObject *)type);
     }
 #endif
 
-    n_unnamed_members = 0;
-    for (i = 0; desc->fields[i].name != NULL; ++i)
-        if (desc->fields[i].name == PyStructSequence_UnnamedField)
-            n_unnamed_members++;
-    n_members = i;
+    /* PyTypeObject has already been initialized */
+    if (Py_REFCNT(type) != 0) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
 
-    memcpy(type, &_struct_sequence_template, sizeof(PyTypeObject));
-    type->tp_base = &PyTuple_Type;
     type->tp_name = desc->name;
+    type->tp_basicsize = sizeof(PyStructSequence) - sizeof(PyObject *);
+    type->tp_itemsize = sizeof(PyObject *);
+    type->tp_dealloc = (destructor)structseq_dealloc;
+    type->tp_repr = (reprfunc)structseq_repr;
     type->tp_doc = desc->doc;
+    type->tp_base = &PyTuple_Type;
+    type->tp_methods = structseq_methods;
+    type->tp_new = structseq_new;
+    type->tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC;
 
-    members = PyMem_NEW(PyMemberDef, n_members-n_unnamed_members+1);
+    n_members = count_members(desc, &n_unnamed_members);
+    members = PyMem_NEW(PyMemberDef, n_members - n_unnamed_members + 1);
     if (members == NULL) {
         PyErr_NoMemory();
         return -1;
     }
-
-    for (i = k = 0; i < n_members; ++i) {
-        if (desc->fields[i].name == PyStructSequence_UnnamedField)
-            continue;
-        members[k].name = desc->fields[i].name;
-        members[k].type = T_OBJECT;
-        members[k].offset = offsetof(PyStructSequence, ob_item)
-          + i * sizeof(PyObject*);
-        members[k].flags = READONLY;
-        members[k].doc = desc->fields[i].doc;
-        k++;
-    }
-    members[k].name = NULL;
-
+    initialize_members(desc, members, n_members);
     type->tp_members = members;
 
-    if (PyType_Ready(type) < 0)
+    if (PyType_Ready(type) < 0) {
+        PyMem_FREE(members);
         return -1;
+    }
     Py_INCREF(type);
 
-    dict = type->tp_dict;
-#define SET_DICT_FROM_SIZE(key, value)                          \
-    do {                                                        \
-        v = PyLong_FromSsize_t(value);                          \
-        if (v == NULL)                                          \
-            return -1;                                          \
-        if (PyDict_SetItemString(dict, key, v) < 0) {           \
-            Py_DECREF(v);                                       \
-            return -1;                                          \
-        }                                                       \
-        Py_DECREF(v);                                           \
-    } while (0)
-
-    SET_DICT_FROM_SIZE(visible_length_key, desc->n_in_sequence);
-    SET_DICT_FROM_SIZE(real_length_key, n_members);
-    SET_DICT_FROM_SIZE(unnamed_fields_key, n_unnamed_members);
+    if (initialize_structseq_dict(
+            desc, type->tp_dict, n_members, n_unnamed_members) < 0) {
+        PyMem_FREE(members);
+        Py_DECREF(type);
+        return -1;
+    }
 
     return 0;
 }
@@ -406,19 +409,57 @@ PyStructSequence_InitType(PyTypeObject *type, PyStructSequence_Desc *desc)
     (void)PyStructSequence_InitType2(type, desc);
 }
 
-PyTypeObject*
+PyTypeObject *
 PyStructSequence_NewType(PyStructSequence_Desc *desc)
 {
-    PyTypeObject *result;
-
-    result = (PyTypeObject*)PyType_GenericAlloc(&PyType_Type, 0);
-    if (result == NULL)
+    PyMemberDef *members;
+    PyObject *bases;
+    PyTypeObject *type;
+    PyType_Slot slots[7];
+    PyType_Spec spec;
+    Py_ssize_t n_members, n_unnamed_members;
+
+    /* Initialize MemberDefs */
+    n_members = count_members(desc, &n_unnamed_members);
+    members = PyMem_NEW(PyMemberDef, n_members - n_unnamed_members + 1);
+    if (members == NULL) {
+        PyErr_NoMemory();
         return NULL;
-    if (PyStructSequence_InitType2(result, desc) < 0) {
-        Py_DECREF(result);
+    }
+    initialize_members(desc, members, n_members);
+
+    /* Initialize Slots */
+    slots[0] = (PyType_Slot){Py_tp_dealloc, (destructor)structseq_dealloc};
+    slots[1] = (PyType_Slot){Py_tp_repr, (reprfunc)structseq_repr};
+    slots[2] = (PyType_Slot){Py_tp_doc, (void *)desc->doc};
+    slots[3] = (PyType_Slot){Py_tp_methods, structseq_methods};
+    slots[4] = (PyType_Slot){Py_tp_new, structseq_new};
+    slots[5] = (PyType_Slot){Py_tp_members, members};
+    slots[6] = (PyType_Slot){0, 0};
+
+    /* Initialize Spec */
+    /* The name in this PyType_Spec is statically allocated so it is */
+    /* expected that it'll outlive the PyType_Spec */
+    spec.name = desc->name;
+    spec.basicsize = sizeof(PyStructSequence) - sizeof(PyObject *);
+    spec.itemsize = sizeof(PyObject *);
+    spec.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC;
+    spec.slots = slots;
+
+    bases = PyTuple_Pack(1, &PyTuple_Type);
+    type = (PyTypeObject *)PyType_FromSpecWithBases(&spec, bases);
+    Py_DECREF(bases);
+    PyMem_FREE(members);
+    if (type == NULL) {
         return NULL;
     }
-    return result;
+
+    if (initialize_structseq_dict(
+            desc, type->tp_dict, n_members, n_unnamed_members) < 0) {
+        return NULL;
+    }
+
+    return type;
 }
 
 int _PyStructSequence_Init(void)