1 files changed, 117 insertions, 25 deletions

diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c
index 3c33255..7cc889f 100644
--- a/Objects/obmalloc.c
+++ b/Objects/obmalloc.c
@@ -5,6 +5,7 @@
 #ifdef PYMALLOC_DEBUG   /* WITH_PYMALLOC && PYMALLOC_DEBUG */
 /* Forward declaration */
 static void* _PyMem_DebugMalloc(void *ctx, size_t size);
+static void* _PyMem_DebugCalloc(void *ctx, size_t nelem, size_t elsize);
 static void _PyMem_DebugFree(void *ctx, void *p);
 static void* _PyMem_DebugRealloc(void *ctx, void *ptr, size_t size);
 
@@ -43,6 +44,7 @@ static void _PyMem_DebugCheckAddress(char api_id, const void *p);
 
 /* Forward declaration */
 static void* _PyObject_Malloc(void *ctx, size_t size);
+static void* _PyObject_Calloc(void *ctx, size_t nelem, size_t elsize);
 static void _PyObject_Free(void *ctx, void *p);
 static void* _PyObject_Realloc(void *ctx, void *ptr, size_t size);
 #endif
@@ -51,7 +53,7 @@ static void* _PyObject_Realloc(void *ctx, void *ptr, size_t size);
 static void *
 _PyMem_RawMalloc(void *ctx, size_t size)
 {
-    /* PyMem_Malloc(0) means malloc(1). Some systems would return NULL
+    /* PyMem_RawMalloc(0) means malloc(1). Some systems would return NULL
        for malloc(0), which would be treated as an error. Some platforms would
        return a pointer with no memory behind it, which would break pymalloc.
        To solve these problems, allocate an extra byte. */
@@ -61,6 +63,20 @@ _PyMem_RawMalloc(void *ctx, size_t size)
 }
 
 static void *
+_PyMem_RawCalloc(void *ctx, size_t nelem, size_t elsize)
+{
+    /* PyMem_RawCalloc(0, 0) means calloc(1, 1). Some systems would return NULL
+       for calloc(0, 0), which would be treated as an error. Some platforms
+       would return a pointer with no memory behind it, which would break
+       pymalloc.  To solve these problems, allocate an extra byte. */
+    if (nelem == 0 || elsize == 0) {
+        nelem = 1;
+        elsize = 1;
+    }
+    return calloc(nelem, elsize);
+}
+
+static void *
 _PyMem_RawRealloc(void *ctx, void *ptr, size_t size)
 {
     if (size == 0)
@@ -123,9 +139,9 @@ _PyObject_ArenaFree(void *ctx, void *ptr, size_t size)
 #endif
 
 
-#define PYRAW_FUNCS _PyMem_RawMalloc, _PyMem_RawRealloc, _PyMem_RawFree
+#define PYRAW_FUNCS _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree
 #ifdef WITH_PYMALLOC
-#  define PYOBJ_FUNCS _PyObject_Malloc, _PyObject_Realloc, _PyObject_Free
+#  define PYOBJ_FUNCS _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free
 #else
 #  define PYOBJ_FUNCS PYRAW_FUNCS
 #endif
@@ -135,7 +151,7 @@ _PyObject_ArenaFree(void *ctx, void *ptr, size_t size)
 typedef struct {
     /* We tag each block with an API ID in order to tag API violations */
     char api_id;
-    PyMemAllocator alloc;
+    PyMemAllocatorEx alloc;
 } debug_alloc_api_t;
 static struct {
     debug_alloc_api_t raw;
@@ -147,10 +163,10 @@ static struct {
     {'o', {NULL, PYOBJ_FUNCS}}
     };
 
-#define PYDBG_FUNCS _PyMem_DebugMalloc, _PyMem_DebugRealloc, _PyMem_DebugFree
+#define PYDBG_FUNCS _PyMem_DebugMalloc, _PyMem_DebugCalloc, _PyMem_DebugRealloc, _PyMem_DebugFree
 #endif
 
-static PyMemAllocator _PyMem_Raw = {
+static PyMemAllocatorEx _PyMem_Raw = {
 #ifdef PYMALLOC_DEBUG
     &_PyMem_Debug.raw, PYDBG_FUNCS
 #else
@@ -158,7 +174,7 @@ static PyMemAllocator _PyMem_Raw = {
 #endif
     };
 
-static PyMemAllocator _PyMem = {
+static PyMemAllocatorEx _PyMem = {
 #ifdef PYMALLOC_DEBUG
     &_PyMem_Debug.mem, PYDBG_FUNCS
 #else
@@ -166,7 +182,7 @@ static PyMemAllocator _PyMem = {
 #endif
     };
 
-static PyMemAllocator _PyObject = {
+static PyMemAllocatorEx _PyObject = {
 #ifdef PYMALLOC_DEBUG
     &_PyMem_Debug.obj, PYDBG_FUNCS
 #else
@@ -193,9 +209,10 @@ void
 PyMem_SetupDebugHooks(void)
 {
 #ifdef PYMALLOC_DEBUG
-    PyMemAllocator alloc;
+    PyMemAllocatorEx alloc;
 
     alloc.malloc = _PyMem_DebugMalloc;
+    alloc.calloc = _PyMem_DebugCalloc;
     alloc.realloc = _PyMem_DebugRealloc;
     alloc.free = _PyMem_DebugFree;
 
@@ -220,7 +237,7 @@ PyMem_SetupDebugHooks(void)
 }
 
 void
-PyMem_GetAllocator(PyMemAllocatorDomain domain, PyMemAllocator *allocator)
+PyMem_GetAllocator(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)
 {
     switch(domain)
     {
@@ -228,16 +245,17 @@ PyMem_GetAllocator(PyMemAllocatorDomain domain, PyMemAllocator *allocator)
     case PYMEM_DOMAIN_MEM: *allocator = _PyMem; break;
     case PYMEM_DOMAIN_OBJ: *allocator = _PyObject; break;
     default:
-        /* unknown domain */
+        /* unknown domain: set all attributes to NULL */
         allocator->ctx = NULL;
         allocator->malloc = NULL;
+        allocator->calloc = NULL;
         allocator->realloc = NULL;
         allocator->free = NULL;
     }
 }
 
 void
-PyMem_SetAllocator(PyMemAllocatorDomain domain, PyMemAllocator *allocator)
+PyMem_SetAllocator(PyMemAllocatorDomain domain, PyMemAllocatorEx *allocator)
 {
     switch(domain)
     {
@@ -272,10 +290,18 @@ PyMem_RawMalloc(size_t size)
      */
     if (size > (size_t)PY_SSIZE_T_MAX)
         return NULL;
-
     return _PyMem_Raw.malloc(_PyMem_Raw.ctx, size);
 }
 
+void *
+PyMem_RawCalloc(size_t nelem, size_t elsize)
+{
+    /* see PyMem_RawMalloc() */
+    if (elsize != 0 && nelem > (size_t)PY_SSIZE_T_MAX / elsize)
+        return NULL;
+    return _PyMem_Raw.calloc(_PyMem_Raw.ctx, nelem, elsize);
+}
+
 void*
 PyMem_RawRealloc(void *ptr, size_t new_size)
 {
@@ -300,6 +326,15 @@ PyMem_Malloc(size_t size)
 }
 
 void *
+PyMem_Calloc(size_t nelem, size_t elsize)
+{
+    /* see PyMem_RawMalloc() */
+    if (elsize != 0 && nelem > (size_t)PY_SSIZE_T_MAX / elsize)
+        return NULL;
+    return _PyMem.calloc(_PyMem.ctx, nelem, elsize);
+}
+
+void *
 PyMem_Realloc(void *ptr, size_t new_size)
 {
     /* see PyMem_RawMalloc() */
@@ -352,6 +387,15 @@ PyObject_Malloc(size_t size)
 }
 
 void *
+PyObject_Calloc(size_t nelem, size_t elsize)
+{
+    /* see PyMem_RawMalloc() */
+    if (elsize != 0 && nelem > (size_t)PY_SSIZE_T_MAX / elsize)
+        return NULL;
+    return _PyObject.calloc(_PyObject.ctx, nelem, elsize);
+}
+
+void *
 PyObject_Realloc(void *ptr, size_t new_size)
 {
     /* see PyMem_RawMalloc() */
@@ -1122,8 +1166,9 @@ int Py_ADDRESS_IN_RANGE(void *P, poolp pool) Py_NO_INLINE;
  */
 
 static void *
-_PyObject_Malloc(void *ctx, size_t nbytes)
+_PyObject_Alloc(int use_calloc, void *ctx, size_t nelem, size_t elsize)
 {
+    size_t nbytes;
     block *bp;
     poolp pool;
     poolp next;
@@ -1131,6 +1176,9 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
 
     _Py_AllocatedBlocks++;
 
+    assert(nelem <= PY_SSIZE_T_MAX / elsize);
+    nbytes = nelem * elsize;
+
 #ifdef WITH_VALGRIND
     if (UNLIKELY(running_on_valgrind == -1))
         running_on_valgrind = RUNNING_ON_VALGRIND;
@@ -1138,9 +1186,9 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
         goto redirect;
 #endif
 
-    /*
-     * This implicitly redirects malloc(0).
-     */
+    if (nelem == 0 || elsize == 0)
+        goto redirect;
+
     if ((nbytes - 1) < SMALL_REQUEST_THRESHOLD) {
         LOCK();
         /*
@@ -1158,6 +1206,8 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
             assert(bp != NULL);
             if ((pool->freeblock = *(block **)bp) != NULL) {
                 UNLOCK();
+                if (use_calloc)
+                    memset(bp, 0, nbytes);
                 return (void *)bp;
             }
             /*
@@ -1170,6 +1220,8 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
                 pool->nextoffset += INDEX2SIZE(size);
                 *(block **)(pool->freeblock) = NULL;
                 UNLOCK();
+                if (use_calloc)
+                    memset(bp, 0, nbytes);
                 return (void *)bp;
             }
             /* Pool is full, unlink from used pools. */
@@ -1178,6 +1230,8 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
             next->prevpool = pool;
             pool->nextpool = next;
             UNLOCK();
+            if (use_calloc)
+                memset(bp, 0, nbytes);
             return (void *)bp;
         }
 
@@ -1257,6 +1311,8 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
                 assert(bp != NULL);
                 pool->freeblock = *(block **)bp;
                 UNLOCK();
+                if (use_calloc)
+                    memset(bp, 0, nbytes);
                 return (void *)bp;
             }
             /*
@@ -1272,6 +1328,8 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
             pool->freeblock = bp + size;
             *(block **)(pool->freeblock) = NULL;
             UNLOCK();
+            if (use_calloc)
+                memset(bp, 0, nbytes);
             return (void *)bp;
         }
 
@@ -1281,7 +1339,7 @@ _PyObject_Malloc(void *ctx, size_t nbytes)
         pool = (poolp)usable_arenas->pool_address;
         assert((block*)pool <= (block*)usable_arenas->address +
                                ARENA_SIZE - POOL_SIZE);
-        pool->arenaindex = usable_arenas - arenas;
+        pool->arenaindex = (uint)(usable_arenas - arenas);
         assert(&arenas[pool->arenaindex] == usable_arenas);
         pool->szidx = DUMMY_SIZE_IDX;
         usable_arenas->pool_address += POOL_SIZE;
@@ -1311,13 +1369,29 @@ redirect:
      * has been reached.
      */
     {
-        void *result = PyMem_RawMalloc(nbytes);
+        void *result;
+        if (use_calloc)
+            result = PyMem_RawCalloc(nelem, elsize);
+        else
+            result = PyMem_RawMalloc(nbytes);
         if (!result)
             _Py_AllocatedBlocks--;
         return result;
     }
 }
 
+static void *
+_PyObject_Malloc(void *ctx, size_t nbytes)
+{
+    return _PyObject_Alloc(0, ctx, 1, nbytes);
+}
+
+static void *
+_PyObject_Calloc(void *ctx, size_t nelem, size_t elsize)
+{
+    return _PyObject_Alloc(1, ctx, nelem, elsize);
+}
+
 /* free */
 
 ATTRIBUTE_NO_ADDRESS_SAFETY_ANALYSIS
@@ -1561,7 +1635,7 @@ _PyObject_Realloc(void *ctx, void *p, size_t nbytes)
 #endif
 
     if (p == NULL)
-        return _PyObject_Malloc(ctx, nbytes);
+        return _PyObject_Alloc(0, ctx, 1, nbytes);
 
 #ifdef WITH_VALGRIND
     /* Treat running_on_valgrind == -1 the same as 0 */
@@ -1589,7 +1663,7 @@ _PyObject_Realloc(void *ctx, void *p, size_t nbytes)
             }
             size = nbytes;
         }
-        bp = _PyObject_Malloc(ctx, nbytes);
+        bp = _PyObject_Alloc(0, ctx, 1, nbytes);
         if (bp != NULL) {
             memcpy(bp, p, size);
             _PyObject_Free(ctx, p);
@@ -1745,7 +1819,7 @@ p[2*S+n+S: 2*S+n+2*S]
 */
 
 static void *
-_PyMem_DebugMalloc(void *ctx, size_t nbytes)
+_PyMem_DebugAlloc(int use_calloc, void *ctx, size_t nbytes)
 {
     debug_alloc_api_t *api = (debug_alloc_api_t *)ctx;
     uchar *p;           /* base address of malloc'ed block */
@@ -1758,7 +1832,10 @@ _PyMem_DebugMalloc(void *ctx, size_t nbytes)
         /* overflow:  can't represent total as a Py_ssize_t */
         return NULL;
 
-    p = (uchar *)api->alloc.malloc(api->alloc.ctx, total);
+    if (use_calloc)
+        p = (uchar *)api->alloc.calloc(api->alloc.ctx, 1, total);
+    else
+        p = (uchar *)api->alloc.malloc(api->alloc.ctx, total);
     if (p == NULL)
         return NULL;
 
@@ -1767,7 +1844,7 @@ _PyMem_DebugMalloc(void *ctx, size_t nbytes)
     p[SST] = (uchar)api->api_id;
     memset(p + SST + 1, FORBIDDENBYTE, SST-1);
 
-    if (nbytes > 0)
+    if (nbytes > 0 && !use_calloc)
         memset(p + 2*SST, CLEANBYTE, nbytes);
 
     /* at tail, write pad (SST bytes) and serialno (SST bytes) */
@@ -1778,6 +1855,21 @@ _PyMem_DebugMalloc(void *ctx, size_t nbytes)
     return p + 2*SST;
 }
 
+static void *
+_PyMem_DebugMalloc(void *ctx, size_t nbytes)
+{
+    return _PyMem_DebugAlloc(0, ctx, nbytes);
+}
+
+static void *
+_PyMem_DebugCalloc(void *ctx, size_t nelem, size_t elsize)
+{
+    size_t nbytes;
+    assert(elsize == 0 || nelem <= PY_SSIZE_T_MAX / elsize);
+    nbytes = nelem * elsize;
+    return _PyMem_DebugAlloc(1, ctx, nbytes);
+}
+
 /* The debug free first checks the 2*SST bytes on each end for sanity (in
    particular, that the FORBIDDENBYTEs with the api ID are still intact).
    Then fills the original bytes with DEADBYTE.
@@ -1811,7 +1903,7 @@ _PyMem_DebugRealloc(void *ctx, void *p, size_t nbytes)
     int i;
 
     if (p == NULL)
-        return _PyMem_DebugMalloc(ctx, nbytes);
+        return _PyMem_DebugAlloc(0, ctx, nbytes);
 
     _PyMem_DebugCheckAddress(api->api_id, p);
     bumpserialno();