From afa7d9b100321b60e1e58b9f81bb6616870179ca Mon Sep 17 00:00:00 2001 From: David Mak 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 e0a72e8c4..ae84b0a68 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, + ndims: Option, + 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, u32, Vec)>>;