bpo-30789: Use a single memory block for co_extra. (#2555)

* bpo-30789: Use a single memory block for co_extra.

* Address review comments.

1 files changed, 15 insertions, 39 deletions

diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index 1979606..de8d16c 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -10,7 +10,7 @@
 /* Holder for co_extra information */
 typedef struct {
     Py_ssize_t ce_size;
-    void **ce_extras;
+    void *ce_extras[1];
 } _PyCodeObjectExtra;
 
 /* all_name_chars(s): true iff all chars in s are valid NAME_CHARS */
@@ -427,7 +427,6 @@ code_dealloc(PyCodeObject *co)
             }
         }
 
-        PyMem_Free(co_extra->ce_extras);
         PyMem_Free(co_extra);
     }
 
@@ -455,12 +454,13 @@ code_sizeof(PyCodeObject *co, void *unused)
     Py_ssize_t res = _PyObject_SIZE(Py_TYPE(co));
     _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra*) co->co_extra;
 
-    if (co->co_cell2arg != NULL && co->co_cellvars != NULL)
+    if (co->co_cell2arg != NULL && co->co_cellvars != NULL) {
         res += PyTuple_GET_SIZE(co->co_cellvars) * sizeof(Py_ssize_t);
-
-    if (co_extra != NULL)
-        res += co_extra->ce_size * sizeof(co_extra->ce_extras[0]);
-
+    }
+    if (co_extra != NULL) {
+        res += sizeof(_PyCodeObjectExtra) +
+               (co_extra->ce_size-1) * sizeof(co_extra->ce_extras[0]);
+    }
     return PyLong_FromSsize_t(res);
 }
 
@@ -863,48 +863,24 @@ _PyCode_SetExtra(PyObject *code, Py_ssize_t index, void *extra)
     PyCodeObject *o = (PyCodeObject*) code;
     _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra *) o->co_extra;
 
-    if (co_extra == NULL) {
-        co_extra = PyMem_Malloc(sizeof(_PyCodeObjectExtra));
+    if (co_extra == NULL || co_extra->ce_size <= index) {
+        Py_ssize_t i = (co_extra == NULL ? 0 : co_extra->ce_size);
+        co_extra = PyMem_Realloc(
+                co_extra,
+                sizeof(_PyCodeObjectExtra) +
+                (interp->co_extra_user_count-1) * sizeof(void*));
         if (co_extra == NULL) {
             return -1;
         }
-
-        co_extra->ce_extras = PyMem_Malloc(
-            interp->co_extra_user_count * sizeof(void*));
-        if (co_extra->ce_extras == NULL) {
-            PyMem_Free(co_extra);
-            return -1;
-        }
-
-        co_extra->ce_size = interp->co_extra_user_count;
-
-        for (Py_ssize_t i = 0; i < co_extra->ce_size; i++) {
+        for (; i < interp->co_extra_user_count; i++) {
             co_extra->ce_extras[i] = NULL;
         }
-
-        o->co_extra = co_extra;
-    }
-    else if (co_extra->ce_size <= index) {
-        void** ce_extras = PyMem_Realloc(
-            co_extra->ce_extras, interp->co_extra_user_count * sizeof(void*));
-
-        if (ce_extras == NULL) {
-            return -1;
-        }
-
-        for (Py_ssize_t i = co_extra->ce_size;
-             i < interp->co_extra_user_count;
-             i++) {
-            ce_extras[i] = NULL;
-        }
-
-        co_extra->ce_extras = ce_extras;
         co_extra->ce_size = interp->co_extra_user_count;
+        o->co_extra = co_extra;
     }
 
     if (co_extra->ce_extras[index] != NULL) {
         freefunc free = interp->co_extra_freefuncs[index];
-
         if (free != NULL) {
             free(co_extra->ce_extras[index]);
         }