fixed array indexed assignemnt

LuisaGroup · Feb 8, 2024 · 5ec6f57 · 5ec6f57
1 parent 695c992
commit 5ec6f57
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git 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
@@ -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);