fix: capturing deep composites

crates/rustc_codegen_spirv/src/linker/destructure_composites.rs

@@ -14,29 +14,32 @@ pub fn destructure_composites(function: &mut Function) {
         .all_inst_iter()
         .filter_map(|inst| match inst.class.opcode {
             Op::CompositeConstruct => Some((inst.result_id.unwrap(), inst.clone())),
-            Op::CompositeInsert if inst.operands.len() == 3 => {
-                Some((inst.result_id.unwrap(), inst.clone()))
-            }
+            Op::CompositeInsert => Some((inst.result_id.unwrap(), inst.clone())),
             _ => None,
         })
         .collect();
     for inst in function.all_inst_iter_mut() {
-        if inst.class.opcode == Op::CompositeExtract && inst.operands.len() == 2 {
+        if inst.class.opcode == Op::CompositeExtract {
             let mut composite = inst.operands[0].unwrap_id_ref();
-            let index = inst.operands[1].unwrap_literal_int32();
+            let indices = &inst.operands[1..];
 
             let origin = loop {
                 if let Some(inst) = reference.get(&composite) {
                     match inst.class.opcode {
                         Op::CompositeInsert => {
-                            let insert_index = inst.operands[2].unwrap_literal_int32();
-                            if insert_index == index {
+                            let insert_indices = &inst.operands[2..];
+                            if insert_indices == indices {
                                 break Some(inst.operands[0].unwrap_id_ref());
                             }
                             composite = inst.operands[1].unwrap_id_ref();
                         }
                         Op::CompositeConstruct => {
-                            break inst.operands.get(index as usize).map(|o| o.unwrap_id_ref());
+                            if indices.len() == 1 {
+                                break inst
+                                    .operands
+                                    .get(indices[0].unwrap_literal_int32() as usize)
+                                    .map(|o| o.unwrap_id_ref());
+                            }
                         }
                         _ => unreachable!(),
                     }

tests/ui/lang/control_flow/closure_capture_invalid_01.rs

@@ -0,0 +1,32 @@
+// build-fail
+// compile-flags: -Ctarget-feature=+RayTracingKHR,+RayQueryKHR,+ext:SPV_KHR_ray_tracing,+ext:SPV_KHR_ray_query
+
+#[cfg(not(target_arch = "spirv"))]
+use spirv_std::macros::spirv;
+use spirv_std::ray_tracing::{AccelerationStructure, RayFlags, RayQuery};
+use spirv_std::{glam::Vec3, ray_query};
+
+#[inline]
+pub fn some_fun(mut c: impl FnMut() -> bool) {
+    c();
+}
+#[spirv(intersection)]
+pub fn voxels_gbuffer_is(
+    #[spirv(descriptor_set = 1, binding = 3)] acceleration_structure: &AccelerationStructure,
+) {
+    ray_query!(let mut handle);
+    some_fun(|| {
+        unsafe {
+            handle.initialize(
+                acceleration_structure,
+                RayFlags::TERMINATE_ON_FIRST_HIT,
+                0xFF,
+                Vec3::new(0.0, 0.0, 0.0),
+                0.0,
+                Vec3::new(0.0, 0.0, 0.0),
+                10000.0,
+            );
+            true
+        }
+    });
+}

tests/ui/lang/control_flow/closure_capture_invalid_01.stderr

@@ -0,0 +1,8 @@
+error: error:0:0 - In Logical addressing, variables may not allocate a pointer type
+         %29 = OpVariable %_ptr_Function__ptr_UniformConstant_13 Function
+  |
+  = note: spirv-val failed
+  = note: module `$TEST_BUILD_DIR/lang/control_flow/closure_capture_invalid_01.stage-id.spv.dir/module`
+
+error: aborting due to previous error
+

tests/ui/lang/control_flow/closure_capture_invalid_02.rs

@@ -0,0 +1,43 @@
+// build-fail
+// compile-flags: -Ctarget-feature=+RayTracingKHR,+RayQueryKHR,+ext:SPV_KHR_ray_tracing,+ext:SPV_KHR_ray_query
+
+#[cfg(not(target_arch = "spirv"))]
+use spirv_std::macros::spirv;
+use spirv_std::{glam::Vec3, ray_query};
+use spirv_std::ray_tracing::{AccelerationStructure, CommittedIntersection, RayFlags, RayQuery};
+use spirv_std::arch::IndexUnchecked;
+
+pub fn some_fun(mut c: impl FnMut(u32) -> bool) {
+    let mut i = 0;
+    while i < 10 {
+        if c(i) {
+            return;
+        }
+        i += 1;
+    }
+}
+#[spirv(intersection)]
+pub fn voxels_gbuffer_is(
+    #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] vertices: &[u32],
+    #[spirv(descriptor_set = 1, binding = 3)] acceleration_structure: &AccelerationStructure,
+) {
+    some_fun(|i| {
+        unsafe {
+            ray_query!(let mut handle);
+            handle.initialize(
+                acceleration_structure,
+                RayFlags::TERMINATE_ON_FIRST_HIT,
+                0xFF,
+                Vec3::new(0.0, 0.0, 0.0),
+                0.0,
+                Vec3::new(0.0, 0.0, 0.0),
+                10000.0,
+            );
+            let ind = unsafe {
+                vertices.index_unchecked(0)
+            };
+            ind + i == 2
+        }
+    });
+}
+

tests/ui/lang/control_flow/closure_capture_invalid_02.stderr

@@ -0,0 +1,8 @@
+error: error:0:0 - In Logical addressing, variables may not allocate a pointer type
+         %45 = OpVariable %_ptr_Function__ptr_UniformConstant_29 Function
+  |
+  = note: spirv-val failed
+  = note: module `$TEST_BUILD_DIR/lang/control_flow/closure_capture_invalid_02.stage-id.spv.dir/module`
+
+error: aborting due to previous error
+