Repository: numpy/numpy
Pull request: MAINT: add PyArray_ZeroContiguousBuffer helper and use it in PyArray_Resize
Fixes #26859 

Problem statement:
BUG: crash after resize on a StringDType array
### Describe the issue:

Dear numpy team,
A simple `StringDType` array resize seems to corrupt memory:
```
    import numpy as np
    from numpy.dtypes import StringDType
    a = np.array([''], dtype=StringDType())
    a.resize(2) # corrupts memory at a[1]
    print(a[1]) # crashes or raises MemoryError or whatever
```
I think this is because `PyArray_Resize` fearlessly writes a PyObject pointer into the array when a dtype has the `NPY_ITEM_REFCOUNT` flag set:
```
    PyObject *zero = PyLong_FromLong(0);
    char* optr = PyArray_BYTES(self) + oldnbytes;
    Py_INCREF(zero);
    memcpy(optr, &zero, sizeof(zero));
```
I suspect `PyArray_Resize` should do this only when the dtype is `NPY_OBJECT`,
but I leave it to the experts who know all the implications.
Thanks!



### Reproduce the code example:

```python
import numpy as np
from numpy.dtypes import StringDType
a = np.array([''], dtype=StringDType())
a.resize(2)
print(a[1])
```


### Error message:

```shell
>>> print(a[1])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
MemoryError: Failed to load string in StringDType getitem
```


### Python and NumPy Versions:

numpy 2.0.0
python 3.12.3 (tags/v3.12.3:f6650f9, Apr  9 2024, 14:05:25) [MSC v.1938 64 bit (AMD64)]

### Runtime Environment:

_No response_

### Context for the issue:

_No response_

Hints from issue discussion:
Thanks for the report! See #26934 for a fix.

Commit message: BUG: properly zero-fill the buffer

Patch to review:
diff --git a/numpy/_core/src/multiarray/ctors.c b/numpy/_core/src/multiarray/ctors.c
index 8ee9d28c3086..5c1a78daf0c5 100644
--- a/numpy/_core/src/multiarray/ctors.c
+++ b/numpy/_core/src/multiarray/ctors.c
@@ -822,6 +822,12 @@ PyArray_NewFromDescr_int(
 
 
     if (data == NULL) {
+        /* This closely follows PyArray_ZeroContiguousBuffer. We can't use
+         *  that because here we need to allocate after checking if there is
+         *  custom zeroing logic and that function accepts an already-allocated
+         *  array
+         */
+
         /* float errors do not matter and we do not release GIL */
         NPY_ARRAYMETHOD_FLAGS zero_flags;
         PyArrayMethod_GetTraverseLoop *get_fill_zero_loop =
diff --git a/numpy/_core/src/multiarray/refcount.c b/numpy/_core/src/multiarray/refcount.c
index 1bc693532646..4fc19979bac5 100644
--- a/numpy/_core/src/multiarray/refcount.c
+++ b/numpy/_core/src/multiarray/refcount.c
@@ -18,8 +18,8 @@
 #include "iterators.h"
 #include "dtypemeta.h"
 #include "refcount.h"
-
 #include "npy_config.h"
+#include "templ_common.h" /* for npy_mul_sizes_with_overflow */
 
 
 
@@ -56,6 +56,61 @@ PyArray_ClearBuffer(
 }
 
 
+/*
+ * Helper function to zero an array buffer.
+ *
+ * Here "zeroing" means an abstract zeroing operation, which for an
+ * array of references might be something more complicated than
+ * zero-filling the buffer.
+ *
+ * Failure (returns -1) indicates some sort of programming or logical
+ * error and should not happen for a data type that has been set up
+ * correctly. In principle a sufficiently weird dtype might run out of
+ * memory but in practice this likely won't happen.
+ */
+NPY_NO_EXPORT int
+PyArray_ZeroContiguousBuffer(
+        PyArray_Descr *descr, char *data,
+        npy_intp stride, npy_intp size, int aligned)
+{
+    NPY_traverse_info zero_info;
+    NPY_traverse_info_init(&zero_info);
+    /* Flags unused: float errors do not matter and we do not release GIL */
+    NPY_ARRAYMETHOD_FLAGS flags_unused;
+    PyArrayMethod_GetTraverseLoop *get_fill_zero_loop =
+            NPY_DT_SLOTS(NPY_DTYPE(descr))->get_fill_zero_loop;
+    if (get_fill_zero_loop != NULL) {
+        if (get_fill_zero_loop(
+                    NULL, descr, aligned, descr->elsize, &(zero_info.func),
+                    &(zero_info.auxdata), &flags_unused) < 0) {
+            goto fail;
+        }
+    }
+    else {
+        npy_intp nbytes;
+        if (!npy_mul_sizes_with_overflow(&nbytes, size, stride)) {
+            memset(data, 0, nbytes);
+            NPY_traverse_info_xfree(&zero_info);
+            return 0;
+        }
+        else {
+            PyErr_SetString(PyExc_OverflowError,
+                            "Integer overflow in computing resized buffer size");
+            goto fail;
+        }
+    }
+
+    int res = zero_info.func(
+            NULL, descr, data, size, stride, zero_info.auxdata);
+    NPY_traverse_info_xfree(&zero_info);
+    return res;
+
+  fail:
+    NPY_traverse_info_xfree(&zero_info);
+    return -1;
+}
+
+
 /*
  * Helper function to clear whole array.  It seems plausible that we should
  * be able to get away with assuming the array is contiguous.
diff --git a/numpy/_core/src/multiarray/refcount.h b/numpy/_core/src/multiarray/refcount.h
index d9f472b2697e..41c428f321e4 100644
--- a/numpy/_core/src/multiarray/refcount.h
+++ b/numpy/_core/src/multiarray/refcount.h
@@ -6,6 +6,11 @@ PyArray_ClearBuffer(
         PyArray_Descr *descr, char *data,
         npy_intp stride, npy_intp size, int aligned);
 
+NPY_NO_EXPORT int
+PyArray_ZeroContiguousBuffer(
+        PyArray_Descr *descr, char *data,
+        npy_intp stride, npy_intp size, int aligned);
+
 NPY_NO_EXPORT int
 PyArray_ClearArray(PyArrayObject *arr);
 
diff --git a/numpy/_core/src/multiarray/shape.c b/numpy/_core/src/multiarray/shape.c
index c33272de4eb5..340fe7289ac8 100644
--- a/numpy/_core/src/multiarray/shape.c
+++ b/numpy/_core/src/multiarray/shape.c
@@ -7,22 +7,17 @@
 
 #include "numpy/arrayobject.h"
 #include "numpy/arrayscalars.h"
-
 #include "numpy/npy_math.h"
 
 #include "npy_config.h"
-
-
-
 #include "arraywrap.h"
 #include "ctors.h"
-
 #include "shape.h"
-
 #include "npy_static_data.h" /* for interned strings */
 #include "templ_common.h" /* for npy_mul_sizes_with_overflow */
 #include "common.h" /* for convert_shape_to_string */
 #include "alloc.h"
+#include "refcount.h"
 
 static int
 _fix_unknown_dimension(PyArray_Dims *newshape, PyArrayObject *arr);
@@ -31,9 +26,6 @@ static int
 _attempt_nocopy_reshape(PyArrayObject *self, int newnd, const npy_intp *newdims,
                         npy_intp *newstrides, int is_f_order);
 
-static void
-_putzero(char *optr, PyObject *zero, PyArray_Descr *dtype);
-
 /*NUMPY_API
  * Resize (reallocate data).  Only works if nothing else is referencing this
  * array and it is contiguous.  If refcheck is 0, then the reference count is
@@ -141,20 +133,14 @@ PyArray_Resize(PyArrayObject *self, PyArray_Dims *newshape, int refcheck,
     }
 
     if (newnbytes > oldnbytes && PyArray_ISWRITEABLE(self)) {
-        /* Fill new memory with zeros */
-        if (PyDataType_FLAGCHK(PyArray_DESCR(self), NPY_ITEM_REFCOUNT)) {
-            PyObject *zero = PyLong_FromLong(0);
-            char *optr;
-            optr = PyArray_BYTES(self) + oldnbytes;
-            npy_intp n_new = newsize - oldsize;
-            for (npy_intp i = 0; i < n_new; i++) {
-                _putzero((char *)optr, zero, PyArray_DESCR(self));
-                optr += elsize;
-            }
-            Py_DECREF(zero);
-        }
-        else{
-            memset(PyArray_BYTES(self) + oldnbytes, 0, newnbytes - oldnbytes);
+        /* Fill new memory with zeros (PyLong zero for object arrays) */
+        npy_intp stride = elsize;
+        npy_intp size = newsize - oldsize;
+        char *data = PyArray_BYTES(self) + oldnbytes;
+        int aligned = PyArray_ISALIGNED(self);
+        if (PyArray_ZeroContiguousBuffer(PyArray_DESCR(self), data,
+                                         stride, size, aligned) < 0) {
+            return NULL;
         }
     }
 
@@ -338,41 +324,6 @@ PyArray_Reshape(PyArrayObject *self, PyObject *shape)
 }
 
 
-static void
-_putzero(char *optr, PyObject *zero, PyArray_Descr *dtype)
-{
-    if (!PyDataType_FLAGCHK(dtype, NPY_ITEM_REFCOUNT)) {
-        memset(optr, 0, dtype->elsize);
-    }
-    else if (PyDataType_HASFIELDS(dtype)) {
-        PyObject *key, *value, *title = NULL;
-        PyArray_Descr *new;
-        int offset;
-        Py_ssize_t pos = 0;
-        while (PyDict_Next(PyDataType_FIELDS(dtype), &pos, &key, &value)) {
-            if (NPY_TITLE_KEY(key, value)) {
-                continue;
-            }
-            if (!PyArg_ParseTuple(value, "Oi|O", &new, &offset, &title)) {
-                return;
-            }
-            _putzero(optr + offset, zero, new);
-        }
-    }
-    else {
-        npy_intp i;
-        npy_intp nsize = dtype->elsize / sizeof(zero);
-
-        for (i = 0; i < nsize; i++) {
-            Py_INCREF(zero);
-            memcpy(optr, &zero, sizeof(zero));
-            optr += sizeof(zero);
-        }
-    }
-    return;
-}
-
-
 /*
  * attempt to reshape an array without copying data
  *
diff --git a/numpy/_core/tests/test_stringdtype.py b/numpy/_core/tests/test_stringdtype.py
index 9ff3224947d9..069a9c00d5de 100644
--- a/numpy/_core/tests/test_stringdtype.py
+++ b/numpy/_core/tests/test_stringdtype.py
@@ -11,7 +11,7 @@
 
 from numpy.dtypes import StringDType
 from numpy._core.tests._natype import pd_NA
-from numpy.testing import assert_array_equal, IS_WASM
+from numpy.testing import assert_array_equal, IS_WASM, IS_PYPY
 
 
 @pytest.fixture
@@ -511,6 +511,15 @@ def test_concatenate(string_list):
     assert_array_equal(np.concatenate([sarr], axis=0), sarr)
 
 
+def test_resize_method(string_list):
+    sarr = np.array(string_list, dtype="T")
+    if IS_PYPY:
+        sarr.resize(len(string_list)+3, refcheck=False)
+    else:
+        sarr.resize(len(string_list)+3)
+    assert_array_equal(sarr, np.array(string_list + ['']*3,  dtype="T"))
+
+
 def test_create_with_copy_none(string_list):
     arr = np.array(string_list, dtype=StringDType())
     # create another stringdtype array with an arena that has a different


Code review comments:
diff_hunk: @@ -56,6 +56,61 @@ PyArray_ClearBuffer(
 }
 
 
+/*
+ * Helper function to zero an array buffer.
+ *
+ * Here "zeroing" means an abstract zeroing operation, which for an
+ * array of references might be something more complicated than
+ * zero-filling the buffer.
+ *
+ * Failure (returns -1) indicates some sort of programming or logical
+ * error and should not happen for a data type that has been set up
+ * correctly. In principle a sufficiently weird dtype might run out of
+ * memory but in practice this likely won't happen.
+ */
+NPY_NO_EXPORT int
+PyArray_ZeroContiguousBuffer(
+        PyArray_Descr *descr, char *data,
+        npy_intp stride, npy_intp size, int aligned)
+{
+    NPY_traverse_info zero_info;
+    NPY_traverse_info_init(&zero_info);
+    /* Flags unused: float errors do not matter and we do not release GIL */
+    NPY_ARRAYMETHOD_FLAGS flags_unused;
+    PyArrayMethod_GetTraverseLoop *get_fill_zero_loop =
+            NPY_DT_SLOTS(NPY_DTYPE(descr))->get_fill_zero_loop;
+    if (get_fill_zero_loop != NULL) {
+        if (get_fill_zero_loop(
+                    NULL, descr, aligned, descr->elsize, &(zero_info.func),
+                    &(zero_info.auxdata), &flags_unused) < 0) {
+            goto fail;
+        }
+    }
+    else {
+        npy_intp nbytes;
+        if (!npy_mul_sizes_with_overflow(&nbytes, size, stride)) {
path: numpy/_core/src/multiarray/refcount.c
text: @user1:
I would be happier with _not_ checking it here, TBH (not because I mind it here, but because I don't want to clutter more code with cargo-culted checks).
A valid NumPy buffer cannot overflow this so if this happens there would be a missing check earlier.  Maybe an `assert` instead if you like to clarify it?

After all, just like ufuncs/other functions the `zero_info.func` is not and should not be required to check this as well.