From c35ad0694902c70db2189c4fa8185bd5173292d0 Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Thu, 30 May 2024 14:26:22 +0800
Subject: [PATCH] core/expr: Add support for 1D slicing of NDArrays

---
 nac3core/src/codegen/expr.rs | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 72e32405..674d7840 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -1712,8 +1712,27 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
         slice.location,
     );
 
-    if let ExprKind::Slice { .. } = &slice.node {
-        return Err(String::from("subscript operator for ndarray not implemented"))
+    if let ExprKind::Slice { lower, upper, step } = &slice.node {
+        let dim0_sz = unsafe {
+            v.dim_sizes().get_typed_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
+        };
+
+        let Some((start, stop, step)) = handle_slice_indices(
+            lower,
+            upper,
+            step,
+            ctx,
+            generator,
+            dim0_sz,
+        )? else { return Ok(None) };
+
+        return Ok(Some(numpy::ndarray_sliced_copy(
+            generator,
+            ctx,
+            ty,
+            v,
+            &[(start, stop, step)],
+        )?.as_ptr_value().into()))
     }
 
     let index = if let Some(index) = generator.gen_expr(ctx, slice)? {