From a481add9aff2372e87afd004870614e03e7fe6fc 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.rs | 14 ++++++++------
 6 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs
index 99998791..b076a0f6 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, ProxyType},
diff --git a/nac3core/src/codegen/builtin_fns.rs b/nac3core/src/codegen/builtin_fns.rs
index 20d35006..f8356f30 100644
--- a/nac3core/src/codegen/builtin_fns.rs
+++ b/nac3core/src/codegen/builtin_fns.rs
@@ -922,7 +922,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 f8462ab2..47f04037 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 f6c4a1eb..970dbf27 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 58869f2c..0029f69b 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.rs b/nac3core/src/codegen/values/ndarray.rs
index 38c6a980..3fcd79d5 100644
--- a/nac3core/src/codegen/values/ndarray.rs
+++ b/nac3core/src/codegen/values/ndarray.rs
@@ -9,10 +9,7 @@ use super::{
     UntypedArrayLikeAccessor, UntypedArrayLikeMutator,
 };
 use crate::codegen::{
-    irrt::{call_ndarray_calc_size, call_ndarray_flatten_index},
-    llvm_intrinsics::call_int_umin,
-    stmt::gen_for_callback_incrementing,
-    types::NDArrayType,
+    irrt, llvm_intrinsics::call_int_umin, stmt::gen_for_callback_incrementing, types::NDArrayType,
     CodeGenContext, CodeGenerator,
 };
 
@@ -296,7 +293,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),
+        )
     }
 }
 
@@ -405,7 +407,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(