From d01f798b5a464cebea6a8b14d77c46987c9b0469 Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Tue, 7 May 2024 18:07:24 -0400
Subject: [PATCH] impl Aggregate for [T;N]

---
 luisa_compute/src/lang.rs | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/luisa_compute/src/lang.rs b/luisa_compute/src/lang.rs
index c8249b4..47c8a5b 100644
--- a/luisa_compute/src/lang.rs
+++ b/luisa_compute/src/lang.rs
@@ -154,6 +154,23 @@ pub trait Aggregate: Sized {
     fn from_nodes<I: Iterator<Item = SafeNodeRef>>(iter: &mut I) -> Self;
 }
 
+impl<const N: usize, T: Aggregate> Aggregate for [T; N] {
+    fn to_nodes(&self, nodes: &mut Vec<SafeNodeRef>) {
+        for x in self {
+            x.to_nodes(nodes);
+        }
+    }
+    fn from_nodes<I: Iterator<Item = SafeNodeRef>>(iter: &mut I) -> Self {
+        unsafe {
+            let mut ret = std::mem::MaybeUninit::<[T; N]>::uninit();
+            for i in 0..N {
+                let x = T::from_nodes(iter);
+                ret.as_mut_ptr().cast::<T>().add(i).write(x);
+            }
+            ret.assume_init()
+        }
+    }
+}
 impl<T: Aggregate> Aggregate for Vec<T> {
     fn to_nodes(&self, nodes: &mut Vec<SafeNodeRef>) {
         let len_node = __new_user_node(nodes.len());
@@ -427,7 +444,8 @@ impl FnRecorder {
             device: None,
             block_size: None,
             pools: pools.clone(),
-            arena: parent.as_ref()
+            arena: parent
+                .as_ref()
                 .map(|p| p.borrow().arena.clone())
                 .unwrap_or_else(|| Rc::new(Bump::new())),
             building_kernel: false,