ndarray: add new API for n-dimensional arrays

3 files changed, 243 insertions, 0 deletions

diff --git a/meson.build b/meson.build
index 73dfee9..3d5b035 100644
--- a/meson.build
+++ b/meson.build
@@ -10,6 +10,7 @@ lib_src = [
     'ell_grid_solve.c',
     'log.c',
     'mg2d.c',
+    'ndarray.c',
     'residual_calc.c',
 ]
 lib_obj = []
diff --git a/ndarray.c b/ndarray.c
new file mode 100644
index 0000000..57112ca
--- /dev/null
+++ b/ndarray.c
@@ -0,0 +1,187 @@
+/*
+ * N-dimensional strided array
+ * Copyright 2018 Anton Khirnov <anton@khirnov.net>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include <errno.h>
+#include <limits.h>
+#include <stddef.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "ndarray.h"
+
+struct NDArrayInternal {
+    size_t     alloc_size;
+    void      *data;
+
+    size_t    *shape;
+    ptrdiff_t *stride;
+};
+
+int mg2di_ndarray_alloc(NDArray **result, unsigned int dims,
+                        const size_t * const size, unsigned int flags)
+{
+    NDArray *a;
+    int ret;
+
+    if (!dims || dims >= INT_MAX)
+        return -EINVAL;
+
+    a = calloc(1, sizeof(*a));
+    if (!a)
+        return -ENOMEM;
+
+    a->priv = calloc(1, sizeof(*a->priv));
+    if (!a->priv) {
+        free(a);
+        return -ENOMEM;
+    }
+
+    a->priv->shape = calloc(dims, sizeof(*a->priv->shape));
+    if (!a->priv->shape) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+    memcpy(a->priv->shape, size, dims * sizeof(*a->priv->shape));
+
+    a->priv->stride = calloc(dims, sizeof(*a->priv->stride));
+    if (!a->priv->stride) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+
+    a->priv->stride[dims - 1] = 1;
+    for (int i = (int)dims - 2; i >= 0; i--)
+        a->priv->stride[i] = size[i + 1] * a->priv->stride[i + 1];
+
+    a->priv->alloc_size = a->priv->stride[0] * size[0] * sizeof(*a->data);
+    ret = posix_memalign((void**)&a->priv->data, 64, a->priv->alloc_size);
+    if (ret != 0) {
+        ret = -ret;
+        goto fail;
+    }
+
+    if (flags & NDARRAY_ALLOC_ZERO)
+        memset(a->priv->data, 0, a->priv->alloc_size);
+
+    a->dims   = dims;
+
+    a->data   = a->priv->data;
+    a->shape  = a->priv->shape;
+    a->stride = a->priv->stride;
+
+    *result = a;
+
+    return 0;
+fail:
+    mg2di_ndarray_free(&a);
+    return ret;
+}
+
+int mg2di_ndarray_slice(NDArray **result, NDArray *src,
+                        const Slice * const slice)
+{
+    NDArray *a;
+    ptrdiff_t offset;
+    int ret;
+
+    a = calloc(1, sizeof(*a));
+    if (!a)
+        return -ENOMEM;
+
+    a->priv = calloc(1, sizeof(*a->priv));
+    if (!a->priv) {
+        free(a);
+        return -ENOMEM;
+    }
+
+    a->priv->shape = calloc(src->dims, sizeof(*a->priv->shape));
+    if (!a->priv->shape) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+    memcpy(a->priv->shape, src->priv->shape, src->dims * sizeof(*a->priv->shape));
+
+    a->priv->stride = calloc(src->dims, sizeof(*a->priv->stride));
+    if (!a->priv->stride) {
+        ret = -ENOMEM;
+        goto fail;
+    }
+    memcpy(a->priv->stride, src->priv->stride, src->dims * sizeof(*a->priv->shape));
+
+    offset = 0;
+    for (unsigned int i = 0; i < src->dims; i++) {
+        const Slice *s = &slice[i];
+        ptrdiff_t start, end;
+
+        if (s->step != 1) {
+            ret = -ENOSYS;
+            goto fail;
+        }
+        if (s->start == PTRDIFF_MAX) {
+            start = 0;
+            end   = src->shape[i];
+        } else {
+            start = s->start;
+            end   = s->end;
+
+            if (start < 0)
+                start = src->shape[i] + start;
+            if (end <= 0)
+                end = src->shape[i] + end;
+
+            if (start >= end || end > src->shape[i]) {
+                ret = -EINVAL;
+                goto fail;
+            }
+
+            offset += src->stride[i] * start;
+            a->priv->shape[i] = end - start;
+        }
+    }
+
+    a->dims   = src->dims;
+
+    a->data   = src->data + offset;
+    a->shape  = a->priv->shape;
+    a->stride = a->priv->stride;
+
+    *result = a;
+    return 0;
+
+fail:
+    mg2di_ndarray_free(&a);
+    return ret;
+}
+
+void mg2di_ndarray_free(NDArray **pa)
+{
+    NDArray *a = *pa;
+
+    if (!a)
+        return;
+
+    if (a->priv) {
+        free(a->priv->data);
+        free(a->priv->shape);
+        free(a->priv->stride);
+    }
+    free(a->priv);
+
+    free(a);
+    *pa = NULL;
+}
diff --git a/ndarray.h b/ndarray.h
new file mode 100644
index 0000000..e264fba
--- /dev/null
+++ b/ndarray.h
@@ -0,0 +1,55 @@
+/*
+ * N-dimensional strided array
+ * Copyright 2018 Anton Khirnov <anton@khirnov.net>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef MG2D_NDARRAY_H
+#define MG2D_NDARRAY_H
+
+#include <stddef.h>
+#include <stdint.h>
+
+typedef struct Slice {
+    ptrdiff_t start;
+    ptrdiff_t end;
+    size_t    step;
+} Slice;
+
+#define SLICE(start, end, step) ((Slice){ start, end, step })
+#define SLICE_NULL ((Slice){ PTRDIFF_MAX })
+
+typedef struct NDArrayInternal NDArrayInternal;
+
+typedef struct NDArray {
+    NDArrayInternal *priv;
+
+    unsigned int     dims;
+
+    double          *data;
+    const size_t    *shape;
+    const ptrdiff_t *stride;
+} NDArray;
+
+#define NDARRAY_ALLOC_ZERO (1 << 0)
+
+int mg2di_ndarray_alloc(NDArray **result, unsigned int dims,
+                        const size_t * const size, unsigned int flags);
+void mg2di_ndarray_free(NDArray **a);
+
+int mg2di_ndarray_slice(NDArray **result, NDArray *src,
+                        const Slice * const slice);
+
+#endif // MG2D_ARRAY_H