From b2b46939b532fe49eb6c24b64ef0147434a22d17 Mon Sep 17 00:00:00 2001
From: lyken <lyken@m-labs.hk>
Date: Sun, 28 Jul 2024 14:23:31 +0800
Subject: [PATCH] core/ndstrides: add NDArray with strides definition

---
 nac3core/irrt/irrt/ndarray/def.hpp        | 44 +++++++++++++++++++++++
 nac3core/irrt/irrt_everything.hpp         |  3 +-
 nac3core/src/codegen/structure/mod.rs     |  1 +
 nac3core/src/codegen/structure/ndarray.rs | 34 ++++++++++++++++++
 4 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 nac3core/irrt/irrt/ndarray/def.hpp
 create mode 100644 nac3core/src/codegen/structure/ndarray.rs

diff --git a/nac3core/irrt/irrt/ndarray/def.hpp b/nac3core/irrt/irrt/ndarray/def.hpp
new file mode 100644
index 00000000..4b2c1216
--- /dev/null
+++ b/nac3core/irrt/irrt/ndarray/def.hpp
@@ -0,0 +1,44 @@
+#pragma once
+
+namespace {
+/**
+ * @brief The NDArray object
+ *
+ * The official numpy implementations: https://github.com/numpy/numpy/blob/735a477f0bc2b5b84d0e72d92f224bde78d4e069/doc/source/reference/c-api/types-and-structures.rst
+ */
+template <typename SizeT>
+struct NDArray {
+    /**
+     * @brief The underlying data this `ndarray` is pointing to.
+     *
+     * Must be set to `nullptr` to indicate that this NDArray's `data` is uninitialized.
+     */
+    uint8_t* data;
+
+    /**
+     * @brief The number of bytes of a single element in `data`.
+     */
+    SizeT itemsize;
+
+    /**
+     * @brief The number of dimensions of this shape.
+     */
+    SizeT ndims;
+
+    /**
+     * @brief The NDArray shape, with length equal to `ndims`.
+     *
+     * Note that it may contain 0.
+     */
+    SizeT* shape;
+
+    /**
+     * @brief Array strides, with length equal to `ndims`
+     *
+     * The stride values are in units of bytes, not number of elements.
+     *
+     * Note that `strides` can have negative values.
+     */
+    SizeT* strides;
+};
+}  // namespace
\ No newline at end of file
diff --git a/nac3core/irrt/irrt_everything.hpp b/nac3core/irrt/irrt_everything.hpp
index 636a011c..d1db2364 100644
--- a/nac3core/irrt/irrt_everything.hpp
+++ b/nac3core/irrt/irrt_everything.hpp
@@ -3,4 +3,5 @@
 #include <irrt/core.hpp>
 #include <irrt/exception.hpp>
 #include <irrt/int_defs.hpp>
-#include <irrt/util.hpp>
\ No newline at end of file
+#include <irrt/ndarray/def.hpp>
+#include <irrt/util.hpp>
diff --git a/nac3core/src/codegen/structure/mod.rs b/nac3core/src/codegen/structure/mod.rs
index 14b596fb..51c9cab0 100644
--- a/nac3core/src/codegen/structure/mod.rs
+++ b/nac3core/src/codegen/structure/mod.rs
@@ -1,2 +1,3 @@
 pub mod cslice;
 pub mod exception;
+pub mod ndarray;
diff --git a/nac3core/src/codegen/structure/ndarray.rs b/nac3core/src/codegen/structure/ndarray.rs
new file mode 100644
index 00000000..a7313c22
--- /dev/null
+++ b/nac3core/src/codegen/structure/ndarray.rs
@@ -0,0 +1,34 @@
+use crate::codegen::*;
+
+pub struct NpArrayFields<'ctx, F: FieldTraversal<'ctx>> {
+    pub data: F::Out<PtrModel<IntModel<Byte>>>,
+    pub itemsize: F::Out<IntModel<SizeT>>,
+    pub ndims: F::Out<IntModel<SizeT>>,
+    pub shape: F::Out<PtrModel<IntModel<SizeT>>>,
+    pub strides: F::Out<PtrModel<IntModel<SizeT>>>,
+}
+
+// TODO: Rename to `NDArray` when the old NDArray is removed.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct NpArray;
+
+impl<'ctx> StructKind<'ctx> for NpArray {
+    type Fields<F: FieldTraversal<'ctx>> = NpArrayFields<'ctx, F>;
+
+    fn traverse_fields<F: FieldTraversal<'ctx>>(&self, traversal: &mut F) -> Self::Fields<F> {
+        Self::Fields {
+            data: traversal.add_auto("data"),
+            itemsize: traversal.add_auto("itemsize"),
+            ndims: traversal.add_auto("ndims"),
+            shape: traversal.add_auto("shape"),
+            strides: traversal.add_auto("strides"),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct NDArrayObject<'ctx> {
+    pub dtype: Type,
+    pub ndims: Type,
+    pub value: Ptr<'ctx, StructModel<NpArray>>,
+}