Issue #21233: Add new C functions: PyMem_RawCalloc(), PyMem_Calloc(),

PyObject_Calloc(), _PyObject_GC_Calloc(). bytes(int) and bytearray(int) are now
using ``calloc()`` instead of ``malloc()`` for large objects which is faster
and use less memory (until the bytearray buffer is filled with data).

3 files changed, 135 insertions, 19 deletions

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 750e90fa4c..a75500426c 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -2710,6 +2710,20 @@ test_pymem_alloc0(PyObject *self)
 {
     void *ptr;
 
+    ptr = PyMem_RawMalloc(0);
+    if (ptr == NULL) {
+        PyErr_SetString(PyExc_RuntimeError, "PyMem_RawMalloc(0) returns NULL");
+        return NULL;
+    }
+    PyMem_RawFree(ptr);
+
+    ptr = PyMem_RawCalloc(0, 0);
+    if (ptr == NULL) {
+        PyErr_SetString(PyExc_RuntimeError, "PyMem_RawCalloc(0, 0) returns NULL");
+        return NULL;
+    }
+    PyMem_RawFree(ptr);
+
     ptr = PyMem_Malloc(0);
     if (ptr == NULL) {
         PyErr_SetString(PyExc_RuntimeError, "PyMem_Malloc(0) returns NULL");
@@ -2717,6 +2731,13 @@ test_pymem_alloc0(PyObject *self)
     }
     PyMem_Free(ptr);
 
+    ptr = PyMem_Calloc(0, 0);
+    if (ptr == NULL) {
+        PyErr_SetString(PyExc_RuntimeError, "PyMem_Calloc(0, 0) returns NULL");
+        return NULL;
+    }
+    PyMem_Free(ptr);
+
     ptr = PyObject_Malloc(0);
     if (ptr == NULL) {
         PyErr_SetString(PyExc_RuntimeError, "PyObject_Malloc(0) returns NULL");
@@ -2724,6 +2745,13 @@ test_pymem_alloc0(PyObject *self)
     }
     PyObject_Free(ptr);
 
+    ptr = PyObject_Calloc(0, 0);
+    if (ptr == NULL) {
+        PyErr_SetString(PyExc_RuntimeError, "PyObject_Calloc(0, 0) returns NULL");
+        return NULL;
+    }
+    PyObject_Free(ptr);
+
     Py_RETURN_NONE;
 }
 
@@ -2731,6 +2759,8 @@ typedef struct {
     PyMemAllocator alloc;
 
     size_t malloc_size;
+    size_t calloc_nelem;
+    size_t calloc_elsize;
     void *realloc_ptr;
     size_t realloc_new_size;
     void *free_ptr;
@@ -2743,6 +2773,14 @@ static void* hook_malloc (void* ctx, size_t size)
     return hook->alloc.malloc(hook->alloc.ctx, size);
 }
 
+static void* hook_calloc (void* ctx, size_t nelem, size_t elsize)
+{
+    alloc_hook_t *hook = (alloc_hook_t *)ctx;
+    hook->calloc_nelem = nelem;
+    hook->calloc_elsize = elsize;
+    return hook->alloc.calloc(hook->alloc.ctx, nelem, elsize);
+}
+
 static void* hook_realloc (void* ctx, void* ptr, size_t new_size)
 {
     alloc_hook_t *hook = (alloc_hook_t *)ctx;
@@ -2765,16 +2803,14 @@ test_setallocators(PyMemAllocatorDomain domain)
     const char *error_msg;
     alloc_hook_t hook;
     PyMemAllocator alloc;
-    size_t size, size2;
+    size_t size, size2, nelem, elsize;
     void *ptr, *ptr2;
 
-    hook.malloc_size = 0;
-    hook.realloc_ptr = NULL;
-    hook.realloc_new_size = 0;
-    hook.free_ptr = NULL;
+    memset(&hook, 0, sizeof(hook));
 
     alloc.ctx = &hook;
     alloc.malloc = &hook_malloc;
+    alloc.calloc = &hook_calloc;
     alloc.realloc = &hook_realloc;
     alloc.free = &hook_free;
     PyMem_GetAllocator(domain, &hook.alloc);
@@ -2831,6 +2867,33 @@ test_setallocators(PyMemAllocatorDomain domain)
         goto fail;
     }
 
+    nelem = 2;
+    elsize = 5;
+    switch(domain)
+    {
+    case PYMEM_DOMAIN_RAW: ptr = PyMem_RawCalloc(nelem, elsize); break;
+    case PYMEM_DOMAIN_MEM: ptr = PyMem_Calloc(nelem, elsize); break;
+    case PYMEM_DOMAIN_OBJ: ptr = PyObject_Calloc(nelem, elsize); break;
+    default: ptr = NULL; break;
+    }
+
+    if (ptr == NULL) {
+        error_msg = "calloc failed";
+        goto fail;
+    }
+
+    if (hook.calloc_nelem != nelem || hook.calloc_elsize != elsize) {
+        error_msg = "calloc invalid nelem or elsize";
+        goto fail;
+    }
+
+    switch(domain)
+    {
+    case PYMEM_DOMAIN_RAW: PyMem_RawFree(ptr); break;
+    case PYMEM_DOMAIN_MEM: PyMem_Free(ptr); break;
+    case PYMEM_DOMAIN_OBJ: PyObject_Free(ptr); break;
+    }
+
     Py_INCREF(Py_None);
     res = Py_None;
     goto finally;
diff --git a/Modules/_tracemalloc.c b/Modules/_tracemalloc.c
index 780e8eda82..e254a9029d 100644
--- a/Modules/_tracemalloc.c
+++ b/Modules/_tracemalloc.c
@@ -476,17 +476,22 @@ tracemalloc_remove_trace(void *ptr)
 }
 
 static void*
-tracemalloc_malloc(void *ctx, size_t size)
+tracemalloc_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
 {
     PyMemAllocator *alloc = (PyMemAllocator *)ctx;
     void *ptr;
 
-    ptr = alloc->malloc(alloc->ctx, size);
+    assert(nelem <= PY_SIZE_MAX / elsize);
+
+    if (use_calloc)
+        ptr = alloc->calloc(alloc->ctx, nelem, elsize);
+    else
+        ptr = alloc->malloc(alloc->ctx, nelem * elsize);
     if (ptr == NULL)
         return NULL;
 
     TABLES_LOCK();
-    if (tracemalloc_add_trace(ptr, size) < 0) {
+    if (tracemalloc_add_trace(ptr, nelem * elsize) < 0) {
         /* Failed to allocate a trace for the new memory block */
         TABLES_UNLOCK();
         alloc->free(alloc->ctx, ptr);
@@ -560,13 +565,16 @@ tracemalloc_free(void *ctx, void *ptr)
 }
 
 static void*
-tracemalloc_malloc_gil(void *ctx, size_t size)
+tracemalloc_alloc_gil(int use_calloc, void *ctx, size_t nelem, size_t elsize)
 {
     void *ptr;
 
     if (get_reentrant()) {
         PyMemAllocator *alloc = (PyMemAllocator *)ctx;
-        return alloc->malloc(alloc->ctx, size);
+        if (use_calloc)
+            return alloc->calloc(alloc->ctx, nelem, elsize);
+        else
+            return alloc->malloc(alloc->ctx, nelem * elsize);
     }
 
     /* Ignore reentrant call. PyObjet_Malloc() calls PyMem_Malloc() for
@@ -574,13 +582,25 @@ tracemalloc_malloc_gil(void *ctx, size_t size)
        allocation twice. */
     set_reentrant(1);
 
-    ptr = tracemalloc_malloc(ctx, size);
+    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
 
     set_reentrant(0);
     return ptr;
 }
 
 static void*
+tracemalloc_malloc_gil(void *ctx, size_t size)
+{
+    return tracemalloc_alloc_gil(0, ctx, 1, size);
+}
+
+static void*
+tracemalloc_calloc_gil(void *ctx, size_t nelem, size_t elsize)
+{
+    return tracemalloc_alloc_gil(1, ctx, nelem, elsize);
+}
+
+static void*
 tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
 {
     void *ptr2;
@@ -614,7 +634,7 @@ tracemalloc_realloc_gil(void *ctx, void *ptr, size_t new_size)
 
 #ifdef TRACE_RAW_MALLOC
 static void*
-tracemalloc_raw_malloc(void *ctx, size_t size)
+tracemalloc_raw_alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
 {
 #ifdef WITH_THREAD
     PyGILState_STATE gil_state;
@@ -623,7 +643,10 @@ tracemalloc_raw_malloc(void *ctx, size_t size)
 
     if (get_reentrant()) {
         PyMemAllocator *alloc = (PyMemAllocator *)ctx;
-        return alloc->malloc(alloc->ctx, size);
+        if (use_calloc)
+            return alloc->calloc(alloc->ctx, nelem, elsize);
+        else
+            return alloc->malloc(alloc->ctx, nelem * elsize);
     }
 
     /* Ignore reentrant call. PyGILState_Ensure() may call PyMem_RawMalloc()
@@ -633,10 +656,10 @@ tracemalloc_raw_malloc(void *ctx, size_t size)
 
 #ifdef WITH_THREAD
     gil_state = PyGILState_Ensure();
-    ptr = tracemalloc_malloc(ctx, size);
+    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
     PyGILState_Release(gil_state);
 #else
-    ptr = tracemalloc_malloc(ctx, size);
+    ptr = tracemalloc_alloc(use_calloc, ctx, nelem, elsize);
 #endif
 
     set_reentrant(0);
@@ -644,6 +667,18 @@ tracemalloc_raw_malloc(void *ctx, size_t size)
 }
 
 static void*
+tracemalloc_raw_malloc(void *ctx, size_t size)
+{
+    return tracemalloc_raw_alloc(0, ctx, 1, size);
+}
+
+static void*
+tracemalloc_raw_calloc(void *ctx, size_t nelem, size_t elsize)
+{
+    return tracemalloc_raw_alloc(1, ctx, nelem, elsize);
+}
+
+static void*
 tracemalloc_raw_realloc(void *ctx, void *ptr, size_t new_size)
 {
 #ifdef WITH_THREAD
@@ -856,6 +891,7 @@ tracemalloc_start(int max_nframe)
 
 #ifdef TRACE_RAW_MALLOC
     alloc.malloc = tracemalloc_raw_malloc;
+    alloc.calloc = tracemalloc_raw_calloc;
     alloc.realloc = tracemalloc_raw_realloc;
     alloc.free = tracemalloc_free;
 
@@ -865,6 +901,7 @@ tracemalloc_start(int max_nframe)
 #endif
 
     alloc.malloc = tracemalloc_malloc_gil;
+    alloc.calloc = tracemalloc_calloc_gil;
     alloc.realloc = tracemalloc_realloc_gil;
     alloc.free = tracemalloc_free;
 
diff --git a/Modules/gcmodule.c b/Modules/gcmodule.c
index 6281a7c343..cff5d097af 100644
--- a/Modules/gcmodule.c
+++ b/Modules/gcmodule.c
@@ -1703,15 +1703,19 @@ PyObject_GC_UnTrack(void *op)
         _PyObject_GC_UNTRACK(op);
 }
 
-PyObject *
-_PyObject_GC_Malloc(size_t basicsize)
+static PyObject *
+_PyObject_GC_Alloc(int use_calloc, size_t basicsize)
 {
     PyObject *op;
     PyGC_Head *g;
+    size_t size;
     if (basicsize > PY_SSIZE_T_MAX - sizeof(PyGC_Head))
         return PyErr_NoMemory();
-    g = (PyGC_Head *)PyObject_MALLOC(
-        sizeof(PyGC_Head) + basicsize);
+    size = sizeof(PyGC_Head) + basicsize;
+    if (use_calloc)
+        g = (PyGC_Head *)PyObject_Calloc(1, size);
+    else
+        g = (PyGC_Head *)PyObject_Malloc(size);
     if (g == NULL)
         return PyErr_NoMemory();
     g->gc.gc_refs = 0;
@@ -1731,6 +1735,18 @@ _PyObject_GC_Malloc(size_t basicsize)
 }
 
 PyObject *
+_PyObject_GC_Malloc(size_t basicsize)
+{
+    return _PyObject_GC_Alloc(0, basicsize);
+}
+
+PyObject *
+_PyObject_GC_Calloc(size_t basicsize)
+{
+    return _PyObject_GC_Alloc(1, basicsize);
+}
+
+PyObject *
 _PyObject_GC_New(PyTypeObject *tp)
 {
     PyObject *op = _PyObject_GC_Malloc(_PyObject_SIZE(tp));