From 6a92edee9bd679d1f71e66b2cc96456bead219fd Mon Sep 17 00:00:00 2001 From: Anton Khirnov Date: Wed, 17 Apr 2019 10:18:45 +0200 Subject: ndarray: support stepsizes other than 1 in slicing --- ndarray.c | 12 +++++++----- ndarray.h | 2 +- 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 }) -- cgit v1.2.3