core/ndstrides: move functions to numpy_new/util.rs

nac3core/src/codegen/expr.rs

@@ -18,6 +18,7 @@ use crate::{
             call_memcpy_generic,
         },
         need_sret, numpy,
+        numpy_new::util::alloca_ndarray,
         stmt::{
             gen_for_callback_incrementing, gen_if_callback, gen_if_else_expr_callback, gen_raise,
             gen_var,
@@ -42,10 +43,8 @@ use nac3parser::ast::{
     self, Boolop, Cmpop, Comprehension, Constant, Expr, ExprKind, Located, Location, Operator,
     StrRef, Unaryop,
 };
-use ndarray::{
+
-    allocation::alloca_ndarray,
+use ndarray::indexing::{call_nac3_ndarray_index, RustNDIndex};
-    indexing::{call_nac3_ndarray_index, RustNDIndex},
-};
 
 use super::{
     model::*,

nac3core/src/codegen/irrt/ndarray/allocation.rs

@@ -1,83 +0,0 @@
-use crate::codegen::model::*;
-use crate::codegen::util::array_writer::ArrayWriter;
-use crate::codegen::{structure::ndarray::NpArray, CodeGenContext, CodeGenerator};
-
-use super::basic::{
-    call_nac3_ndarray_nbytes, call_nac3_ndarray_set_strides_by_shape,
-    call_nac3_ndarray_util_assert_shape_no_negative,
-};
-
-/**
-    Allocate an ndarray on the stack given its `ndims`.
-
-    `shape` and `strides` will be automatically allocated on the stack.
-
-    The returned ndarray's content will be:
-    - `data`: `nullptr`
-    - `itemsize`: **uninitialized** value
-    - `ndims`: initialized value, set to the input `ndims`
-    - `shape`: initialized pointer to an allocated stack with **uninitialized** values
-    - `strides`: initialized pointer to an allocated stack with **uninitialized** values
-*/
-pub fn alloca_ndarray<'ctx, G>(
-    generator: &mut G,
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    ndims: Int<'ctx, SizeT>,
-    name: &str,
-) -> Result<Ptr<'ctx, StructModel<NpArray>>, String>
-where
-    G: CodeGenerator + ?Sized,
-{
-    let tyctx = generator.type_context(ctx.ctx);
-
-    let sizet_model = IntModel(SizeT);
-    let ndarray_model = StructModel(NpArray);
-    let ndarray_data_model = PtrModel(IntModel(Byte));
-
-    // Setup ndarray
-    let ndarray_ptr = ndarray_model.alloca(tyctx, ctx, name);
-    let shape = sizet_model.array_alloca(tyctx, ctx, ndims.value, "shape");
-    let strides = sizet_model.array_alloca(tyctx, ctx, ndims.value, "strides");
-
-    ndarray_ptr.gep(ctx, |f| f.data).store(ctx, ndarray_data_model.nullptr(tyctx, ctx.ctx));
-    ndarray_ptr.gep(ctx, |f| f.ndims).store(ctx, ndims);
-    ndarray_ptr.gep(ctx, |f| f.shape).store(ctx, shape);
-    ndarray_ptr.gep(ctx, |f| f.strides).store(ctx, strides);
-
-    Ok(ndarray_ptr)
-}
-
-/// Initialize an ndarray's `shape` and asserts on.
-/// `shape`'s values and prohibit illegal inputs like negative dimensions.
-pub fn init_ndarray_shape<'ctx, G: CodeGenerator + ?Sized>(
-    generator: &mut G,
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    pndarray: Ptr<'ctx, StructModel<NpArray>>,
-    shape_writer: &ArrayWriter<'ctx, G, SizeT, IntModel<SizeT>>,
-) -> Result<(), String> {
-    let tyctx = generator.type_context(ctx.ctx);
-    let shape = pndarray.gep(ctx, |f| f.shape).load(tyctx, ctx, "shape");
-    (shape_writer.write)(generator, ctx, shape)?;
-    call_nac3_ndarray_util_assert_shape_no_negative(generator, ctx, shape_writer.len, shape);
-    Ok(())
-}
-
-/// Initialize an ndarray's `data` by allocating a buffer on the stack.
-/// The allocated data buffer is considered to be *owned* by the ndarray.
-///
-/// `strides` of the ndarray will also be updated with `set_strides_by_shape`.
-///
-/// `shape` and `itemsize` of the ndarray ***must*** be initialized first.
-pub fn init_ndarray_data_by_alloca<'ctx, G: CodeGenerator + ?Sized>(
-    generator: &mut G,
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    pndarray: Ptr<'ctx, StructModel<NpArray>>,
-) {
-    let tyctx = generator.type_context(ctx.ctx);
-    let ndarray_data_model = IntModel(Byte);
-
-    let nbytes = call_nac3_ndarray_nbytes(generator, ctx, pndarray);
-    let data = ndarray_data_model.array_alloca(tyctx, ctx, nbytes.value, "data");
-    pndarray.gep(ctx, |f| f.data).store(ctx, data);
-    call_nac3_ndarray_set_strides_by_shape(generator, ctx, pndarray);
-}

nac3core/src/codegen/irrt/ndarray/mod.rs

@@ -1,4 +1,3 @@
-pub mod allocation;
 pub mod basic;
 pub mod indexing;
 pub mod reshape;

nac3core/src/codegen/numpy_new/control.rs

@@ -1,49 +0,0 @@
-use crate::codegen::{
-    irrt::ndarray::basic::{call_nac3_ndarray_get_nth_pelement, call_nac3_ndarray_size},
-    model::*,
-    stmt::BreakContinueHooks,
-    structure::ndarray::NpArray,
-    util::control::gen_model_for,
-    CodeGenContext, CodeGenerator,
-};
-
-/// Iterate through all elements in an ndarray.
-///
-/// `body` is given the index of an element and an opaque pointer (as an `uint8_t*`, you might want to cast it) to the element.
-///
-/// Short-circuiting is possible with the given [`BreakContinueHooks`].
-pub fn gen_foreach_ndarray_elements<'ctx, G, F>(
-    generator: &mut G,
-    ctx: &mut CodeGenContext<'ctx, '_>,
-    pndarray: Ptr<'ctx, StructModel<NpArray>>,
-    body: F,
-) -> Result<(), String>
-where
-    G: CodeGenerator + ?Sized,
-    F: Fn(
-        &mut G,
-        &mut CodeGenContext<'ctx, '_>,
-        BreakContinueHooks<'ctx>,
-        Int<'ctx, SizeT>,
-        Ptr<'ctx, IntModel<Byte>>,
-    ) -> Result<(), String>,
-{
-    // TODO: Make this more efficient - use a special NDArray iterator?
-
-    let tyctx = generator.type_context(ctx.ctx);
-
-    let sizet_model = IntModel(SizeT);
-    let size = call_nac3_ndarray_size(generator, ctx, pndarray);
-
-    gen_model_for(
-        generator,
-        ctx,
-        sizet_model.const_0(tyctx, ctx.ctx),
-        size,
-        sizet_model.const_1(tyctx, ctx.ctx),
-        |generator, ctx, hooks, index| {
-            let pelement = call_nac3_ndarray_get_nth_pelement(generator, ctx, pndarray, index);
-            body(generator, ctx, hooks, index, pelement)
-        },
-    )
-}

nac3core/src/codegen/numpy_new/factory.rs

@@ -7,10 +7,8 @@ use nac3parser::ast::StrRef;
 
 use crate::{
     codegen::{
-        irrt::ndarray::allocation::{
-            alloca_ndarray, init_ndarray_data_by_alloca, init_ndarray_shape,
-        },
         model::*,
+        numpy_new::util::{alloca_ndarray, init_ndarray_data_by_alloca, init_ndarray_shape},
         structure::ndarray::NpArray,
         util::shape::make_shape_writer,
         CodeGenContext, CodeGenerator,
@@ -20,7 +18,7 @@ use crate::{
     typecheck::typedef::{FunSignature, Type},
 };
 
-use super::control::gen_foreach_ndarray_elements;
+use super::util::gen_foreach_ndarray_elements;
 
 /// Helper function to create an ndarray with uninitialized values
 ///

nac3core/src/codegen/numpy_new/mod.rs

@@ -1,3 +1,3 @@
-pub mod control;
 pub mod factory;
+pub mod util;
 pub mod view;

nac3core/src/codegen/numpy_new/util.rs

@@ -0,0 +1,184 @@
+use inkwell::{types::BasicType, values::BasicValueEnum};
+
+use crate::{
+    codegen::{
+        irrt::ndarray::basic::{
+            call_nac3_ndarray_get_nth_pelement, call_nac3_ndarray_nbytes,
+            call_nac3_ndarray_set_strides_by_shape, call_nac3_ndarray_size,
+            call_nac3_ndarray_util_assert_shape_no_negative,
+        },
+        model::*,
+        stmt::BreakContinueHooks,
+        structure::ndarray::NpArray,
+        util::{array_writer::ArrayWriter, control::gen_model_for},
+        CodeGenContext, CodeGenerator,
+    },
+    toplevel::numpy::unpack_ndarray_var_tys,
+    typecheck::typedef::{Type, TypeEnum},
+};
+
+/**
+    Allocate an ndarray on the stack given its `ndims`.
+
+    `shape` and `strides` will be automatically allocated on the stack.
+
+    The returned ndarray's content will be:
+    - `data`: `nullptr`
+    - `itemsize`: **uninitialized** value
+    - `ndims`: initialized value, set to the input `ndims`
+    - `shape`: initialized pointer to an allocated stack with **uninitialized** values
+    - `strides`: initialized pointer to an allocated stack with **uninitialized** values
+*/
+pub fn alloca_ndarray<'ctx, G>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    ndims: Int<'ctx, SizeT>,
+    name: &str,
+) -> Result<Ptr<'ctx, StructModel<NpArray>>, String>
+where
+    G: CodeGenerator + ?Sized,
+{
+    let tyctx = generator.type_context(ctx.ctx);
+
+    let sizet_model = IntModel(SizeT);
+    let ndarray_model = StructModel(NpArray);
+    let ndarray_data_model = PtrModel(IntModel(Byte));
+
+    // Setup ndarray
+    let ndarray_ptr = ndarray_model.alloca(tyctx, ctx, name);
+    let shape = sizet_model.array_alloca(tyctx, ctx, ndims.value, "shape");
+    let strides = sizet_model.array_alloca(tyctx, ctx, ndims.value, "strides");
+
+    ndarray_ptr.gep(ctx, |f| f.data).store(ctx, ndarray_data_model.nullptr(tyctx, ctx.ctx));
+    ndarray_ptr.gep(ctx, |f| f.ndims).store(ctx, ndims);
+    ndarray_ptr.gep(ctx, |f| f.shape).store(ctx, shape);
+    ndarray_ptr.gep(ctx, |f| f.strides).store(ctx, strides);
+
+    Ok(ndarray_ptr)
+}
+
+/// Initialize an ndarray's `shape` and asserts on.
+/// `shape`'s values and prohibit illegal inputs like negative dimensions.
+pub fn init_ndarray_shape<'ctx, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    pndarray: Ptr<'ctx, StructModel<NpArray>>,
+    shape_writer: &ArrayWriter<'ctx, G, SizeT, IntModel<SizeT>>,
+) -> Result<(), String> {
+    let tyctx = generator.type_context(ctx.ctx);
+    let shape = pndarray.gep(ctx, |f| f.shape).load(tyctx, ctx, "shape");
+    (shape_writer.write)(generator, ctx, shape)?;
+    call_nac3_ndarray_util_assert_shape_no_negative(generator, ctx, shape_writer.len, shape);
+    Ok(())
+}
+
+/// Initialize an ndarray's `data` by allocating a buffer on the stack.
+/// The allocated data buffer is considered to be *owned* by the ndarray.
+///
+/// `strides` of the ndarray will also be updated with `set_strides_by_shape`.
+///
+/// `shape` and `itemsize` of the ndarray ***must*** be initialized first.
+pub fn init_ndarray_data_by_alloca<'ctx, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    pndarray: Ptr<'ctx, StructModel<NpArray>>,
+) {
+    let tyctx = generator.type_context(ctx.ctx);
+    let ndarray_data_model = IntModel(Byte);
+
+    let nbytes = call_nac3_ndarray_nbytes(generator, ctx, pndarray);
+    let data = ndarray_data_model.array_alloca(tyctx, ctx, nbytes.value, "data");
+    pndarray.gep(ctx, |f| f.data).store(ctx, data);
+    call_nac3_ndarray_set_strides_by_shape(generator, ctx, pndarray);
+}
+
+/// Convert `input` to an ndarray - behaves similarly to `np.asarray`.
+///
+/// Returns the ndarray interpretation of `input` and **the element type** of the ndarray.
+///
+/// Here are the exact details:
+/// - If `input` is an ndarray, the function returns back the **same** ndarray and the `dtype`
+/// of the ndarray.
+/// - If `input` is not an ndarray, the function creates an ndarray with a single element `input`,
+/// and returns the created ndarray and `input_ty`. Note that the created ndarray's `ndims` will
+/// be `0` (an *unsized* ndarray).
+pub fn as_ndarray<'ctx, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    input: BasicValueEnum<'ctx>,
+    input_ty: Type,
+) -> Result<(Ptr<'ctx, StructModel<NpArray>>, Type), String> {
+    let tyctx = generator.type_context(ctx.ctx);
+    let sizet_model = IntModel(SizeT);
+    let pbyte_model = PtrModel(IntModel(Byte));
+    let pndarray_model = PtrModel(StructModel(NpArray));
+
+    let input_ty_enum = ctx.unifier.get_ty(input_ty);
+    match &*input_ty_enum {
+        TypeEnum::TObj { obj_id, .. }
+            if *obj_id == ctx.primitives.ndarray.obj_id(&ctx.unifier).unwrap() =>
+        {
+            let pndarray = pndarray_model.check_value(tyctx, ctx.ctx, input).unwrap();
+            let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, input_ty);
+            Ok((pndarray, elem_ty))
+        }
+        _ => {
+            let ndims = sizet_model.const_0(tyctx, ctx.ctx);
+            let pndarray = alloca_ndarray(generator, ctx, ndims, "ndarray")?;
+
+            // We have to put `input` onto the stack to get a data pointer.
+            let data = ctx.builder.build_alloca(input.get_type(), "as_ndarray_scalar").unwrap();
+            ctx.builder.build_store(data, input).unwrap();
+
+            let data = pbyte_model.transmute(tyctx, ctx, data, "data");
+            pndarray.gep(ctx, |f| f.data).store(ctx, data);
+
+            let itemsize = input.get_type().size_of().unwrap();
+            let itemsize = sizet_model.check_value(tyctx, ctx.ctx, itemsize).unwrap();
+            pndarray.gep(ctx, |f| f.itemsize).store(ctx, itemsize);
+
+            Ok((pndarray, input_ty))
+        }
+    }
+}
+
+/// Iterate through all elements in an ndarray.
+///
+/// `body` is given the index of an element and an opaque pointer (as an `uint8_t*`, you might want to cast it) to the element.
+///
+/// Short-circuiting is possible with the given [`BreakContinueHooks`].
+pub fn gen_foreach_ndarray_elements<'ctx, G, F>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    pndarray: Ptr<'ctx, StructModel<NpArray>>,
+    body: F,
+) -> Result<(), String>
+where
+    G: CodeGenerator + ?Sized,
+    F: Fn(
+        &mut G,
+        &mut CodeGenContext<'ctx, '_>,
+        BreakContinueHooks<'ctx>,
+        Int<'ctx, SizeT>,
+        Ptr<'ctx, IntModel<Byte>>,
+    ) -> Result<(), String>,
+{
+    // TODO: Make this more efficient - use a special NDArray iterator?
+
+    let tyctx = generator.type_context(ctx.ctx);
+
+    let sizet_model = IntModel(SizeT);
+    let size = call_nac3_ndarray_size(generator, ctx, pndarray);
+
+    gen_model_for(
+        generator,
+        ctx,
+        sizet_model.const_0(tyctx, ctx.ctx),
+        size,
+        sizet_model.const_1(tyctx, ctx.ctx),
+        |generator, ctx, hooks, index| {
+            let pelement = call_nac3_ndarray_get_nth_pelement(generator, ctx, pndarray, index);
+            body(generator, ctx, hooks, index, pelement)
+        },
+    )
+}

nac3core/src/codegen/numpy_new/view.rs

@@ -4,7 +4,6 @@ use nac3parser::ast::StrRef;
 use crate::{
     codegen::{
         irrt::ndarray::{
-            allocation::{alloca_ndarray, init_ndarray_shape},
             basic::{
                 call_nac3_ndarray_is_c_contiguous, call_nac3_ndarray_nbytes,
                 call_nac3_ndarray_set_strides_by_shape, call_nac3_ndarray_size,
@@ -12,6 +11,7 @@ use crate::{
             reshape::call_nac3_ndarray_resolve_and_check_new_shape,
         },
         model::*,
+        numpy_new::util::{alloca_ndarray, init_ndarray_shape},
         structure::ndarray::NpArray,
         util::{array_writer::ArrayWriter, shape::make_shape_writer},
         CodeGenContext, CodeGenerator,