From a176c3eb705ebe2e127c10a0612cb5ffe25782fb Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Wed, 29 May 2024 15:07:28 +0800
Subject: [PATCH] core/irrt: Change handle_slice_indices to instead take length
 of object

So that all other array-like datatypes (e.g. ndarray) can also take
advantage of it.
---
 nac3core/src/codegen/expr.rs     | 24 ++++++++++++++++--------
 nac3core/src/codegen/irrt/mod.rs |  3 +--
 nac3core/src/codegen/stmt.rs     | 23 ++++++++++++++++-------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 6d8b4ac1..72e32405 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -2264,10 +2264,14 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
                     let ty = ctx.get_llvm_type(generator, *ty);
                     if let ExprKind::Slice { lower, upper, step } = &slice.node {
                         let one = int32.const_int(1, false);
-                        let Some((start, end, step)) =
-                            handle_slice_indices(lower, upper, step, ctx, generator, v)? else {
-                            return Ok(None)
-                        };
+                        let Some((start, end, step)) = handle_slice_indices(
+                                lower,
+                                upper,
+                                step,
+                                ctx,
+                                generator,
+                                v.load_size(ctx, None),
+                        )? else { return Ok(None) };
                         let length = calculate_len_for_slice_range(
                             generator,
                             ctx,
@@ -2289,10 +2293,14 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
                             step,
                         );
                         let res_array_ret = allocate_list(generator, ctx, ty, length, Some("ret"));
-                        let Some(res_ind) =
-                            handle_slice_indices(&None, &None, &None, ctx, generator, res_array_ret)? else {
-                            return Ok(None)
-                        };
+                        let Some(res_ind) = handle_slice_indices(
+                            &None,
+                            &None,
+                            &None,
+                            ctx,
+                            generator,
+                            res_array_ret.load_size(ctx, None),
+                        )? else { return Ok(None) };
                         list_slice_assignment(
                             generator,
                             ctx,
diff --git a/nac3core/src/codegen/irrt/mod.rs b/nac3core/src/codegen/irrt/mod.rs
index fbf0edc5..500a3e6d 100644
--- a/nac3core/src/codegen/irrt/mod.rs
+++ b/nac3core/src/codegen/irrt/mod.rs
@@ -175,12 +175,11 @@ pub fn handle_slice_indices<'ctx, G: CodeGenerator>(
     step: &Option<Box<Expr<Option<Type>>>>,
     ctx: &mut CodeGenContext<'ctx, '_>,
     generator: &mut G,
-    list: ListValue<'ctx>,
+    length: IntValue<'ctx>,
 ) -> Result<Option<(IntValue<'ctx>, IntValue<'ctx>, IntValue<'ctx>)>, String> {
     let int32 = ctx.ctx.i32_type();
     let zero = int32.const_zero();
     let one = int32.const_int(1, false);
-    let length = list.load_size(ctx, Some("length"));
     let length = ctx.builder.build_int_truncate_or_bit_cast(length, int32, "leni32").unwrap();
     Ok(Some(match (start, end, step) {
         (s, e, None) => (
diff --git a/nac3core/src/codegen/stmt.rs b/nac3core/src/codegen/stmt.rs
index e100f310..2fe22176 100644
--- a/nac3core/src/codegen/stmt.rs
+++ b/nac3core/src/codegen/stmt.rs
@@ -240,10 +240,14 @@ pub fn gen_assign<'ctx, G: CodeGenerator>(
                 .to_basic_value_enum(ctx, generator, ls.custom.unwrap())?
                 .into_pointer_value();
             let ls = ListValue::from_ptr_val(ls, llvm_usize, None);
-            let Some((start, end, step)) =
-                handle_slice_indices(lower, upper, step, ctx, generator, ls)? else {
-                return Ok(())
-            };
+            let Some((start, end, step)) = handle_slice_indices(
+                lower,
+                upper,
+                step,
+                ctx,
+                generator,
+                ls.load_size(ctx, None),
+            )? else { return Ok(()) };
             let value = value
                 .to_basic_value_enum(ctx, generator, target.custom.unwrap())?
                 .into_pointer_value();
@@ -257,9 +261,14 @@ pub fn gen_assign<'ctx, G: CodeGenerator>(
             };
 
             let ty = ctx.get_llvm_type(generator, ty);
-            let Some(src_ind) = handle_slice_indices(&None, &None, &None, ctx, generator, value)? else {
-                return Ok(())
-            };
+            let Some(src_ind) = handle_slice_indices(
+                &None,
+                &None,
+                &None,
+                ctx,
+                generator,
+                value.load_size(ctx, None),
+            )? else { return Ok(()) };
             list_slice_assignment(generator, ctx, ty, ls, (start, end, step), value, src_ind);
         }
         _ => {