ndarray: support stepsizes other than 1 in slicing

2 files changed, 8 insertions, 6 deletions

diff --git a/ndarray.c b/ndarray.c
index 504bfb7..b7c30b9 100644
--- a/ndarray.c
+++ b/ndarray.c
@@ -128,8 +128,8 @@ int mg2di_ndarray_slice(NDArray **result, NDArray *src,
         const Slice *s = &slice[i];
         ptrdiff_t start, end;
 
-        if (s->step != 1) {
-            ret = -ENOSYS;
+        if (s->step < 1) {
+            ret = s->step ? -ENOSYS : -EINVAL;
             goto fail;
         }
         if (s->start == PTRDIFF_MAX) {
@@ -148,10 +148,12 @@ int mg2di_ndarray_slice(NDArray **result, NDArray *src,
                 ret = -EINVAL;
                 goto fail;
             }
-
-            offset += src->stride[i] * start;
-            a->priv->shape[i] = end - start;
         }
+
+        offset += src->stride[i] * start;
+
+        a->priv->stride[i] *= s->step;
+        a->priv->shape[i]   = (end - start) / s->step;
     }
 
     a->dims   = src->dims;
diff --git a/ndarray.h b/ndarray.h
index 9e1b8ba..42cb64f 100644
--- a/ndarray.h
+++ b/ndarray.h
@@ -25,7 +25,7 @@
 typedef struct Slice {
     ptrdiff_t start;
     ptrdiff_t end;
-    size_t    step;
+    ptrdiff_t step;
 } Slice;
 
 #define SLICE(start, end, step) ((Slice){ start, end, step })