From 70e1349d329b7a297cd02500291a0237964e5b9e Mon Sep 17 00:00:00 2001
From: lyken <lyken@m-labs.hk>
Date: Tue, 20 Aug 2024 15:17:10 +0800
Subject: [PATCH] core/ndstrides: implement np_identity() and np_eye()

---
 nac3core/src/codegen/numpy.rs                 | 41 ++++++++-------
 .../src/codegen/object/ndarray/factory.rs     | 52 ++++++++++++++++++-
 2 files changed, 73 insertions(+), 20 deletions(-)

diff --git a/nac3core/src/codegen/numpy.rs b/nac3core/src/codegen/numpy.rs
index a975672d..8a7f69e7 100644
--- a/nac3core/src/codegen/numpy.rs
+++ b/nac3core/src/codegen/numpy.rs
@@ -1905,15 +1905,23 @@ pub fn gen_ndarray_eye<'ctx>(
         ))
     }?;
 
-    call_ndarray_eye_impl(
-        generator,
-        context,
-        context.primitives.float,
-        nrows_arg.into_int_value(),
-        ncols_arg.into_int_value(),
-        offset_arg.into_int_value(),
-    )
-    .map(NDArrayValue::into)
+    let (dtype, _) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
+
+    let nrows = Int(Int32)
+        .check_value(generator, context.ctx, nrows_arg)
+        .unwrap()
+        .s_extend(generator, context, SizeT);
+    let ncols = Int(Int32)
+        .check_value(generator, context.ctx, ncols_arg)
+        .unwrap()
+        .s_extend(generator, context, SizeT);
+    let offset = Int(Int32)
+        .check_value(generator, context.ctx, offset_arg)
+        .unwrap()
+        .s_extend(generator, context, SizeT);
+
+    let ndarray = NDArrayObject::make_np_eye(generator, context, dtype, nrows, ncols, offset);
+    Ok(ndarray.instance.value)
 }
 
 /// Generates LLVM IR for `ndarray.identity`.
@@ -1927,20 +1935,15 @@ pub fn gen_ndarray_identity<'ctx>(
     assert!(obj.is_none());
     assert_eq!(args.len(), 1);
 
-    let llvm_usize = generator.get_size_type(context.ctx);
+    let (dtype, _) = unpack_ndarray_var_tys(&mut context.unifier, fun.0.ret);
 
     let n_ty = fun.0.args[0].ty;
     let n_arg = args[0].1.clone().to_basic_value_enum(context, generator, n_ty)?;
 
-    call_ndarray_eye_impl(
-        generator,
-        context,
-        context.primitives.float,
-        n_arg.into_int_value(),
-        n_arg.into_int_value(),
-        llvm_usize.const_zero(),
-    )
-    .map(NDArrayValue::into)
+    let n = Int(Int32).check_value(generator, context.ctx, n_arg).unwrap();
+    let n = n.s_extend(generator, context, SizeT);
+    let ndarray = NDArrayObject::make_np_identity(generator, context, dtype, n);
+    Ok(ndarray.instance.value)
 }
 
 /// Generates LLVM IR for `ndarray.copy`.
diff --git a/nac3core/src/codegen/object/ndarray/factory.rs b/nac3core/src/codegen/object/ndarray/factory.rs
index a1c0af78..04ee79ae 100644
--- a/nac3core/src/codegen/object/ndarray/factory.rs
+++ b/nac3core/src/codegen/object/ndarray/factory.rs
@@ -1,4 +1,4 @@
-use inkwell::values::BasicValueEnum;
+use inkwell::{values::BasicValueEnum, IntPredicate};
 
 use crate::{
     codegen::{
@@ -123,4 +123,54 @@ impl<'ctx> NDArrayObject<'ctx> {
         let fill_value = ndarray_one_value(generator, ctx, dtype);
         NDArrayObject::make_np_full(generator, ctx, dtype, ndims, shape, fill_value)
     }
+
+    /// Create an ndarray like `np.eye`.
+    pub fn make_np_eye<G: CodeGenerator + ?Sized>(
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        dtype: Type,
+        nrows: Instance<'ctx, Int<SizeT>>,
+        ncols: Instance<'ctx, Int<SizeT>>,
+        offset: Instance<'ctx, Int<SizeT>>,
+    ) -> Self {
+        let ndzero = ndarray_zero_value(generator, ctx, dtype);
+        let ndone = ndarray_one_value(generator, ctx, dtype);
+
+        let ndarray = NDArrayObject::alloca_dynamic_shape(generator, ctx, dtype, &[nrows, ncols]);
+
+        // Create data and make the matrix like look np.eye()
+        ndarray.create_data(generator, ctx);
+        ndarray
+            .foreach(generator, ctx, |generator, ctx, _hooks, nditer| {
+                // NOTE: rows and cols can never be zero here, since this ndarray's `np.size` would be zero
+                // and this loop would not execute.
+
+                // Load up `row_i` and `col_i` from indices.
+                let row_i = nditer.get_indices().get_index_const(generator, ctx, 0);
+                let col_i = nditer.get_indices().get_index_const(generator, ctx, 1);
+
+                let be_one = row_i.add(ctx, offset).compare(ctx, IntPredicate::EQ, col_i);
+                let value = ctx.builder.build_select(be_one.value, ndone, ndzero, "value").unwrap();
+
+                let p = nditer.get_pointer(generator, ctx);
+                ctx.builder.build_store(p, value).unwrap();
+
+                Ok(())
+            })
+            .unwrap();
+
+        ndarray
+    }
+
+    /// Create an ndarray like `np.identity`.
+    pub fn make_np_identity<G: CodeGenerator + ?Sized>(
+        generator: &mut G,
+        ctx: &mut CodeGenContext<'ctx, '_>,
+        dtype: Type,
+        size: Instance<'ctx, Int<SizeT>>,
+    ) -> Self {
+        // Convenient implementation
+        let offset = Int(SizeT).const_0(generator, ctx.ctx);
+        NDArrayObject::make_np_eye(generator, ctx, dtype, size, size, offset)
+    }
 }