ENH: iter: Add support for custom dtypes by wrapping copyswap when necessary

3 files changed, 363 insertions, 165 deletions

diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c
index 21607cf48..c98453111 100644
--- a/numpy/core/src/multiarray/dtype_transfer.c
+++ b/numpy/core/src/multiarray/dtype_transfer.c
@@ -17,6 +17,31 @@
 
 #define NPY_LOWLEVEL_BUFFER_BLOCKSIZE  128
 
+/*
+ * Returns a transfer function which DECREFs any references in src_type.
+ *
+ * Returns NPY_SUCCEED or NPY_FAIL.
+ */
+static int
+get_decsrcref_transfer_function(int aligned,
+                            npy_intp src_stride,
+                            PyArray_Descr *src_dtype,
+                            PyArray_StridedTransferFn *outstransfer,
+                            void **outtransferdata);
+
+/*
+ * Returns a transfer function which zeros out the dest values.
+ *
+ * Returns NPY_SUCCEED or NPY_FAIL.
+ */
+static int
+get_setdstzero_transfer_function(int aligned,
+                            npy_intp dst_stride,
+                            PyArray_Descr *dst_dtype,
+                            PyArray_StridedTransferFn *outstransfer,
+                            void **outtransferdata);
+
+
 /*************************** COPY REFERENCES *******************************/
 
 /* Moves references from src to dst */
@@ -327,7 +352,7 @@ _strided_to_strided_contig_align_wrap_init_dest(char *dst, npy_intp dst_stride,
  * Returns NPY_SUCCEED or NPY_FAIL.
  */
 NPY_NO_EXPORT int
-PyArray_WrapAlignedContigTransferFunction(
+wrap_aligned_contig_transfer_function(
             npy_intp src_itemsize, npy_intp dst_itemsize,
             PyArray_StridedTransferFn tobuffer, void *todata,
             PyArray_StridedTransferFn frombuffer, void *fromdata,
@@ -378,6 +403,92 @@ PyArray_WrapAlignedContigTransferFunction(
     return NPY_SUCCEED;
 }
 
+/*************************** WRAP DTYPE COPY/SWAP *************************/
+/* Does a simple aligned cast */
+typedef struct {
+    void *freefunc, *copyfunc;
+    PyArray_CopySwapNFunc *copyswapn;
+    int swap;
+    PyArrayObject *arr;
+} _wrap_copy_swap_data;
+
+/* wrap copy swap data free function */
+void _wrap_copy_swap_data_free(_wrap_copy_swap_data *data)
+{
+    Py_DECREF(data->arr);
+    PyArray_free(data);
+}
+
+/* wrap copy swap data copy function */
+_wrap_copy_swap_data *_wrap_copy_swap_data_copy(_wrap_copy_swap_data *data)
+{
+    _wrap_copy_swap_data *newdata = 
+        (_wrap_copy_swap_data *)PyArray_malloc(sizeof(_wrap_copy_swap_data));
+    if (newdata == NULL) {
+        return NULL;
+    }
+
+    memcpy(newdata, data, sizeof(_wrap_copy_swap_data));
+    Py_INCREF(newdata->arr);
+
+    return newdata;
+}
+
+static void
+_strided_to_strided_wrap_copy_swap(char *dst, npy_intp dst_stride,
+                        char *src, npy_intp src_stride,
+                        npy_intp N, npy_intp NPY_UNUSED(src_itemsize),
+                        void *data)
+{
+    _wrap_copy_swap_data *d = (_wrap_copy_swap_data *)data;
+
+    d->copyswapn(dst, dst_stride, src, src_stride, N, d->swap, d->arr);
+}
+
+/* This only gets used for custom data types */
+static int
+wrap_copy_swap_function(int aligned,
+                npy_intp src_stride, npy_intp dst_stride,
+                PyArray_Descr *dtype,
+                int should_swap,
+                PyArray_StridedTransferFn *outstransfer,
+                void **outtransferdata)
+{
+    _wrap_copy_swap_data *data;
+    npy_intp shape = 1;
+
+    /* Allocate the data for the copy swap */
+    data = (_wrap_copy_swap_data *)PyArray_malloc(sizeof(_wrap_copy_swap_data));
+    if (data == NULL) {
+        PyErr_NoMemory();
+        *outstransfer = NULL;
+        *outtransferdata = NULL;
+        return NPY_FAIL;
+    }
+
+    data->freefunc = &_wrap_copy_swap_data_free;
+    data->copyfunc = &_wrap_copy_swap_data_copy;
+    data->copyswapn = dtype->f->copyswapn;
+    data->swap = should_swap;
+
+    /*
+     * TODO: This is a hack so the copyswap functions have an array.
+     *       The copyswap functions shouldn't need that.
+     */
+    Py_INCREF(dtype);
+    data->arr = (PyArrayObject *)PyArray_NewFromDescr(&PyArray_Type, dtype,
+                            1, &shape, NULL, NULL, 0, NULL);
+    if (data->arr == NULL) {
+        PyArray_free(data);
+        return NPY_FAIL;
+    }
+
+    *outstransfer = &_strided_to_strided_wrap_copy_swap;
+    *outtransferdata = data;
+
+    return NPY_SUCCEED;
+}
+
 /*************************** DTYPE CAST FUNCTIONS *************************/
 
 /* Does a simple aligned cast */
@@ -461,10 +572,8 @@ _aligned_contig_to_contig_cast(char *dst, npy_intp NPY_UNUSED(dst_stride),
                         void *data)
 {
     _strided_cast_data *d = (_strided_cast_data *)data;
-    PyArray_VectorUnaryFunc *castfunc = d->castfunc;
-    PyArrayObject *aip = d->aip, *aop = d->aop;
 
-    castfunc(src, dst, N, aip, aop);
+    d->castfunc(src, dst, N, d->aip, d->aop);
 }
 
 static int
@@ -476,6 +585,7 @@ get_cast_transfer_function(int aligned,
                             void **outtransferdata)
 {
     _strided_cast_data *data;
+    void *todata = NULL, *fromdata = NULL;
     PyArray_VectorUnaryFunc *castfunc;
     npy_intp shape = 1, src_itemsize = src_dtype->elsize,
             dst_itemsize = dst_dtype->elsize;
@@ -554,7 +664,18 @@ get_cast_transfer_function(int aligned,
         PyArray_StridedTransferFn tobuffer, frombuffer, casttransfer;
 
         /* Get the copy/swap operation from src */
-        if (src_itemsize == 1 || PyArray_ISNBO(src_dtype->byteorder)) {
+
+        /* If it's a custom data type, wrap its copy swap function */
+        if (src_dtype->type_num >= NPY_NTYPES) {
+            tobuffer = NULL;
+            wrap_copy_swap_function(aligned,
+                                src_stride, src_itemsize,
+                                src_dtype,
+                                !PyArray_ISNBO(src_dtype->byteorder),
+                                &tobuffer, &todata);
+        }
+        /* A straight copy */
+        else if (src_itemsize == 1 || PyArray_ISNBO(src_dtype->byteorder)) {
             tobuffer = PyArray_GetStridedCopyFn(aligned,
                                         src_stride, src_itemsize,
                                         src_itemsize);
@@ -573,7 +694,18 @@ get_cast_transfer_function(int aligned,
         }
 
         /* Get the copy/swap operation to dst */
-        if (dst_itemsize == 1 || PyArray_ISNBO(dst_dtype->byteorder)) {
+
+        /* If it's a custom data type, wrap its copy swap function */
+        if (dst_dtype->type_num >= NPY_NTYPES) {
+            frombuffer = NULL;
+            wrap_copy_swap_function(aligned,
+                                dst_itemsize, dst_stride,
+                                dst_dtype,
+                                !PyArray_ISNBO(dst_dtype->byteorder),
+                                &frombuffer, &fromdata);
+        }
+        /* A straight copy */
+        else if (dst_itemsize == 1 || PyArray_ISNBO(dst_dtype->byteorder)) {
             if (dst_dtype->type_num == NPY_OBJECT) {
                 frombuffer = &_strided_to_strided_move_references;
             }
@@ -598,6 +730,8 @@ get_cast_transfer_function(int aligned,
 
         if (frombuffer == NULL || tobuffer == NULL) {
             PyArray_FreeStridedTransferData(data);
+            PyArray_FreeStridedTransferData(todata);
+            PyArray_FreeStridedTransferData(fromdata);
             return NPY_FAIL;
         }
 
@@ -611,10 +745,10 @@ get_cast_transfer_function(int aligned,
         }
 
         /* Wrap it all up in a new transfer function + data */
-        if (PyArray_WrapAlignedContigTransferFunction(
+        if (wrap_aligned_contig_transfer_function(
                             src_itemsize, dst_itemsize,
-                            tobuffer, NULL,
-                            frombuffer, NULL,
+                            tobuffer, todata,
+                            frombuffer, fromdata,
                             casttransfer, data,
                             PyDataType_FLAGCHK(dst_dtype, NPY_NEEDS_INIT),
                             outstransfer, outtransferdata) != NPY_SUCCEED) {
@@ -805,7 +939,7 @@ get_one_to_n_transfer_function(int aligned,
 
     /* If the src object will need a DECREF, set src_dtype */
     if (move_references && PyDataType_REFCHK(src_dtype)) {
-        if (PyArray_GetDecSrcRefTransferFunction(aligned,
+        if (get_decsrcref_transfer_function(aligned,
                             src_stride,
                             src_dtype,
                             &stransfer_finish_src,
@@ -1000,6 +1134,7 @@ typedef struct {
     PyArray_StridedTransferFn stransfer;
     void *data;
     npy_intp src_N, dst_N, src_itemsize, dst_itemsize;
+    /* TODO: Switch to using decsrcref transfer functions instead */
     /* If this is non-NULL the source type has references needing a decref */
     PyArray_Descr *src_dtype;
     /* If this is non-NULL, the dest type has references needing a decref */
@@ -1090,15 +1225,17 @@ _strided_to_strided_subarray_broadcast(char *dst, npy_intp dst_stride,
                 }
                 else {
                     char *tmp = dst + i*dst_subitemsize;
-                    if (dst_dtype) {
+                    if (dst_dtype != NULL) {
                         PyArray_Item_XDECREF(tmp, dst_dtype);
                     }
                     memset(tmp, 0, dst_subitemsize);
                 }
             }
 
-            for (i = 0; i < src_subN; ++i) {
-                PyArray_Item_XDECREF(src + i*src_subitemsize, src_dtype);
+            if (src_dtype != NULL) {
+                for (i = 0; i < src_subN; ++i) {
+                    PyArray_Item_XDECREF(src + i*src_subitemsize, src_dtype);
+                }
             }
 
             src += src_stride;
@@ -1262,15 +1399,9 @@ get_subarray_transfer_function(int aligned,
     }
 
     /*
-     * Just a straight one-element copy.  If the source size isn't 1,
-     * we copy the element at index 0.  If the source data type is
-     * a reference and we're moving references, a DECREF for each
-     * source element would also be needed, so the general case will be
-     * used below
+     * Just a straight one-element copy.
      */
-    if (dst_size == 1 && (src_size == 1 ||
-                          !move_references ||
-                          !PyDataType_REFCHK(src_dtype))) {
+    if (dst_size == 1 && src_size == 1) {
         PyDimMem_FREE(src_shape.ptr);
         PyDimMem_FREE(dst_shape.ptr);
 
@@ -1494,7 +1625,7 @@ get_fields_transfer_function(int aligned,
          * another transfer function to do that.
          */
         if (move_references && PyDataType_REFCHK(src_dtype)) {
-            if (PyArray_GetDecSrcRefTransferFunction(0,
+            if (get_decsrcref_transfer_function(0,
                                     src_stride,
                                     src_dtype,
                                     &fields[field_count].stransfer,
@@ -1562,7 +1693,7 @@ get_fields_transfer_function(int aligned,
         }
         fields[0].src_offset = src_offset;
         fields[0].dst_offset = 0;
-        fields[0].src_itemsize = src_dtype->elsize;
+        fields[0].src_itemsize = src_fld_dtype->elsize;
 
         /*
          * If the references should be removed from src, add
@@ -1579,7 +1710,7 @@ get_fields_transfer_function(int aligned,
                     return NPY_FAIL;
                 }
                 if (PyDataType_REFCHK(src_fld_dtype)) {
-                    if (PyArray_GetDecSrcRefTransferFunction(0,
+                    if (get_decsrcref_transfer_function(0,
                                 src_stride,
                                 src_fld_dtype,
                                 &fields[field_count].stransfer,
@@ -1685,7 +1816,7 @@ get_fields_transfer_function(int aligned,
                 }
             }
             else {
-                if (PyArray_GetSetDstZeroTransferFunction(0,
+                if (get_setdstzero_transfer_function(0,
                                             dst_stride,
                                             dst_fld_dtype,
                                             &fields[i].stransfer,
@@ -1723,7 +1854,7 @@ get_fields_transfer_function(int aligned,
                         return NPY_FAIL;
                     }
                     if (PyDataType_REFCHK(src_fld_dtype)) {
-                        if (PyArray_GetDecSrcRefTransferFunction(0,
+                        if (get_decsrcref_transfer_function(0,
                                     src_stride,
                                     src_fld_dtype,
                                     &fields[field_count].stransfer,
@@ -1795,7 +1926,7 @@ get_decsrcref_fields_transfer_function(int aligned,
             return NPY_FAIL;
         }
         if (PyDataType_REFCHK(src_fld_dtype)) {
-            if (PyArray_GetDecSrcRefTransferFunction(0,
+            if (get_decsrcref_transfer_function(0,
                                     src_stride,
                                     src_fld_dtype,
                                     &fields[field_count].stransfer,
@@ -1859,7 +1990,7 @@ get_setdestzero_fields_transfer_function(int aligned,
             PyArray_free(data);
             return NPY_FAIL;
         }
-        if (PyArray_GetSetDstZeroTransferFunction(0,
+        if (get_setdstzero_transfer_function(0,
                                 dst_stride,
                                 dst_fld_dtype,
                                 &fields[i].stransfer,
@@ -1883,112 +2014,6 @@ get_setdestzero_fields_transfer_function(int aligned,
     return NPY_SUCCEED;
 }
 
-NPY_NO_EXPORT int
-PyArray_GetDTypeTransferFunction(int aligned,
-                            npy_intp src_stride, npy_intp dst_stride,
-                            PyArray_Descr *src_dtype, PyArray_Descr *dst_dtype,
-                            int move_references,
-                            PyArray_StridedTransferFn *outstransfer,
-                            void **outtransferdata)
-{
-    npy_intp src_itemsize = src_dtype->elsize,
-             dst_itemsize = dst_dtype->elsize;
-    int src_type_num = src_dtype->type_num,
-        dst_type_num = dst_dtype->type_num;
-
-    /* First look at the possibilities of just a copy or swap */
-    if (src_itemsize == dst_itemsize && src_dtype->kind == dst_dtype->kind &&
-                src_type_num < NPY_NTYPES && dst_type_num < NPY_NTYPES &&
-                !PyDataType_HASFIELDS(src_dtype) &&
-                !PyDataType_HASFIELDS(dst_dtype) &&
-                src_dtype->subarray == NULL && dst_dtype->subarray == NULL) {
-        /* The special types, which have no byte-order */
-        switch (src_type_num) {
-            case NPY_VOID:
-            case NPY_STRING:
-            case NPY_UNICODE:
-                *outstransfer = PyArray_GetStridedCopyFn(0,
-                                    src_stride, dst_stride,
-                                    src_itemsize);
-                *outtransferdata = NULL;
-                return NPY_SUCCEED;
-            case NPY_OBJECT:
-                if (move_references) {
-                    *outstransfer = &_strided_to_strided_move_references;
-                    *outtransferdata = NULL;
-                }
-                else {
-                    *outstransfer = &_strided_to_strided_copy_references;
-                    *outtransferdata = NULL;
-                }
-                return NPY_SUCCEED;
-        }
-
-        /* This is a straight copy */
-        if (src_itemsize == 1 || PyArray_ISNBO(src_dtype->byteorder) ==
-                                 PyArray_ISNBO(dst_dtype->byteorder)) {
-            *outstransfer = PyArray_GetStridedCopyFn(aligned,
-                                        src_stride, dst_stride,
-                                        src_itemsize);
-            *outtransferdata = NULL;
-            return (*outstransfer == NULL) ? NPY_FAIL : NPY_SUCCEED;
-        }
-        /* This is a straight copy + byte swap */
-        else if (!PyTypeNum_ISCOMPLEX(src_type_num)) {
-            *outstransfer = PyArray_GetStridedCopySwapFn(aligned,
-                                        src_stride, dst_stride,
-                                        src_itemsize);
-            *outtransferdata = NULL;
-            return (*outstransfer == NULL) ? NPY_FAIL : NPY_SUCCEED;
-        }
-        /* This is a straight copy + element pair byte swap */
-        else {
-            *outstransfer = PyArray_GetStridedCopySwapPairFn(aligned,
-                                        src_stride, dst_stride,
-                                        src_itemsize);
-            *outtransferdata = NULL;
-            return (*outstransfer == NULL) ? NPY_FAIL : NPY_SUCCEED;
-        }
-    }
-
-    /* Handle subarrays */
-    if (src_dtype->subarray != NULL || dst_dtype->subarray != NULL) {
-        return get_subarray_transfer_function(aligned,
-                        src_stride, dst_stride,
-                        src_dtype, dst_dtype,
-                        move_references,
-                        outstransfer, outtransferdata);
-    }
-
-    /* Handle fields */
-    if (PyDataType_HASFIELDS(src_dtype) ||
-                PyDataType_HASFIELDS(dst_dtype)) {
-        return get_fields_transfer_function(aligned,
-                        src_stride, dst_stride,
-                        src_dtype, dst_dtype,
-                        move_references,
-                        outstransfer, outtransferdata);
-    }
-
-    /* Check for different-sized strings, unicodes, or voids */
-    if (src_type_num == dst_type_num) switch (src_type_num) {
-        case NPY_STRING:
-        case NPY_UNICODE:
-        case NPY_VOID:
-            return PyArray_GetStridedZeroPadCopyFn(0,
-                                    src_stride, dst_stride,
-                                    src_dtype->elsize, dst_dtype->elsize,
-                                    outstransfer, outtransferdata);
-    }
-
-    /* Otherwise a cast is necessary */
-    return get_cast_transfer_function(aligned,
-                    src_stride, dst_stride,
-                    src_dtype, dst_dtype,
-                    move_references,
-                    outstransfer, outtransferdata);
-}
-
 /*************************** DEST SETZERO *******************************/
 
 /* Sets dest to zero */
@@ -2068,7 +2093,7 @@ _null_to_strided_reference_setzero(char *dst,
 }
 
 NPY_NO_EXPORT int
-PyArray_GetSetDstZeroTransferFunction(int aligned,
+get_setdstzero_transfer_function(int aligned,
                             npy_intp dst_stride,
                             PyArray_Descr *dst_dtype,
                             PyArray_StridedTransferFn *outstransfer,
@@ -2123,7 +2148,7 @@ PyArray_GetSetDstZeroTransferFunction(int aligned,
         PyDimMem_FREE(dst_shape.ptr);
 
         /* Get a function for contiguous dst of the subarray type */
-        if (PyArray_GetSetDstZeroTransferFunction(aligned,
+        if (get_setdstzero_transfer_function(aligned,
                                 dst_dtype->subarray->base->elsize,
                                 dst_dtype->subarray->base,
                                 &stransfer, &data) != NPY_SUCCEED) {
@@ -2183,7 +2208,7 @@ _strided_to_null_dec_src_ref_reference(char *NPY_UNUSED(dst),
 
 
 NPY_NO_EXPORT int
-PyArray_GetDecSrcRefTransferFunction(int aligned,
+get_decsrcref_transfer_function(int aligned,
                             npy_intp src_stride,
                             PyArray_Descr *src_dtype,
                             PyArray_StridedTransferFn *outstransfer,
@@ -2220,7 +2245,7 @@ PyArray_GetDecSrcRefTransferFunction(int aligned,
         PyDimMem_FREE(src_shape.ptr);
         
         /* Get a function for contiguous src of the subarray type */
-        if (PyArray_GetDecSrcRefTransferFunction(aligned,
+        if (get_decsrcref_transfer_function(aligned,
                                 src_dtype->subarray->base->elsize,
                                 src_dtype->subarray->base,
                                 &stransfer, &data) != NPY_SUCCEED) {
@@ -2247,3 +2272,161 @@ PyArray_GetDecSrcRefTransferFunction(int aligned,
     }
 }
 
+/********************* MAIN DTYPE TRANSFER FUNCTION ***********************/
+
+NPY_NO_EXPORT int
+PyArray_GetDTypeTransferFunction(int aligned,
+                            npy_intp src_stride, npy_intp dst_stride,
+                            PyArray_Descr *src_dtype, PyArray_Descr *dst_dtype,
+                            int move_references,
+                            PyArray_StridedTransferFn *outstransfer,
+                            void **outtransferdata)
+{
+    npy_intp src_itemsize, dst_itemsize;
+    int src_type_num, dst_type_num;
+
+    /*
+     * If one of the dtypes is NULL, we give back either a src decref
+     * function or a dst setzero function
+     */
+    if (dst_dtype == NULL) {
+        if (move_references) {
+            return get_decsrcref_transfer_function(aligned,
+                                src_dtype->elsize,
+                                src_dtype,
+                                outstransfer, outtransferdata);
+        }
+        else {
+            *outstransfer = &_dec_src_ref_nop;
+            *outtransferdata = NULL;
+        }
+    }
+    else if (src_dtype == NULL) {
+        return get_setdstzero_transfer_function(aligned,
+                                dst_dtype->elsize,
+                                dst_dtype,
+                                outstransfer, outtransferdata);
+    }
+
+    src_itemsize = src_dtype->elsize;
+    dst_itemsize = dst_dtype->elsize;
+    src_type_num = src_dtype->type_num;
+    dst_type_num = dst_dtype->type_num;
+
+    /* First look at the possibilities of just a copy or swap */
+    if (src_itemsize == dst_itemsize && src_dtype->kind == dst_dtype->kind &&
+                !PyDataType_HASFIELDS(src_dtype) &&
+                !PyDataType_HASFIELDS(dst_dtype) &&
+                src_dtype->subarray == NULL && dst_dtype->subarray == NULL) {
+        /* A custom data type requires that we use its copy/swap */
+        if (src_type_num >= NPY_NTYPES || dst_type_num >= NPY_NTYPES) {
+            /*
+             * If the sizes and kinds are identical, but they're different
+             * custom types, then get a cast function
+             */
+            if (src_type_num != dst_type_num) {
+                return get_cast_transfer_function(aligned,
+                                src_stride, dst_stride,
+                                src_dtype, dst_dtype,
+                                move_references,
+                                outstransfer, outtransferdata);
+            }
+            else {
+                return wrap_copy_swap_function(aligned,
+                                src_stride, dst_stride,
+                                src_dtype,
+                                PyArray_ISNBO(src_dtype->byteorder) !=
+                                        PyArray_ISNBO(dst_dtype->byteorder),
+                                outstransfer, outtransferdata);
+            }
+
+
+        }
+
+        /* The special types, which have no byte-order */
+        switch (src_type_num) {
+            case NPY_VOID:
+            case NPY_STRING:
+            case NPY_UNICODE:
+                *outstransfer = PyArray_GetStridedCopyFn(0,
+                                    src_stride, dst_stride,
+                                    src_itemsize);
+                *outtransferdata = NULL;
+                return NPY_SUCCEED;
+            case NPY_OBJECT:
+                if (move_references) {
+                    *outstransfer = &_strided_to_strided_move_references;
+                    *outtransferdata = NULL;
+                }
+                else {
+                    *outstransfer = &_strided_to_strided_copy_references;
+                    *outtransferdata = NULL;
+                }
+                return NPY_SUCCEED;
+        }
+
+        /* This is a straight copy */
+        if (src_itemsize == 1 || PyArray_ISNBO(src_dtype->byteorder) ==
+                                 PyArray_ISNBO(dst_dtype->byteorder)) {
+            *outstransfer = PyArray_GetStridedCopyFn(aligned,
+                                        src_stride, dst_stride,
+                                        src_itemsize);
+            *outtransferdata = NULL;
+            return (*outstransfer == NULL) ? NPY_FAIL : NPY_SUCCEED;
+        }
+        /* This is a straight copy + byte swap */
+        else if (!PyTypeNum_ISCOMPLEX(src_type_num)) {
+            *outstransfer = PyArray_GetStridedCopySwapFn(aligned,
+                                        src_stride, dst_stride,
+                                        src_itemsize);
+            *outtransferdata = NULL;
+            return (*outstransfer == NULL) ? NPY_FAIL : NPY_SUCCEED;
+        }
+        /* This is a straight copy + element pair byte swap */
+        else {
+            *outstransfer = PyArray_GetStridedCopySwapPairFn(aligned,
+                                        src_stride, dst_stride,
+                                        src_itemsize);
+            *outtransferdata = NULL;
+            return (*outstransfer == NULL) ? NPY_FAIL : NPY_SUCCEED;
+        }
+    }
+
+    /* Handle subarrays */
+    if (src_dtype->subarray != NULL || dst_dtype->subarray != NULL) {
+        return get_subarray_transfer_function(aligned,
+                        src_stride, dst_stride,
+                        src_dtype, dst_dtype,
+                        move_references,
+                        outstransfer, outtransferdata);
+    }
+
+    /* Handle fields */
+    if (PyDataType_HASFIELDS(src_dtype) ||
+                PyDataType_HASFIELDS(dst_dtype)) {
+        return get_fields_transfer_function(aligned,
+                        src_stride, dst_stride,
+                        src_dtype, dst_dtype,
+                        move_references,
+                        outstransfer, outtransferdata);
+    }
+
+    /* Check for different-sized strings, unicodes, or voids */
+    if (src_type_num == dst_type_num) switch (src_type_num) {
+        case NPY_STRING:
+        case NPY_UNICODE:
+        case NPY_VOID:
+            return PyArray_GetStridedZeroPadCopyFn(0,
+                                    src_stride, dst_stride,
+                                    src_dtype->elsize, dst_dtype->elsize,
+                                    outstransfer, outtransferdata);
+    }
+
+    /* Otherwise a cast is necessary */
+    return get_cast_transfer_function(aligned,
+                    src_stride, dst_stride,
+                    src_dtype, dst_dtype,
+                    move_references,
+                    outstransfer, outtransferdata);
+}
+
diff --git a/numpy/core/src/multiarray/lowlevel_strided_loops.h b/numpy/core/src/multiarray/lowlevel_strided_loops.h
index ef962038a..89fcd72d6 100644
--- a/numpy/core/src/multiarray/lowlevel_strided_loops.h
+++ b/numpy/core/src/multiarray/lowlevel_strided_loops.h
@@ -2,6 +2,12 @@
 #define __LOWLEVEL_STRIDED_LOOPS_H
 
 /*
+ * NOTE: This API should remain private for the time being, to allow
+ *       for further refinement.  I think the 'aligned' mechanism
+ *       needs changing, for example.
+ */
+
+/*
  * This function pointer is for functions that transfer an arbitrarily strided
  * input to a an arbitrarily strided output.  It may be a fully general
  * function, or a specialized function when the strides or item size
@@ -62,6 +68,9 @@ PyArray_GetStridedCopyFn(npy_intp aligned, npy_intp src_stride,
  * and swapping strided memory.  This assumes each element is a single
  * value to be swapped.
  *
+ * For information on the 'aligned', 'src_stride' and 'dst_stride' parameters
+ * see above.
+ *
  * Parameters are as for PyArray_GetStridedCopyFn.
  */
 NPY_NO_EXPORT PyArray_StridedTransferFn
@@ -73,6 +82,9 @@ PyArray_GetStridedCopySwapFn(npy_intp aligned, npy_intp src_stride,
  * and swapping strided memory.  This assumes each element is a pair
  * of values, each of which needs to be swapped.
  *
+ * For information on the 'aligned', 'src_stride' and 'dst_stride' parameters
+ * see above.
+ *
  * Parameters are as for PyArray_GetStridedCopyFn.
  */
 NPY_NO_EXPORT PyArray_StridedTransferFn
@@ -84,6 +96,9 @@ PyArray_GetStridedCopySwapPairFn(npy_intp aligned, npy_intp src_stride,
  * the data from source to dest, truncating it if the data doesn't
  * fit, and padding with zero bytes if there's too much space.
  *
+ * For information on the 'aligned', 'src_stride' and 'dst_stride' parameters
+ * see above.
+ *
  * Returns NPY_SUCCEED or NPY_FAIL
  */
 NPY_NO_EXPORT int
@@ -94,36 +109,15 @@ PyArray_GetStridedZeroPadCopyFn(int aligned,
                             void **outtransferdata);
 
 /*
- * Returns a transfer function which DECREFs any references in src_type.
- *
- * Returns NPY_SUCCEED or NPY_FAIL.
- */
-NPY_NO_EXPORT int
-PyArray_GetDecSrcRefTransferFunction(int aligned,
-                            npy_intp src_stride,
-                            PyArray_Descr *src_dtype,
-                            PyArray_StridedTransferFn *outstransfer,
-                            void **outtransferdata);
-
-/*
- * Returns a transfer function which zeros out the dest values.
- *
- * Returns NPY_SUCCEED or NPY_FAIL.
- */
-NPY_NO_EXPORT int
-PyArray_GetSetDstZeroTransferFunction(int aligned,
-                            npy_intp dst_stride,
-                            PyArray_Descr *dst_dtype,
-                            PyArray_StridedTransferFn *outstransfer,
-                            void **outtransferdata);
-
-/*
  * If it's possible, gives back a transfer function which casts and/or
  * byte swaps data with the dtype 'src_dtype' into data with the dtype
  * 'dst_dtype'.  If the outtransferdata is populated with a non-NULL value,
  * it must be deallocated with the ``PyArray_FreeStridedTransferData``
  * function when the transfer function is no longer required.
  *
+ * For information on the 'aligned', 'src_stride' and 'dst_stride' parameters
+ * see above.
+ *
  * If move_references is 1, and 'src_dtype' has references,
  * the source references will get a DECREF after the reference value is
  * cast to the dest type.
@@ -133,6 +127,13 @@ PyArray_GetSetDstZeroTransferFunction(int aligned,
  *          first destination reference will get the value and all the rest
  *          will get NULL.
  *
+ * If you pass NULL to src_dtype, you get a transfer function which does
+ * not touch the src values, and sets the dst values to zeros.
+ *
+ * If you pass NULL to dst_dtype, you either get a no-op, or if you
+ * also set move_references to 1, you get a transfer function which
+ * decrements the references in src.
+ *
  * Returns NPY_SUCCEED or NPY_FAIL.
  */
 NPY_NO_EXPORT int
diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 99c6decd9..f11fd7a60 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -4252,10 +4252,16 @@ npyiter_allocate_transfer_functions(NpyIter *iter)
             }
             /* If no write back but there are references make a decref fn */
             else if (PyDataType_REFCHK(op_dtype[iiter])) {
-                if (PyArray_GetDecSrcRefTransferFunction(
+                /*
+                 * By passing NULL to dst_type and setting move_references
+                 * to 1, we get back a function that just decrements the
+                 * src references.
+                 */
+                if (PyArray_GetDTypeTransferFunction(
                                         (flags&NPY_OP_ITFLAG_ALIGNED) != 0,
-                                        op_dtype[iiter]->elsize,
-                                        op_dtype[iiter],
+                                        op_dtype[iiter]->elsize, 0,
+                                        op_dtype[iiter], NULL,
+                                        1,
                                         &stransfer,
                                         &transferdata) != NPY_SUCCEED) {
                     goto fail;
@@ -4493,9 +4499,17 @@ npyiter_copy_from_buffers(NpyIter *iter)
             /* Decrement refs only if the pointer was pointing to the buffer */
             npy_intp delta = (ptrs[iiter] - buffer);
             if (0 <= delta && delta <= transfersize*dtypes[iiter]->elsize) {
+                /* Decrement refs */
                 stransfer(NULL, 0, buffer, dtypes[iiter]->elsize,
                             transfersize, dtypes[iiter]->elsize,
                             transferdata);
+                /*
+                 * Zero out the memory for safety.  For instance,
+                 * if during iteration some Python code copied an
+                 * array pointing into the buffer, it will get None
+                 * values for its references after this.
+                 */
+                memset(buffer, 0, dtypes[iiter]->elsize*transfersize);
             }
         }
     }