From 8033980554dced5bce5462194e9a3972c25af360 Mon Sep 17 00:00:00 2001
From: Jim Blandy <jimb@red-bean.com>
Date: Mon, 25 Mar 2024 19:11:23 -0700
Subject: [PATCH] [naga] Adjust RayQuery statements in override processing.

---
 naga/src/back/pipeline_constants.rs | 20 ++++++++++++++++++--
 naga/tests/snapshots.rs             |  5 +++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/naga/src/back/pipeline_constants.rs b/naga/src/back/pipeline_constants.rs
index d41eeedef27..c1fd2d02cc3 100644
--- a/naga/src/back/pipeline_constants.rs
+++ b/naga/src/back/pipeline_constants.rs
@@ -633,7 +633,7 @@ fn adjust_stmt(new_pos: &[Handle<Expression>], stmt: &mut Statement) {
         Statement::Call {
             ref mut arguments,
             ref mut result,
-            ..
+            function: _,
         } => {
             for argument in arguments.iter_mut() {
                 adjust(argument);
@@ -642,8 +642,24 @@ fn adjust_stmt(new_pos: &[Handle<Expression>], stmt: &mut Statement) {
                 adjust(e);
             }
         }
-        Statement::RayQuery { ref mut query, .. } => {
+        Statement::RayQuery {
+            ref mut query,
+            ref mut fun,
+        } => {
             adjust(query);
+            match *fun {
+                crate::RayQueryFunction::Initialize {
+                    ref mut acceleration_structure,
+                    ref mut descriptor,
+                } => {
+                    adjust(acceleration_structure);
+                    adjust(descriptor);
+                }
+                crate::RayQueryFunction::Proceed { ref mut result } => {
+                    adjust(result);
+                }
+                crate::RayQueryFunction::Terminate => {}
+            }
         }
         Statement::Break | Statement::Continue | Statement::Kill | Statement::Barrier(_) => {}
     }
diff --git a/naga/tests/snapshots.rs b/naga/tests/snapshots.rs
index 151e8b3da31..94c50c79759 100644
--- a/naga/tests/snapshots.rs
+++ b/naga/tests/snapshots.rs
@@ -466,6 +466,7 @@ fn write_output_spv(
             );
         }
     } else {
+        assert!(pipeline_constants.is_empty());
         write_output_spv_inner(input, module, info, &options, None, "spvasm");
     }
 }
@@ -857,6 +858,10 @@ fn convert_wgsl() {
             "overrides-atomicCompareExchangeWeak",
             Targets::IR | Targets::SPIRV,
         ),
+        (
+            "overrides-ray-query",
+            Targets::IR | Targets::SPIRV | Targets::METAL,
+        ),
     ];
 
     for &(name, targets) in inputs.iter() {