MAINT: refactor inner loop in test

3 files changed, 60 insertions, 46 deletions

diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index f04de088f..430c1873e 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -111,6 +111,24 @@ typedef int (PyUFunc_MaskedInnerLoopSelectionFunc)(
                             NpyAuxData **out_innerloopdata,
                             int *out_needs_api);
 
+typedef struct _ufunc_extension {
+        /* New in PyUFuncObject version 1 and above */
+
+        /*
+         * for each core_num_dim_ix distinct dimension names,
+         * the possible "frozen" size (-1 if not frozen).
+         */
+        npy_intp *core_dim_sizes;
+
+        /*
+         * for each distinct core dimension, a set of flags OR'd together
+         * e.g., UFUNC_CORE_DIM_CAN_IGNORE if signature has ?
+         *       UFUNC_CORE_DIM_SIZE_UNSET for non-frozen dimensions.
+         */
+        npy_uint32 *core_dim_flags;
+
+} ufunc_extension;
+
 typedef struct _tagPyUFuncObject {
         PyObject_HEAD
         /*
@@ -188,9 +206,10 @@ typedef struct _tagPyUFuncObject {
         /*
          * This was blocked off to be the "new" inner loop selector in 1.7,
          * but this was never implemented. (This is also why the above
-         * selector is called the "legacy" selector.)
+         * selector is called the "legacy" selector.) Repurposed in 1.16 to
+         * point to s_extension
          */
-        void *reserved2;
+        ufunc_extension * extension;
         /*
          * A function which returns a masked inner loop for the ufunc.
          */
@@ -210,20 +229,7 @@ typedef struct _tagPyUFuncObject {
          */
         npy_uint32 iter_flags;
 
-        /* New in version 1 and above */
-
-        /*
-         * for each core_num_dim_ix distinct dimension names,
-         * the possible "frozen" size (-1 if not frozen).
-         */
-        npy_intp *core_dim_sizes;
-
-        /*
-         * for each distinct core dimension, a set of flags OR'd together
-         * e.g., UFUNC_CORE_DIM_CAN_IGNORE if signature has ?
-         *       UFUNC_CORE_DIM_SIZE_UNSET for non-frozen dimensions.
-         */
-        npy_uint32 *core_dim_flags;
+        ufunc_extension s_extension;
 
 } PyUFuncObject;
 
diff --git a/numpy/core/src/umath/_umath_tests.c.src b/numpy/core/src/umath/_umath_tests.c.src
index 36f2d52d6..e03ab12ec 100644
--- a/numpy/core/src/umath/_umath_tests.c.src
+++ b/numpy/core/src/umath/_umath_tests.c.src
@@ -217,16 +217,20 @@ static void
     INIT_OUTER_LOOP_3
     npy_intp is1=steps[0], is2=steps[1], os = steps[2];
     BEGIN_OUTER_LOOP_3
-        char *ip1=args[0], *ip2=args[1], *op=args[2];
+        @typ@ i1_x = *(@typ@ *)(args[0] + 0*is1);
+        @typ@ i1_y = *(@typ@ *)(args[0] + 1*is1);
+        @typ@ i1_z = *(@typ@ *)(args[0] + 2*is1);
+
+        @typ@ i2_x = *(@typ@ *)(args[1] + 0*is2);
+        @typ@ i2_y = *(@typ@ *)(args[1] + 1*is2);
+        @typ@ i2_z = *(@typ@ *)(args[1] + 2*is2);
+        char *op = args[2];
 
-        *(@typ@ *)op = *(@typ@ *)(ip1 + is1) * *(@typ@ *)(ip2 + 2*is2) -
-                       *(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)(ip2 + is2);
+        *(@typ@ *)op = i1_y * i2_z - i1_z * i2_y;
         op += os;
-        *(@typ@ *)op = *(@typ@ *)(ip1 + 2*is1) * *(@typ@ *)ip2 -
-                       *(@typ@ *)ip1 * *(@typ@ *)(ip2 + 2*is2);
+        *(@typ@ *)op = i1_z * i2_x - i1_x * i2_z;
         op += os;
-        *(@typ@ *)op = *(@typ@ *)ip1 * *(@typ@ *)(ip2 + is2) -
-                       *(@typ@ *)(ip1 + is1) * *(@typ@ *)ip2;
+        *(@typ@ *)op = i1_x * i2_y - i1_y * i2_x;
     END_OUTER_LOOP
 }
 
@@ -531,13 +535,13 @@ UMath_Tests_test_signature(PyObject *NPY_UNUSED(dummy), PyObject *args)
         Py_INCREF(Py_None);
         core_dim_ixs = Py_None;
     }
-    if (f->core_dim_flags != NULL) {
+    if (f->extension->core_dim_flags != NULL) {
         core_dim_flags = PyTuple_New(f->core_num_dim_ix);
         if (core_dim_flags == NULL) {
             goto fail;
         }
         for (i = 0; i < f->core_num_dim_ix; i++) {
-            PyObject * val = PyLong_FromLong(f->core_dim_flags[i]);
+            PyObject * val = PyLong_FromLong(f->extension->core_dim_flags[i]);
             PyTuple_SET_ITEM(core_dim_flags, i, val);
         }
     }
@@ -545,13 +549,13 @@ UMath_Tests_test_signature(PyObject *NPY_UNUSED(dummy), PyObject *args)
         Py_INCREF(Py_None);
         core_dim_flags = Py_None;
     }
-    if (f->core_dim_sizes != NULL) {
+    if (f->extension->core_dim_sizes != NULL) {
         core_dim_sizes = PyTuple_New(f->core_num_dim_ix);
         if (core_dim_sizes == NULL) {
             goto fail;
         }
         for (i = 0; i < f->core_num_dim_ix; i++) {
-            PyObject * val = PyLong_FromLong(f->core_dim_sizes[i]);
+            PyObject * val = PyLong_FromLong(f->extension->core_dim_sizes[i]);
             PyTuple_SET_ITEM(core_dim_sizes, i, val);
         }
     }
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 09344f602..9f2bb3fb6 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -572,17 +572,18 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
     ufunc->core_offsets = PyArray_malloc(sizeof(int) * ufunc->nargs);
     /* The next three items will be shrunk later */
     ufunc->core_dim_ixs = PyArray_malloc(sizeof(int) * len);
-    ufunc->core_dim_flags = PyArray_malloc(sizeof(npy_uint32) * len);
-    ufunc->core_dim_sizes = PyArray_malloc(sizeof(npy_intp) * len);
+    ufunc->extension->core_dim_sizes = PyArray_malloc(sizeof(npy_intp) * len);
+    ufunc->extension->core_dim_flags = PyArray_malloc(sizeof(npy_uint32) * len);
 
     if (ufunc->core_num_dims == NULL || ufunc->core_dim_ixs == NULL ||
-        ufunc->core_offsets == NULL || ufunc->core_dim_flags == NULL ||
-        ufunc->core_dim_sizes == NULL) {
+        ufunc->core_offsets == NULL ||
+        ufunc->extension->core_dim_sizes == NULL ||
+        ufunc->extension->core_dim_flags == NULL) {
         PyErr_NoMemory();
         goto fail;
     }
     for (i = 0; i < len; i++) {
-        ufunc->core_dim_flags[i] = 0;
+        ufunc->extension->core_dim_flags[i] = 0;
     }
 
     i = _next_non_white_space(signature, 0);
@@ -638,7 +639,7 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
              */
             for(ix = 0; ix < ufunc->core_num_dim_ix; ix++) {
                 if (frozen_size > 0 ?
-                    frozen_size == ufunc->core_dim_sizes[ix] :
+                    frozen_size == ufunc->extension->core_dim_sizes[ix] :
                     _is_same_name(signature + i, var_names[ix])) {
                     break;
                 }
@@ -649,20 +650,20 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
             if (ix == ufunc->core_num_dim_ix) {
                 ufunc->core_num_dim_ix++;
                 var_names[ix] = signature + i;
-                ufunc->core_dim_sizes[ix] = frozen_size;
+                ufunc->extension->core_dim_sizes[ix] = frozen_size;
                 if (frozen_size < 0) {
-                    ufunc->core_dim_flags[ix] |= UFUNC_CORE_DIM_SIZE_UNSET;
+                    ufunc->extension->core_dim_flags[ix] |= UFUNC_CORE_DIM_SIZE_UNSET;
                 }
                 if (can_ignore) {
-                    ufunc->core_dim_flags[ix] |= UFUNC_CORE_DIM_CAN_IGNORE;
+                    ufunc->extension->core_dim_flags[ix] |= UFUNC_CORE_DIM_CAN_IGNORE;
                 }
             } else {
-                if (can_ignore && !(ufunc->core_dim_flags[ix] &
+                if (can_ignore && !(ufunc->extension->core_dim_flags[ix] &
                                     UFUNC_CORE_DIM_CAN_IGNORE)) {
                     parse_error = "? cannot be used, name already seen without ?";
                     goto fail;
                 }
-                if (!can_ignore && (ufunc->core_dim_flags[ix] &
+                if (!can_ignore && (ufunc->extension->core_dim_flags[ix] &
                                     UFUNC_CORE_DIM_CAN_IGNORE)) {
                     parse_error = "? must be used, name already seen with ?";
                     goto fail;
@@ -709,9 +710,11 @@ _parse_signature(PyUFuncObject *ufunc, const char *signature)
     }
     ufunc->core_dim_ixs = PyArray_realloc(ufunc->core_dim_ixs,
             sizeof(int) * cur_core_dim);
-    ufunc->core_dim_sizes = PyArray_realloc(ufunc->core_dim_sizes,
+    ufunc->extension->core_dim_sizes = PyArray_realloc(
+            ufunc->extension->core_dim_sizes,
             sizeof(npy_intp) * ufunc->core_num_dim_ix);
-    ufunc->core_dim_flags = PyArray_realloc(ufunc->core_dim_flags,
+    ufunc->extension->core_dim_flags = PyArray_realloc(
+            ufunc->extension->core_dim_flags,
             sizeof(npy_uint32) * ufunc->core_num_dim_ix);
 
     /* check for trivial core-signature, e.g. "(),()->()" */
@@ -2489,10 +2492,10 @@ _initialize_variable_parts(PyUFuncObject *ufunc,
     for (i = 0; i < ufunc->nargs; i++) {
         op_core_num_dims[i] = ufunc->core_num_dims[i];
     }
-    if (ufunc->version == 1) {
+    if (ufunc->version == 1 && ufunc->extension != NULL) {
         for (i = 0; i < ufunc->core_num_dim_ix; i++) {
-            core_dim_sizes[i] = ufunc->core_dim_sizes[i];
-            core_dim_flags[i] = ufunc->core_dim_flags[i];
+            core_dim_sizes[i] = ufunc->extension->core_dim_sizes[i];
+            core_dim_flags[i] = ufunc->extension->core_dim_flags[i];
         }
     }
     else if (ufunc->version == 0) {
@@ -2503,7 +2506,7 @@ _initialize_variable_parts(PyUFuncObject *ufunc,
     }
     else {
         PyErr_Format(PyExc_TypeError,
-                     "'%s': unrecognized version number %d.",
+                     "'%s': unrecognized version number %d or corrupted data.",
                      ufunc_get_name_cstr(ufunc), ufunc->version);
         return -1;
     }
@@ -4870,6 +4873,7 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
     memset(ufunc, 0, sizeof(PyUFuncObject));
     *((int*)&ufunc->version) = UFUNC_VERSION;
     PyObject_Init((PyObject *)ufunc, &PyUFunc_Type);
+    ufunc->extension = &ufunc->s_extension;
 
     ufunc->nin = nin;
     ufunc->nout = nout;
@@ -5246,8 +5250,8 @@ ufunc_dealloc(PyUFuncObject *ufunc)
 {
     PyArray_free(ufunc->core_num_dims);
     PyArray_free(ufunc->core_dim_ixs);
-    PyArray_free(ufunc->core_dim_flags);
-    PyArray_free(ufunc->core_dim_sizes);
+    PyArray_free(ufunc->extension->core_dim_sizes);
+    PyArray_free(ufunc->extension->core_dim_flags);
     PyArray_free(ufunc->core_offsets);
     PyArray_free(ufunc->core_signature);
     PyArray_free(ufunc->ptr);