ndarray: add a function for copying arrays

2 files changed, 34 insertions, 0 deletions

diff --git a/ndarray.c b/ndarray.c
index 57112ca..a733955 100644
--- a/ndarray.c
+++ b/ndarray.c
@@ -185,3 +185,35 @@ void mg2di_ndarray_free(NDArray **pa)
     free(a);
     *pa = NULL;
 }
+
+static int copy_axis(NDArray *dst, const NDArray *src, unsigned int axis,
+                     ptrdiff_t offset_dst, ptrdiff_t offset_src)
+{
+    if (dst->shape[axis] != src->shape[axis])
+        return -EINVAL;
+
+    if (axis == dst->dims - 1) {
+        if (dst->stride[axis] == 1 && src->stride[axis] == 1)
+            memcpy(dst->data + offset_dst, src->data + offset_src, sizeof(*dst->data) * dst->shape[axis]);
+        else {
+            for (size_t idx = 0; idx < dst->shape[axis]; idx++)
+                dst->data[offset_dst + idx * dst->stride[axis]] = src->data[offset_src + idx * src->stride[axis]];
+        }
+        return 0;
+    }
+
+    for (size_t idx = 0; idx < dst->shape[axis]; idx++)
+        copy_axis(dst, src, axis + 1, offset_dst + idx * dst->stride[axis], offset_src + idx * src->stride[axis]);
+
+    return 0;
+}
+
+int mg2di_ndarray_copy(NDArray *dst, const NDArray *src)
+{
+    const unsigned int dims = src->dims;
+
+    if (dims != dst->dims)
+        return -EINVAL;
+
+    return copy_axis(dst, src, 0, 0, 0);
+}
diff --git a/ndarray.h b/ndarray.h
index e264fba..f10ac80 100644
--- a/ndarray.h
+++ b/ndarray.h
@@ -52,4 +52,6 @@ void mg2di_ndarray_free(NDArray **a);
 int mg2di_ndarray_slice(NDArray **result, NDArray *src,
                         const Slice * const slice);
 
+int mg2di_ndarray_copy(NDArray *dst, const NDArray *src);
+
 #endif // MG2D_ARRAY_H