From afa7d9b100321b60e1e58b9f81bb6616870179ca Mon Sep 17 00:00:00 2001
From: David Mak <csdavidmak@ust.hk>
Date: Thu, 21 Dec 2023 14:43:19 +0800
Subject: [PATCH] core: Implement helper for creation of generic ndarray

---
 nac3core/src/typecheck/typedef/mod.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/nac3core/src/typecheck/typedef/mod.rs b/nac3core/src/typecheck/typedef/mod.rs
index e0a72e8..ae84b0a 100644
--- a/nac3core/src/typecheck/typedef/mod.rs
+++ b/nac3core/src/typecheck/typedef/mod.rs
@@ -205,6 +205,27 @@ impl TypeEnum {
             TypeEnum::TFunc { .. } => "TFunc",
         }
     }
+
+    /// Returns a [TypeEnum] representing a generic `ndarray` type.
+    ///
+    /// * `dtype` - The datatype of the `ndarray`, or `None` if the datatype is generic.
+    /// * `ndims` - The number of dimensions of the `ndarray`, or `None` if the number of dimensions is generic.
+    #[must_use]
+    pub fn ndarray(
+        unifier: &mut Unifier,
+        dtype: Option<Type>,
+        ndims: Option<Type>,
+        primitives: &PrimitiveStore
+    ) -> TypeEnum {
+        let dtype = dtype.unwrap_or_else(|| unifier.get_fresh_var(Some("T".into()), None).0);
+        let ndims = ndims
+            .unwrap_or_else(|| unifier.get_fresh_const_generic_var(primitives.usize(), Some("N".into()), None).0);
+
+        TypeEnum::TNDArray {
+            ty: dtype,
+            ndims,
+        }
+    }
 }
 
 pub type SharedUnifier = Arc<Mutex<(UnificationTable<TypeEnum>, u32, Vec<Call>)>>;