MAINT: Break out duplicated logic for strided casts.

1 files changed, 13 insertions, 52 deletions

diff --git a/numpy/core/src/multiarray/dtype_transfer.c b/numpy/core/src/multiarray/dtype_transfer.c
index b26d5ac89..6c5fe2aed 100644
--- a/numpy/core/src/multiarray/dtype_transfer.c
+++ b/numpy/core/src/multiarray/dtype_transfer.c
@@ -312,6 +312,7 @@ typedef struct {
     NpyAuxData *wrappeddata, *todata, *fromdata;
     npy_intp src_itemsize, dst_itemsize;
     char *bufferin, *bufferout;
+    npy_bool init_dest;
 } _align_wrap_data;
 
 /* transfer data free function */
@@ -372,6 +373,8 @@ static NpyAuxData *_align_wrap_data_clone(NpyAuxData *data)
         }
     }
 
+    newdata->init_dest = d->init_dest;
+
     return (NpyAuxData *)newdata;
 }
 
@@ -391,12 +394,17 @@ _strided_to_strided_contig_align_wrap(char *dst, npy_intp dst_stride,
             *todata = d->todata,
             *fromdata = d->fromdata;
     char *bufferin = d->bufferin, *bufferout = d->bufferout;
+    npy_bool init_dest = d->init_dest;
 
     for(;;) {
         if (N > NPY_LOWLEVEL_BUFFER_BLOCKSIZE) {
             tobuffer(bufferin, inner_src_itemsize, src, src_stride,
                                     NPY_LOWLEVEL_BUFFER_BLOCKSIZE,
                                     src_itemsize, todata);
+            if (init_dest) {
+                memset(bufferout, 0,
+                        dst_itemsize*NPY_LOWLEVEL_BUFFER_BLOCKSIZE);
+            }
             wrapped(bufferout, dst_itemsize, bufferin, inner_src_itemsize,
                                     NPY_LOWLEVEL_BUFFER_BLOCKSIZE,
                                     inner_src_itemsize, wrappeddata);
@@ -410,52 +418,9 @@ _strided_to_strided_contig_align_wrap(char *dst, npy_intp dst_stride,
         else {
             tobuffer(bufferin, inner_src_itemsize, src, src_stride, N,
                                             src_itemsize, todata);
-            wrapped(bufferout, dst_itemsize, bufferin, inner_src_itemsize, N,
-                                            inner_src_itemsize, wrappeddata);
-            frombuffer(dst, dst_stride, bufferout, dst_itemsize, N,
-                                            dst_itemsize, fromdata);
-            return;
-        }
-    }
-}
-
-static void
-_strided_to_strided_contig_align_wrap_init_dest(char *dst, npy_intp dst_stride,
-                        char *src, npy_intp src_stride,
-                        npy_intp N, npy_intp src_itemsize,
-                        NpyAuxData *data)
-{
-    _align_wrap_data *d = (_align_wrap_data *)data;
-    PyArray_StridedUnaryOp *wrapped = d->wrapped,
-            *tobuffer = d->tobuffer,
-            *frombuffer = d->frombuffer;
-    npy_intp inner_src_itemsize = d->src_itemsize,
-             dst_itemsize = d->dst_itemsize;
-    NpyAuxData *wrappeddata = d->wrappeddata,
-            *todata = d->todata,
-            *fromdata = d->fromdata;
-    char *bufferin = d->bufferin, *bufferout = d->bufferout;
-
-    for(;;) {
-        if (N > NPY_LOWLEVEL_BUFFER_BLOCKSIZE) {
-            tobuffer(bufferin, inner_src_itemsize, src, src_stride,
-                                    NPY_LOWLEVEL_BUFFER_BLOCKSIZE,
-                                    src_itemsize, todata);
-            memset(bufferout, 0, dst_itemsize*NPY_LOWLEVEL_BUFFER_BLOCKSIZE);
-            wrapped(bufferout, dst_itemsize, bufferin, inner_src_itemsize,
-                                    NPY_LOWLEVEL_BUFFER_BLOCKSIZE,
-                                    inner_src_itemsize, wrappeddata);
-            frombuffer(dst, dst_stride, bufferout, dst_itemsize,
-                                    NPY_LOWLEVEL_BUFFER_BLOCKSIZE,
-                                    dst_itemsize, fromdata);
-            N -= NPY_LOWLEVEL_BUFFER_BLOCKSIZE;
-            src += NPY_LOWLEVEL_BUFFER_BLOCKSIZE*src_stride;
-            dst += NPY_LOWLEVEL_BUFFER_BLOCKSIZE*dst_stride;
-        }
-        else {
-            tobuffer(bufferin, inner_src_itemsize, src, src_stride, N,
-                                            src_itemsize, todata);
-            memset(bufferout, 0, dst_itemsize*N);
+            if (init_dest) {
+                memset(bufferout, 0, dst_itemsize*N);
+            }
             wrapped(bufferout, dst_itemsize, bufferin, inner_src_itemsize, N,
                                             inner_src_itemsize, wrappeddata);
             frombuffer(dst, dst_stride, bufferout, dst_itemsize, N,
@@ -519,14 +484,10 @@ wrap_aligned_contig_transfer_function(
     data->bufferin = (char *)data + basedatasize;
     data->bufferout = data->bufferin +
                 NPY_LOWLEVEL_BUFFER_BLOCKSIZE*src_itemsize;
+    data->init_dest = (npy_bool) init_dest;
 
     /* Set the function and data */
-    if (init_dest) {
-        *out_stransfer = &_strided_to_strided_contig_align_wrap_init_dest;
-    }
-    else {
-        *out_stransfer = &_strided_to_strided_contig_align_wrap;
-    }
+    *out_stransfer = &_strided_to_strided_contig_align_wrap;
     *out_transferdata = (NpyAuxData *)data;
 
     return NPY_SUCCEED;