From 3d0a1d281c215cab7ec0261fe20f054c175b024d Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Fri, 22 Nov 2024 16:38:57 +0800
Subject: [PATCH] [core] Expose irrt::ndarray

---
 nac3artiq/src/codegen.rs                   |  2 +-
 nac3core/src/codegen/builtin_fns.rs        |  3 ++-
 nac3core/src/codegen/expr.rs               |  2 +-
 nac3core/src/codegen/irrt/mod.rs           |  3 +--
 nac3core/src/codegen/numpy.rs              |  7 +++++--
 nac3core/src/codegen/values/ndarray/mod.rs | 11 ++++++++---
 6 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs
index 44d287d668..6340eb90e7 100644
--- a/nac3artiq/src/codegen.rs
+++ b/nac3artiq/src/codegen.rs
@@ -15,7 +15,7 @@ use pyo3::{
 use nac3core::{
     codegen::{
         expr::{destructure_range, gen_call},
-        irrt::call_ndarray_calc_size,
+        irrt::ndarray::call_ndarray_calc_size,
         llvm_intrinsics::{call_int_smax, call_memcpy_generic, call_stackrestore, call_stacksave},
         stmt::{gen_block, gen_for_callback_incrementing, gen_if_callback, gen_with},
         types::NDArrayType,
diff --git a/nac3core/src/codegen/builtin_fns.rs b/nac3core/src/codegen/builtin_fns.rs
index 3e6a37cd1b..f29a49eaa9 100644
--- a/nac3core/src/codegen/builtin_fns.rs
+++ b/nac3core/src/codegen/builtin_fns.rs
@@ -892,7 +892,8 @@ pub fn call_numpy_max_min<'ctx, G: CodeGenerator + ?Sized>(
             let llvm_elem_ty = ctx.get_llvm_type(generator, elem_ty);
 
             let n = NDArrayValue::from_pointer_value(n, llvm_elem_ty, llvm_usize, None);
-            let n_sz = irrt::call_ndarray_calc_size(generator, ctx, &n.shape(), (None, None));
+            let n_sz =
+                irrt::ndarray::call_ndarray_calc_size(generator, ctx, &n.shape(), (None, None));
             if ctx.registry.llvm_options.opt_level == OptimizationLevel::None {
                 let n_sz_eqz = ctx
                     .builder
diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 2cd56700ab..7e69e4d8c9 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -2842,7 +2842,7 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
                     llvm_i1.const_zero(),
                 );
 
-                let ndarray_num_elems = call_ndarray_calc_size(
+                let ndarray_num_elems = ndarray::call_ndarray_calc_size(
                     generator,
                     ctx,
                     &ndarray.shape().as_slice_value(ctx, generator),
diff --git a/nac3core/src/codegen/irrt/mod.rs b/nac3core/src/codegen/irrt/mod.rs
index f6c4a1eb9f..970dbf2703 100644
--- a/nac3core/src/codegen/irrt/mod.rs
+++ b/nac3core/src/codegen/irrt/mod.rs
@@ -13,12 +13,11 @@ use super::{CodeGenContext, CodeGenerator};
 use crate::{symbol_resolver::SymbolResolver, typecheck::typedef::Type};
 pub use list::*;
 pub use math::*;
-pub use ndarray::*;
 pub use slice::*;
 
 mod list;
 mod math;
-mod ndarray;
+pub mod ndarray;
 mod slice;
 
 #[must_use]
diff --git a/nac3core/src/codegen/numpy.rs b/nac3core/src/codegen/numpy.rs
index c3631f5644..ddaa515446 100644
--- a/nac3core/src/codegen/numpy.rs
+++ b/nac3core/src/codegen/numpy.rs
@@ -9,8 +9,11 @@ use nac3parser::ast::{Operator, StrRef};
 use super::{
     expr::gen_binop_expr_with_values,
     irrt::{
-        calculate_len_for_slice_range, call_ndarray_calc_broadcast,
-        call_ndarray_calc_broadcast_index, call_ndarray_calc_nd_indices, call_ndarray_calc_size,
+        calculate_len_for_slice_range,
+        ndarray::{
+            call_ndarray_calc_broadcast, call_ndarray_calc_broadcast_index,
+            call_ndarray_calc_nd_indices, call_ndarray_calc_size,
+        },
     },
     llvm_intrinsics::{self, call_memcpy_generic},
     macros::codegen_unreachable,
diff --git a/nac3core/src/codegen/values/ndarray/mod.rs b/nac3core/src/codegen/values/ndarray/mod.rs
index 6480511371..2583d45fd3 100644
--- a/nac3core/src/codegen/values/ndarray/mod.rs
+++ b/nac3core/src/codegen/values/ndarray/mod.rs
@@ -9,7 +9,7 @@ use super::{
     UntypedArrayLikeAccessor, UntypedArrayLikeMutator,
 };
 use crate::codegen::{
-    irrt::{call_ndarray_calc_size, call_ndarray_flatten_index},
+    irrt,
     llvm_intrinsics::call_int_umin,
     stmt::gen_for_callback_incrementing,
     types::{structure::StructField, NDArrayType},
@@ -291,7 +291,12 @@ impl<'ctx> ArrayLikeValue<'ctx> for NDArrayDataProxy<'ctx, '_> {
         ctx: &CodeGenContext<'ctx, '_>,
         generator: &G,
     ) -> IntValue<'ctx> {
-        call_ndarray_calc_size(generator, ctx, &self.as_slice_value(ctx, generator), (None, None))
+        irrt::ndarray::call_ndarray_calc_size(
+            generator,
+            ctx,
+            &self.as_slice_value(ctx, generator),
+            (None, None),
+        )
     }
 }
 
@@ -400,7 +405,7 @@ impl<'ctx, Index: UntypedArrayLikeAccessor<'ctx>> ArrayLikeIndexer<'ctx, Index>
             indices_elem_ty.get_bit_width()
         );
 
-        let index = call_ndarray_flatten_index(generator, ctx, *self.0, indices);
+        let index = irrt::ndarray::call_ndarray_flatten_index(generator, ctx, *self.0, indices);
         let sizeof_elem = ctx
             .builder
             .build_int_truncate_or_bit_cast(