From 5ec6f579026665a1884e53fb92ffda6b6342bf8e Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Thu, 8 Feb 2024 12:32:48 -0500
Subject: [PATCH] fixed array indexed assignemnt

---
 luisa_compute/src/lang/types/array.rs | 16 ++++++++++++++++
 luisa_compute/tests/misc.rs           |  2 +-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/luisa_compute/src/lang/types/array.rs b/luisa_compute/src/lang/types/array.rs
index a9c8c36..7f9812f 100644
--- a/luisa_compute/src/lang/types/array.rs
+++ b/luisa_compute/src/lang/types/array.rs
@@ -93,6 +93,22 @@ impl<T: Value, const N: usize, X: IntoIndex> Index<X> for ArrayExpr<T, N> {
         ._ref()
     }
 }
+
+impl<T: Value, const N: usize, X: IntoIndex> Index<X> for ArrayVar<T, N> {
+    type Output = Var<T>;
+    fn index(&self, i: X) -> &Self::Output {
+        let i = i.to_u64();
+
+        // TODO: Add need_runtime_check()?
+        if need_runtime_check() {
+            check_index_lt_usize(i, N);
+        }
+        let i = i.node().get();
+        let self_node = self.0.node().get();
+        Var::<T>::from_node(__current_scope(|b| b.gep_chained(self_node, &[i], T::type_())).into())
+            ._ref()
+    }
+}
 impl<T: Value, const N: usize, X: IntoIndex> Index<X> for ArrayAtomicRef<T, N> {
     type Output = AtomicRef<T>;
     fn index(&self, i: X) -> &Self::Output {
diff --git a/luisa_compute/tests/misc.rs b/luisa_compute/tests/misc.rs
index 74b0d99..a0d36d2 100644
--- a/luisa_compute/tests/misc.rs
+++ b/luisa_compute/tests/misc.rs
@@ -922,7 +922,7 @@ fn array_read_write() {
             let arr = Var::<[i32; 4]>::zeroed();
             let i = i32::var_zeroed();
             while i < 4 {
-                arr.write(i.as_u32(), tid.as_i32() + i);
+                *arr[i.cast_u32()] = tid.as_i32() + i;
                 *i += 1;
             }
             buf_x.write(tid, arr);