Merge pull request #28032 from seberg/alloc-fixup

seberg · web-flow · commit 35b2c4a68e78 · 2024-12-19T08:39:07.000+01:00
BUG,MAINT: Fix size bug in new alloc helper and use it in one more place
diff --git a/numpy/_core/src/multiarray/alloc.h b/numpy/_core/src/multiarray/alloc.h
@@ -61,7 +61,7 @@ static inline void
 _npy_init_workspace(
     void **buf, void *static_buf, size_t static_buf_size, size_t elsize, size_t size)
 {
-    if (NPY_LIKELY(size <= static_buf_size / elsize)) {
+    if (NPY_LIKELY(size <= static_buf_size)) {
         *buf = static_buf;
     }
     else {
diff --git a/numpy/_core/src/multiarray/common_dtype.c b/numpy/_core/src/multiarray/common_dtype.c
@@ -7,6 +7,7 @@
 #include "numpy/npy_common.h"
 #include "numpy/arrayobject.h"
 
+#include "alloc.h"
 #include "convert_datatype.h"
 #include "dtypemeta.h"
 #include "abstractdtypes.h"
@@ -211,19 +212,10 @@ PyArray_PromoteDTypeSequence(
     PyArray_DTypeMeta *result = NULL;
 
     /* Copy dtypes so that we can reorder them (only allocate when many) */
-    PyObject *_scratch_stack[NPY_MAXARGS];
-    PyObject **_scratch_heap = NULL;
-    PyArray_DTypeMeta **dtypes = (PyArray_DTypeMeta **)_scratch_stack;
-
-    if (length > NPY_MAXARGS) {
-        _scratch_heap = PyMem_Malloc(length * sizeof(PyObject *));
-        if (_scratch_heap == NULL) {
-            PyErr_NoMemory();
-            return NULL;
-        }
-        dtypes = (PyArray_DTypeMeta **)_scratch_heap;
+    NPY_ALLOC_WORKSPACE(dtypes, PyArray_DTypeMeta *, 16, length);
+    if (dtypes == NULL) {
+        return NULL;
     }
-
     memcpy(dtypes, dtypes_in, length * sizeof(PyObject *));
 
     /*
@@ -311,6 +303,6 @@ PyArray_PromoteDTypeSequence(
     }
 
   finish:
-    PyMem_Free(_scratch_heap);
+    npy_free_workspace(dtypes);
     return result;
 }

Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ static inline void`
`61`	`61`	`_npy_init_workspace(`
`62`	`62`	`void *buf, void static_buf, size_t static_buf_size, size_t elsize, size_t size)`
`63`	`63`	`{`
`64`		`- if (NPY_LIKELY(size <= static_buf_size / elsize)) {`
	`64`	`+ if (NPY_LIKELY(size <= static_buf_size)) {`
`65`	`65`	`*buf = static_buf;`
`66`	`66`	`}`
`67`	`67`	`else {`