2 files changed, 102 insertions, 0 deletions
diff --git a/circuitpython/extmod/ulab/code/numpy/ndarray/ndarray_iter.c b/circuitpython/extmod/ulab/code/numpy/ndarray/ndarray_iter.c
new file mode 100644
index 0000000..8704836
--- /dev/null
+++ b/circuitpython/extmod/ulab/code/numpy/ndarray/ndarray_iter.c
@@ -0,0 +1,66 @@
+
+/*
+ * This file is part of the micropython-ulab project,
+ *
+ * https://github.com/v923z/micropython-ulab
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2021 Zoltán Vörös
+ *
+*/
+
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include "py/obj.h"
+#include "py/runtime.h"
+
+#include "ndarray_iter.h"
+
+#ifdef NDARRAY_HAS_FLATITER
+mp_obj_t ndarray_flatiter_make_new(mp_obj_t self_in) {
+    ndarray_flatiter_t *flatiter = m_new_obj(ndarray_flatiter_t);
+    flatiter->base.type = &ndarray_flatiter_type;
+    flatiter->iternext = ndarray_flatiter_next;
+    flatiter->ndarray = MP_OBJ_TO_PTR(self_in);
+    flatiter->cur = 0;
+    return flatiter;
+}
+
+mp_obj_t ndarray_flatiter_next(mp_obj_t self_in) {
+    ndarray_flatiter_t *self = MP_OBJ_TO_PTR(self_in);
+    ndarray_obj_t *ndarray = MP_OBJ_TO_PTR(self->ndarray);
+    uint8_t *array = (uint8_t *)ndarray->array;
+
+    if(self->cur < ndarray->len) {
+        uint32_t remainder = self->cur;
+        uint8_t i = ULAB_MAX_DIMS - 1;
+        do {
+            size_t div = (remainder / ndarray->shape[i]);
+            array += remainder * ndarray->strides[i];
+            remainder -= div * ndarray->shape[i];
+            i--;
+        } while(i > ULAB_MAX_DIMS - ndarray->ndim);
+        self->cur++;
+        return ndarray_get_item(ndarray, array);
+    }
+    return MP_OBJ_STOP_ITERATION;
+}
+
+mp_obj_t ndarray_new_flatiterator(mp_obj_t flatiter_in, mp_obj_iter_buf_t *iter_buf) {
+    assert(sizeof(ndarray_flatiter_t) <= sizeof(mp_obj_iter_buf_t));
+    ndarray_flatiter_t *iter = (ndarray_flatiter_t *)iter_buf;
+    ndarray_flatiter_t *flatiter = MP_OBJ_TO_PTR(flatiter_in);
+    iter->base.type = &mp_type_polymorph_iter;
+    iter->iternext = ndarray_flatiter_next;
+    iter->ndarray = flatiter->ndarray;
+    iter->cur = 0;
+    return MP_OBJ_FROM_PTR(iter);
+}
+
+mp_obj_t ndarray_get_flatiterator(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) {
+    return ndarray_new_flatiterator(o_in, iter_buf);
+}
+#endif /* NDARRAY_HAS_FLATITER */
diff --git a/circuitpython/extmod/ulab/code/numpy/ndarray/ndarray_iter.h b/circuitpython/extmod/ulab/code/numpy/ndarray/ndarray_iter.h
new file mode 100644
index 0000000..b3fc48d
--- /dev/null
+++ b/circuitpython/extmod/ulab/code/numpy/ndarray/ndarray_iter.h
@@ -0,0 +1,36 @@
+
+/*
+ * This file is part of the micropython-ulab project,
+ *
+ * https://github.com/v923z/micropython-ulab
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2020 Jeff Epler for Adafruit Industries
+ *               2020-2021 Zoltán Vörös
+*/
+
+#ifndef _NDARRAY_ITER_
+#define _NDARRAY_ITER_
+
+#include "py/runtime.h"
+#include "py/binary.h"
+#include "py/obj.h"
+#include "py/objarray.h"
+
+#include "../../ulab.h"
+#include "../../ndarray.h"
+
+// TODO: take simply mp_obj_ndarray_it_t from ndarray.c
+typedef struct _mp_obj_ndarray_flatiter_t {
+    mp_obj_base_t base;
+    mp_fun_1_t iternext;
+    mp_obj_t ndarray;
+    size_t cur;
+} ndarray_flatiter_t;
+
+mp_obj_t ndarray_get_flatiterator(mp_obj_t , mp_obj_iter_buf_t *);
+mp_obj_t ndarray_flatiter_make_new(mp_obj_t );
+mp_obj_t ndarray_flatiter_next(mp_obj_t );
+
+#endif
+\ No newline at end of file