Add support for structured array ufuncs

- Add New api method PyUFunc_RegisterLoopForStructType for creating new ufunc for structured array
- Add arg_dtypes array to PyUFunc_Loop1d object. This keeps track of each struct dtype that the ufunc can handle.
- Add arg_dtypes parameter to ufunc_loop_matches function so that struct dtypes can be compared for NPY_VOID type num

4 files changed, 166 insertions, 20 deletions

diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index 07a87f98d..ea4ac9f07 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -384,6 +384,7 @@ ufunc_funcs_api = {
     # End 1.6 API
     'PyUFunc_DefaultTypeResolver':              39,
     'PyUFunc_ValidateCasting':                  40,
+    'PyUFunc_RegisterLoopForStructType':        41,
 }
 
 # List of all the dicts which define the C API
diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index 686d12c38..75611426c 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -319,6 +319,8 @@ typedef struct _loop1d_info {
         void *data;
         int *arg_types;
         struct _loop1d_info *next;
+        int nargs;
+        PyArray_Descr **arg_dtypes;
 } PyUFunc_Loop1d;
 
 
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index a3a164731..c67874246 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -725,7 +725,7 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
     PyObject *str_key_obj = NULL;
     char *ufunc_name;
 
-    int any_flexible = 0, any_object = 0;
+    int any_flexible = 0, any_object = 0, any_flexible_userloops = 0;
 
     ufunc_name = ufunc->name ? ufunc->name : "<unnamed ufunc>";
 
@@ -764,23 +764,56 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
         if (out_op[i] == NULL) {
             return -1;
         }
+
+        int type_num = PyArray_DESCR(out_op[i])->type_num;
         if (!any_flexible &&
-                PyTypeNum_ISFLEXIBLE(PyArray_DESCR(out_op[i])->type_num)) {
+                PyTypeNum_ISFLEXIBLE(type_num)) {
             any_flexible = 1;
         }
         if (!any_object &&
-                PyTypeNum_ISOBJECT(PyArray_DESCR(out_op[i])->type_num)) {
+                PyTypeNum_ISOBJECT(type_num)) {
             any_object = 1;
         }
+
+        /*
+         * If any operand is a flexible dtype, check to see if any
+         * struct dtype ufuncs are registered. A ufunc has been registered
+         * for a struct dtype if ufunc's arg_dtypes array is not NULL.
+         */
+        if (PyTypeNum_ISFLEXIBLE(type_num) &&
+            !any_flexible_userloops &&
+            ufunc->userloops != NULL) {
+            PyUFunc_Loop1d *funcdata;
+            PyObject *key, *obj;
+            key = PyInt_FromLong(type_num);
+            if (key == NULL) {
+                continue;
+            }
+            obj = PyDict_GetItem(ufunc->userloops, key);
+            Py_DECREF(key);
+            if (obj == NULL) {
+                continue;
+            }
+            funcdata = (PyUFunc_Loop1d *)NpyCapsule_AsVoidPtr(obj);
+            while (funcdata != NULL) {
+                if (funcdata->arg_dtypes != NULL) {
+                    any_flexible_userloops = 1;
+                    break;
+                }
+
+                funcdata = funcdata->next;
+            }
+        }
     }
 
     /*
      * Indicate not implemented if there are flexible objects (structured
-     * type or string) but no object types.
+     * type or string) but no object types and no registered struct
+     * dtype ufuncs.
      *
      * Not sure - adding this increased to 246 errors, 150 failures.
      */
-    if (any_flexible && !any_object) {
+    if (any_flexible && !any_flexible_userloops && !any_object) {
         return -2;
 
     }
@@ -4356,9 +4389,19 @@ cmp_arg_types(int *arg1, int *arg2, int n)
 static NPY_INLINE void
 _free_loop1d_list(PyUFunc_Loop1d *data)
 {
+    int i;
+
     while (data != NULL) {
         PyUFunc_Loop1d *next = data->next;
         PyArray_free(data->arg_types);
+
+        if (data->arg_dtypes != NULL) {
+            for (i = 0; i < data->nargs; i++) {
+                Py_DECREF(data->arg_dtypes[i]);
+            }
+            PyArray_free(data->arg_dtypes);
+        }
+
         PyArray_free(data);
         data = next;
     }
@@ -4383,6 +4426,90 @@ _loop1d_list_free(void *ptr)
 
 /*UFUNC_API*/
 NPY_NO_EXPORT int
+PyUFunc_RegisterLoopForStructType(PyUFuncObject *ufunc,
+                            PyArray_Descr *user_dtype,
+                            PyUFuncGenericFunction function,
+                            PyArray_Descr **arg_dtypes,
+                            void *data)
+{
+    int i;
+    int result = 0;
+    int *arg_typenums;
+    PyObject *key, *cobj;
+
+    if (user_dtype == NULL) {
+        PyErr_SetString(PyExc_TypeError, "unknown user defined struct dtype");
+        return -1;
+    }
+
+    key = PyInt_FromLong((long) user_dtype->type_num);
+    if (key == NULL) {
+        return -1;
+    }
+
+    arg_typenums = PyArray_malloc(ufunc->nargs * sizeof(int));
+    if (arg_dtypes != NULL) {
+        for (i = 0; i < ufunc->nargs; i++) {
+            arg_typenums[i] = arg_dtypes[i]->type_num;
+        }
+    }
+    else {
+        for (i = 0; i < ufunc->nargs; i++) {
+            arg_typenums[i] = user_dtype->type_num;
+        }
+    }
+    
+    result = PyUFunc_RegisterLoopForType(ufunc, user_dtype->type_num, function, arg_typenums, data);
+
+    if (result == 0) {
+        cobj = PyDict_GetItem(ufunc->userloops, key);
+        if (cobj == NULL) {
+            PyErr_SetString(PyExc_KeyError, "userloop for user dtype not found");
+            result = -1;
+        }
+        else {
+            PyUFunc_Loop1d *current, *prev = NULL;
+            int cmp = 1;
+            current = (PyUFunc_Loop1d *)NpyCapsule_AsVoidPtr(cobj);
+            while (current != NULL) {
+                cmp = cmp_arg_types(current->arg_types, arg_typenums, ufunc->nargs);
+                if (cmp >= 0 && current->arg_dtypes == NULL) {
+                    break;
+                }
+                prev = current;
+                current = current->next;
+            }
+            if (cmp == 0 && current->arg_dtypes == NULL) {
+                current->arg_dtypes = PyArray_malloc(ufunc->nargs * sizeof(PyArray_Descr*));
+                if (arg_dtypes != NULL) {
+                    for (i = 0; i < ufunc->nargs; i++) {
+                        current->arg_dtypes[i] = arg_dtypes[i];
+                        Py_INCREF(current->arg_dtypes[i]);
+                    }
+                }
+                else {
+                    for (i = 0; i < ufunc->nargs; i++) {
+                        current->arg_dtypes[i] = user_dtype;
+                        Py_INCREF(current->arg_dtypes[i]);
+                    }
+                }
+                current->nargs = ufunc->nargs;
+            }
+            else {
+                result = -1;
+            }
+        }
+    }
+
+    free(arg_typenums);
+
+    Py_DECREF(key);
+
+    return result;
+}
+
+/*UFUNC_API*/
+NPY_NO_EXPORT int
 PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
                             int usertype,
                             PyUFuncGenericFunction function,
@@ -4396,7 +4523,7 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
     int *newtypes=NULL;
 
     descr=PyArray_DescrFromType(usertype);
-    if ((usertype < NPY_USERDEF) || (descr==NULL)) {
+    if ((usertype < NPY_USERDEF && usertype != NPY_VOID) || (descr==NULL)) {
         PyErr_SetString(PyExc_TypeError, "unknown user-defined type");
         return -1;
     }
@@ -4432,6 +4559,8 @@ PyUFunc_RegisterLoopForType(PyUFuncObject *ufunc,
     funcdata->arg_types = newtypes;
     funcdata->data = data;
     funcdata->next = NULL;
+    funcdata->arg_dtypes = NULL;
+    funcdata->nargs = 0;
 
     /* Get entry for this user-defined type*/
     cobj = PyDict_GetItem(ufunc->userloops, key);
diff --git a/numpy/core/src/umath/ufunc_type_resolution.c b/numpy/core/src/umath/ufunc_type_resolution.c
index 7e3fb5633..a75991c2a 100644
--- a/numpy/core/src/umath/ufunc_type_resolution.c
+++ b/numpy/core/src/umath/ufunc_type_resolution.c
@@ -1188,7 +1188,7 @@ find_userloop(PyUFuncObject *ufunc,
         }
 
         type_num = dtypes[i]->type_num;
-        if (type_num != last_userdef && PyTypeNum_ISUSERDEF(type_num)) {
+        if (type_num != last_userdef && (PyTypeNum_ISUSERDEF(type_num) || type_num == NPY_VOID)) {
             PyObject *key, *obj;
 
             last_userdef = type_num;
@@ -1436,7 +1436,7 @@ ufunc_loop_matches(PyUFuncObject *self,
                     NPY_CASTING output_casting,
                     int any_object,
                     int use_min_scalar,
-                    int *types,
+                    int *types, PyArray_Descr **dtypes,
                     int *out_no_castable_output,
                     char *out_err_src_typecode,
                     char *out_err_dst_typecode)
@@ -1463,7 +1463,18 @@ ufunc_loop_matches(PyUFuncObject *self,
             return 0;
         }
 
-        tmp = PyArray_DescrFromType(types[i]);
+        /*
+         * If type num is NPY_VOID and struct dtypes have been passed in,
+         * use struct dtype object. Otherwise create new dtype object
+         * from type num.
+         */
+        if (types[i] == NPY_VOID && dtypes != NULL) {
+            tmp = dtypes[i];
+            Py_INCREF(tmp);
+        }
+        else {
+            tmp = PyArray_DescrFromType(types[i]);
+        }
         if (tmp == NULL) {
             return -1;
         }
@@ -1525,7 +1536,7 @@ ufunc_loop_matches(PyUFuncObject *self,
 static int
 set_ufunc_loop_data_types(PyUFuncObject *self, PyArrayObject **op,
                     PyArray_Descr **out_dtypes,
-                    int *type_nums)
+                    int *type_nums, PyArray_Descr **dtypes)
 {
     int i, nin = self->nin, nop = nin + self->nout;
 
@@ -1536,11 +1547,14 @@ set_ufunc_loop_data_types(PyUFuncObject *self, PyArrayObject **op,
      * instead of creating a new one, similarly to preserve metadata.
      **/
     for (i = 0; i < nop; ++i) {
+        if (dtypes != NULL) {
+            out_dtypes[i] = dtypes[i];
+            Py_XINCREF(out_dtypes[i]);
         /*
          * Copy the dtype from 'op' if the type_num matches,
          * to preserve metadata.
          */
-        if (op[i] != NULL && PyArray_DESCR(op[i])->type_num == type_nums[i]) {
+        } else if (op[i] != NULL && PyArray_DESCR(op[i])->type_num == type_nums[i]) {
             out_dtypes[i] = ensure_dtype_nbo(PyArray_DESCR(op[i]));
             Py_XINCREF(out_dtypes[i]);
         }
@@ -1603,7 +1617,7 @@ linear_search_userloop_type_resolver(PyUFuncObject *self,
         }
 
         type_num = PyArray_DESCR(op[i])->type_num;
-        if (type_num != last_userdef && PyTypeNum_ISUSERDEF(type_num)) {
+        if (type_num != last_userdef && (PyTypeNum_ISUSERDEF(type_num) || type_num == NPY_VOID)) {
             PyObject *key, *obj;
 
             last_userdef = type_num;
@@ -1623,7 +1637,7 @@ linear_search_userloop_type_resolver(PyUFuncObject *self,
                 switch (ufunc_loop_matches(self, op,
                             input_casting, output_casting,
                             any_object, use_min_scalar,
-                            types,
+                            types, funcdata->arg_dtypes,
                             out_no_castable_output, out_err_src_typecode,
                             out_err_dst_typecode)) {
                     /* Error */
@@ -1631,7 +1645,7 @@ linear_search_userloop_type_resolver(PyUFuncObject *self,
                         return -1;
                     /* Found a match */
                     case 1:
-                        set_ufunc_loop_data_types(self, op, out_dtype, types);
+                        set_ufunc_loop_data_types(self, op, out_dtype, types, funcdata->arg_dtypes);
                         return 1;
                 }
 
@@ -1707,12 +1721,12 @@ type_tuple_userloop_type_resolver(PyUFuncObject *self,
                 switch (ufunc_loop_matches(self, op,
                             casting, casting,
                             any_object, use_min_scalar,
-                            types,
+                            types, NULL,
                             &no_castable_output, &err_src_typecode,
                             &err_dst_typecode)) {
                     /* It works */
                     case 1:
-                        set_ufunc_loop_data_types(self, op, out_dtype, types);
+                        set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
                         return 1;
                     /* Didn't match */
                     case 0:
@@ -1875,7 +1889,7 @@ linear_search_type_resolver(PyUFuncObject *self,
         switch (ufunc_loop_matches(self, op,
                     input_casting, output_casting,
                     any_object, use_min_scalar,
-                    types,
+                    types, NULL,
                     &no_castable_output, &err_src_typecode,
                     &err_dst_typecode)) {
             /* Error */
@@ -1883,7 +1897,7 @@ linear_search_type_resolver(PyUFuncObject *self,
                 return -1;
             /* Found a match */
             case 1:
-                set_ufunc_loop_data_types(self, op, out_dtype, types);
+                set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
                 return 0;
         }
     }
@@ -2081,7 +2095,7 @@ type_tuple_type_resolver(PyUFuncObject *self,
         switch (ufunc_loop_matches(self, op,
                     casting, casting,
                     any_object, use_min_scalar,
-                    types,
+                    types, NULL,
                     &no_castable_output, &err_src_typecode,
                     &err_dst_typecode)) {
             /* Error */
@@ -2089,7 +2103,7 @@ type_tuple_type_resolver(PyUFuncObject *self,
                 return -1;
             /* It worked */
             case 1:
-                set_ufunc_loop_data_types(self, op, out_dtype, types);
+                set_ufunc_loop_data_types(self, op, out_dtype, types, NULL);
                 return 0;
             /* Didn't work */
             case 0: