From fd3d02bff0b4c28e9e3137d49ce9f6671564e52f 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         |  1 +
 nac3core/src/codegen/structure/mod.rs     |  1 +
 nac3core/src/codegen/structure/ndarray.rs | 27 ++++++++++++++
 4 files changed, 73 insertions(+)
 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 fa775263..16ea68e7 100644
--- a/nac3core/irrt/irrt_everything.hpp
+++ b/nac3core/irrt/irrt_everything.hpp
@@ -4,4 +4,5 @@
 #include <irrt/core.hpp>
 #include <irrt/error_context.hpp>
 #include <irrt/int_defs.hpp>
+#include <irrt/ndarray/def.hpp>
 #include <irrt/utils.hpp>
\ No newline at end of file
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..79ea4e7c
--- /dev/null
+++ b/nac3core/src/codegen/structure/ndarray.rs
@@ -0,0 +1,27 @@
+use crate::codegen::*;
+
+pub struct NpArrayFields<F: FieldVisitor> {
+    pub data: F::Field<PtrModel<IntModel<Byte>>>,
+    pub itemsize: F::Field<IntModel<SizeT>>,
+    pub ndims: F::Field<IntModel<SizeT>>,
+    pub shape: F::Field<PtrModel<IntModel<SizeT>>>,
+    pub strides: F::Field<PtrModel<IntModel<SizeT>>>,
+}
+
+// TODO: Rename to `NDArray` when the old NDArray is removed.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct NpArray;
+
+impl StructKind for NpArray {
+    type Fields<F: FieldVisitor> = NpArrayFields<F>;
+
+    fn visit_fields<F: FieldVisitor>(&self, visitor: &mut F) -> Self::Fields<F> {
+        Self::Fields {
+            data: visitor.add("data"),
+            itemsize: visitor.add("itemsize"),
+            ndims: visitor.add("ndims"),
+            shape: visitor.add("shape"),
+            strides: visitor.add("strides"),
+        }
+    }
+}