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( + 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( + &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), + } + } }