From 64b47c314497498d198737ca348bb0ad91b1217c Mon Sep 17 00:00:00 2001
From: lyken <lyken@m-labs.hk>
Date: Tue, 20 Aug 2024 17:01:25 +0800
Subject: [PATCH] core/ndstrides: add ScalarOrNDArray::to_ndarray and
 NDArrayObject::make_unsized

---
 nac3core/src/codegen/object/ndarray/mod.rs | 30 ++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/nac3core/src/codegen/object/ndarray/mod.rs b/nac3core/src/codegen/object/ndarray/mod.rs
index ff771f15..8c63c78c 100644
--- a/nac3core/src/codegen/object/ndarray/mod.rs
+++ b/nac3core/src/codegen/object/ndarray/mod.rs
@@ -460,6 +460,22 @@ impl<'ctx> NDArrayObject<'ctx> {
 
         TupleObject::from_objects(generator, ctx, objects)
     }
+
+    /// Create an unsized ndarray to contain `object`.
+    pub fn make_unsized<G: CodeGenerator + ?Sized>(
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        object: AnyObject<'ctx>,
+    ) -> NDArrayObject<'ctx> {
+        // We have to put the value on the stack to get a data pointer.
+        let data = ctx.builder.build_alloca(object.value.get_type(), "make_unsized").unwrap();
+        ctx.builder.build_store(data, object.value).unwrap();
+        let data = Ptr(Int(Byte)).pointer_cast(generator, ctx, data);
+
+        let ndarray = NDArrayObject::alloca(generator, ctx, object.ty, 0);
+        ndarray.instance.set(ctx, |f| f.data, data);
+        ndarray
+    }
 }
 
 /// A convenience enum for implementing functions that acts on scalars or ndarrays or both.
@@ -500,4 +516,18 @@ impl<'ctx> ScalarOrNDArray<'ctx> {
             ScalarOrNDArray::NDArray(ndarray) => ndarray.instance.value.as_basic_value_enum(),
         }
     }
+
+    /// Convert this [`ScalarOrNDArray`] to an ndarray - behaves like `np.asarray`.
+    /// - If this is an ndarray, the ndarray is returned.
+    /// - If this is a scalar, this function returns new ndarray created with [`NDArrayObject::make_unsized`].
+    pub fn to_ndarray<G: CodeGenerator + ?Sized>(
+        &self,
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+    ) -> NDArrayObject<'ctx> {
+        match self {
+            ScalarOrNDArray::NDArray(ndarray) => *ndarray,
+            ScalarOrNDArray::Scalar(scalar) => NDArrayObject::make_unsized(generator, ctx, *scalar),
+        }
+    }
 }