ENH: core: Convert remaining data copying to new method, add tests for overlapping array assignment

7 files changed, 103 insertions, 122 deletions

diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
index f03821859..d81502831 100644
--- a/numpy/core/src/multiarray/convert_datatype.c
+++ b/numpy/core/src/multiarray/convert_datatype.c
@@ -120,17 +120,17 @@ PyArray_GetCastFunc(PyArray_Descr *descr, int type_num)
         }
 #if PY_VERSION_HEX >= 0x02050000
         ret = PyErr_WarnEx(cls,
-                           "Casting complex values to real discards the imaginary "
-                           "part", 1);
+                           "Casting complex values to real discards "
+                           "the imaginary part", 1);
 #else
         ret = PyErr_Warn(cls,
-                         "Casting complex values to real discards the imaginary "
-                         "part");
+                         "Casting complex values to real discards "
+                         "the imaginary part");
 #endif
         Py_XDECREF(cls);
         if (ret < 0) {
             return NULL;
-	}
+	    }
     }
     if (castfunc) {
         return castfunc;
@@ -177,7 +177,8 @@ PyArray_CanCastSafely(int fromtype, int totype)
     PyArray_Descr *from;
 
     /* Fast table lookup for small type numbers */
-    if ((unsigned int)fromtype < NPY_NTYPES && (unsigned int)totype < NPY_NTYPES) {
+    if ((unsigned int)fromtype < NPY_NTYPES &&
+                                (unsigned int)totype < NPY_NTYPES) {
         return _npy_can_cast_safely_table[fromtype][totype];
     }
 
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index de5ecb937..121277282 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -490,89 +490,6 @@ copy_and_swap(void *dst, void *src, int itemsize, npy_intp numitems,
     }
 }
 
-/*
- * Special-case of PyArray_CopyInto when dst is 1-d
- * and contiguous (and aligned).
- * PyArray_CopyInto requires broadcastable arrays while
- * this one is a flattening operation...
- *
- * TODO: Delete this function when its usage is removed.
- */
-NPY_NO_EXPORT int
-_flat_copyinto(PyObject *dst, PyObject *src, NPY_ORDER order)
-{
-    PyArrayIterObject *it;
-    PyObject *orig_src;
-    void (*myfunc)(char *, npy_intp, char *, npy_intp, npy_intp, int);
-    char *dptr;
-    int axis;
-    int elsize;
-    npy_intp nbytes;
-    NPY_BEGIN_THREADS_DEF;
-
-
-    orig_src = src;
-    if (PyArray_NDIM(src) == 0) {
-        /* Refcount note: src and dst have the same size */
-        PyArray_INCREF((PyArrayObject *)src);
-        PyArray_XDECREF((PyArrayObject *)dst);
-        NPY_BEGIN_THREADS;
-        memcpy(PyArray_BYTES(dst), PyArray_BYTES(src),
-                PyArray_ITEMSIZE(src));
-        NPY_END_THREADS;
-        return 0;
-    }
-
-    axis = PyArray_NDIM(src)-1;
-
-    if (order == PyArray_FORTRANORDER) {
-        if (PyArray_NDIM(src) <= 2) {
-            axis = 0;
-        }
-        /* fall back to a more general method */
-        else {
-            src = PyArray_Transpose((PyArrayObject *)orig_src, NULL);
-        }
-    }
-
-    it = (PyArrayIterObject *)PyArray_IterAllButAxis(src, &axis);
-    if (it == NULL) {
-        if (src != orig_src) {
-            Py_DECREF(src);
-        }
-        return -1;
-    }
-
-    if (PyArray_SAFEALIGNEDCOPY(src)) {
-        myfunc = _strided_byte_copy;
-    }
-    else {
-        myfunc = _unaligned_strided_byte_copy;
-    }
-
-    dptr = PyArray_BYTES(dst);
-    elsize = PyArray_ITEMSIZE(dst);
-    nbytes = elsize * PyArray_DIM(src, axis);
-
-    /* Refcount note: src and dst have the same size */
-    PyArray_INCREF((PyArrayObject *)src);
-    PyArray_XDECREF((PyArrayObject *)dst);
-    NPY_BEGIN_THREADS;
-    while(it->index < it->size) {
-        myfunc(dptr, elsize, it->dataptr, PyArray_STRIDE(src,axis),
-                PyArray_DIM(src,axis), elsize);
-        dptr += nbytes;
-        PyArray_ITER_NEXT(it);
-    }
-    NPY_END_THREADS;
-
-    if (src != orig_src) {
-        Py_DECREF(src);
-    }
-    Py_DECREF(it);
-    return 0;
-}
-
 /* Gets a half-open range [start, end) which contains the array data */
 void _get_memory_extents(PyArrayObject *arr,
                     npy_uintp *out_start, npy_uintp *out_end)
@@ -623,20 +540,21 @@ int _arrays_overlap(PyArrayObject *arr1, PyArrayObject *arr2)
  *
  * This is in general a difficult problem to solve efficiently, because
  * strides can be negative.  Consider "a = np.arange(3); a[::-1] = a", which
- * currently incorrectly produces [0, 1, 0].
+ * previously produced the incorrect [0, 1, 0].
  *
  * Instead of trying to be fancy, we simply check for overlap and make
  * a temporary copy when one exists.
- *
- * A special case is when there is just one dimension with positive
- * strides, and we pass that to CopyInto, which correctly handles
- * it for most cases.  It may still incorrectly handle copying of
- * partially-overlapping data elements, where the data pointer was offset
- * by a fraction of the element size.
  */
 NPY_NO_EXPORT int
 PyArray_MoveInto(PyArrayObject *dst, PyArrayObject *src)
 {
+    /*
+     * A special case is when there is just one dimension with positive
+     * strides, and we pass that to CopyInto, which correctly handles
+     * it for most cases.  It may still incorrectly handle copying of
+     * partially-overlapping data elements, where the data pointer was offset
+     * by a fraction of the element size.
+     */
     if ((PyArray_NDIM(dst) == 1 &&
                         PyArray_NDIM(src) == 1 &&
                         PyArray_STRIDE(dst, 0) > 0 &&
@@ -2255,19 +2173,10 @@ PyArray_EnsureAnyArray(PyObject *op)
     return PyArray_EnsureArray(op);
 }
 
-/*NUMPY_API
- * Copy an Array into another array -- memory must not overlap
- * Does not require src and dest to have "broadcastable" shapes
- * (only the same number of elements).
- *
- * TODO: For NumPy 2.0, this could accept an order parameter which
- *       only allows NPY_CORDER and NPY_FORDER.  Could also rename
- *       this to CopyAsFlat to make the name more intuitive.
- *
- * Returns 0 on success, -1 on error.
- */
+/* TODO: Put the order parameter in PyArray_CopyAnyInto and remove this */
 NPY_NO_EXPORT int
-PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
+PyArray_CopyAnyIntoOrdered(PyArrayObject *dst, PyArrayObject *src,
+                                NPY_ORDER order)
 {
     PyArray_StridedTransferFn *stransfer = NULL;
     void *transferdata = NULL;
@@ -2291,6 +2200,12 @@ PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
         return -1;
     }
 
+    if (!(order == NPY_CORDER || order == NPY_FORTRANORDER)) {
+        PyErr_SetString(PyExc_RuntimeError,
+                "CopyAnyIntoOrdered requires either C or Fortran order");
+        return -1;
+    }
+
     /* If the shapes match, use the more efficient CopyInto */
     if (PyArray_NDIM(dst) == PyArray_NDIM(src) &&
             PyArray_CompareLists(PyArray_DIMS(dst), PyArray_DIMS(src),
@@ -2321,7 +2236,7 @@ PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
     dst_iter = NpyIter_New(dst, NPY_ITER_WRITEONLY|
                                 NPY_ITER_NO_INNER_ITERATION|
                                 NPY_ITER_REFS_OK,
-                                NPY_CORDER,
+                                order,
                                 NPY_NO_CASTING,
                                 NULL, 0, NULL, 0);
     if (dst_iter == NULL) {
@@ -2330,7 +2245,7 @@ PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
     src_iter = NpyIter_New(src, NPY_ITER_READONLY|
                                 NPY_ITER_NO_INNER_ITERATION|
                                 NPY_ITER_REFS_OK,
-                                NPY_CORDER,
+                                order,
                                 NPY_NO_CASTING,
                                 NULL, 0, NULL, 0);
     if (src_iter == NULL) {
@@ -2439,6 +2354,23 @@ PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
 }
 
 /*NUMPY_API
+ * Copy an Array into another array -- memory must not overlap
+ * Does not require src and dest to have "broadcastable" shapes
+ * (only the same number of elements).
+ *
+ * TODO: For NumPy 2.0, this could accept an order parameter which
+ *       only allows NPY_CORDER and NPY_FORDER.  Could also rename
+ *       this to CopyAsFlat to make the name more intuitive.
+ *
+ * Returns 0 on success, -1 on error.
+ */
+NPY_NO_EXPORT int
+PyArray_CopyAnyInto(PyArrayObject *dst, PyArrayObject *src)
+{
+    return PyArray_CopyAnyIntoOrdered(dst, src, NPY_CORDER);
+}
+
+/*NUMPY_API
  * Copy an Array into another array -- memory must not overlap.
  * Broadcast to the destination shape if necessary.
  *
diff --git a/numpy/core/src/multiarray/ctors.h b/numpy/core/src/multiarray/ctors.h
index 5fd1d8e58..eb1586f03 100644
--- a/numpy/core/src/multiarray/ctors.h
+++ b/numpy/core/src/multiarray/ctors.h
@@ -45,10 +45,12 @@ PyArray_CopyAnyInto(PyArrayObject *dest, PyArrayObject *src);
 NPY_NO_EXPORT PyObject *
 PyArray_CheckAxis(PyArrayObject *arr, int *axis, int flags);
 
-/* FIXME: remove those from here */
+/* TODO: Put the order parameter in PyArray_CopyAnyInto and remove this */
 NPY_NO_EXPORT int
-_flat_copyinto(PyObject *dst, PyObject *src, NPY_ORDER order);
+PyArray_CopyAnyIntoOrdered(PyArrayObject *dst, PyArrayObject *src,
+                                NPY_ORDER order);
 
+/* FIXME: remove those from here */
 NPY_NO_EXPORT size_t
 _array_fill_strides(intp *strides, intp *dims, int nd, size_t itemsize,
                     int inflag, int *objflags);
diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index 5990f8c3d..bb996e3bc 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -1155,8 +1155,7 @@ iter_array(PyArrayIterObject *it, PyObject *NPY_UNUSED(op))
         if (r == NULL) {
             return NULL;
         }
-        if (_flat_copyinto(r, (PyObject *)it->ao,
-                           PyArray_CORDER) < 0) {
+        if (PyArray_CopyAnyInto(r, it->ao) < 0) {
             Py_DECREF(r);
             return NULL;
         }
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index d3b4dfe6e..e2ccb3f80 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -755,7 +755,7 @@ array_setasflat(PyArrayObject *self, PyObject *args)
         return NULL;
     }
 
-    arr = PyArray_FromAny(arr_in, NULL, 0, 0, 0, NULL);
+    arr = (PyArrayObject *)PyArray_FromAny(arr_in, NULL, 0, 0, 0, NULL);
     if (arr == NULL) {
         return NULL;
     }
diff --git a/numpy/core/src/multiarray/shape.c b/numpy/core/src/multiarray/shape.c
index 671dc1538..f8005d58e 100644
--- a/numpy/core/src/multiarray/shape.c
+++ b/numpy/core/src/multiarray/shape.c
@@ -791,15 +791,15 @@ PyArray_Ravel(PyArrayObject *a, NPY_ORDER fortran)
 NPY_NO_EXPORT PyObject *
 PyArray_Flatten(PyArrayObject *a, NPY_ORDER order)
 {
-    PyObject *ret;
+    PyArrayObject *ret;
     intp size;
 
-    if (order == PyArray_ANYORDER) {
-        order = PyArray_ISFORTRAN(a);
+    if (order == NPY_ANYORDER) {
+        order = PyArray_ISFORTRAN(a) ? NPY_FORTRANORDER : NPY_CORDER;
     }
     size = PyArray_SIZE(a);
     Py_INCREF(a->descr);
-    ret = PyArray_NewFromDescr(Py_TYPE(a),
+    ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(a),
                                a->descr,
                                1, &size,
                                NULL,
@@ -809,11 +809,11 @@ PyArray_Flatten(PyArrayObject *a, NPY_ORDER order)
     if (ret == NULL) {
         return NULL;
     }
-    if (_flat_copyinto(ret, (PyObject *)a, order) < 0) {
+    if (PyArray_CopyAnyIntoOrdered(ret, a, order) < 0) {
         Py_DECREF(ret);
         return NULL;
     }
-    return ret;
+    return (PyObject *)ret;
 }
 
 
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index d30ac13e8..1c5e84537 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -238,7 +238,54 @@ class TestScalarIndexing(TestCase):
         def subscript(x, i): x[i]
         self.assertRaises(IndexError, subscript, a, (newaxis, 0))
         self.assertRaises(IndexError, subscript, a, (newaxis,)*50)
-
+    
+    def test_overlapping_assignment(self):
+        # With positive strides
+        a = np.arange(4)
+        a[:-1] = a[1:]
+        assert_equal(a, [1,2,3,3])
+
+        a = np.arange(4)
+        a[1:] = a[:-1]
+        assert_equal(a, [0,0,1,2])
+
+        # With positive and negative strides
+        a = np.arange(4)
+        a[:] = a[::-1]
+        assert_equal(a, [3,2,1,0])
+
+        a = np.arange(6).reshape(2,3)
+        a[::-1,:] = a[:,::-1]
+        assert_equal(a, [[5,4,3],[2,1,0]])
+
+        a = np.arange(6).reshape(2,3)
+        a[::-1,::-1] = a[:,::-1]
+        assert_equal(a, [[3,4,5],[0,1,2]])
+
+        # With just one element overlapping
+        a = np.arange(5)
+        a[:3] = a[2:]
+        assert_equal(a, [2,3,4,3,4])
+
+        a = np.arange(5)
+        a[2:] = a[:3]
+        assert_equal(a, [0,1,0,1,2])
+
+        a = np.arange(5)
+        a[2::-1] = a[2:]
+        assert_equal(a, [4,3,2,3,4])
+
+        a = np.arange(5)
+        a[2:] = a[2::-1]
+        assert_equal(a, [0,1,2,1,0])
+
+        a = np.arange(5)
+        a[2::-1] = a[:1:-1]
+        assert_equal(a, [2,3,4,3,4])
+
+        a = np.arange(5)
+        a[:1:-1] = a[2::-1]
+        assert_equal(a, [0,1,0,1,2])
 
 class TestCreation(TestCase):
     def test_from_attribute(self):