From 1b0e6513aa17943a1e3331933699f2947a21bfc7 Mon Sep 17 00:00:00 2001
From: ScanMountGoat <scandanavianmountaingoat@gmail.com>
Date: Mon, 11 Nov 2024 16:58:06 -0600
Subject: [PATCH] expose push constant stages, closes #70

---
 example/src/main.rs     |  3 +--
 example/src/shader.rs   |  3 ++-
 wgsl_to_wgpu/src/lib.rs | 32 ++++++++++++++++++++++----------
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/example/src/main.rs b/example/src/main.rs
index 4afdb07..f944cf8 100644
--- a/example/src/main.rs
+++ b/example/src/main.rs
@@ -277,9 +277,8 @@ impl State {
                 color_matrix: glam::Mat4::IDENTITY,
             })
             .unwrap();
-        // TODO: Expose accessed stages for push constants in generated code.
         render_pass.set_push_constants(
-            wgpu::ShaderStages::FRAGMENT,
+            shader::PUSH_CONSTANT_STAGES,
             0,
             &push_constant_bytes.into_inner(),
         );
diff --git a/example/src/shader.rs b/example/src/shader.rs
index 67755ea..f573adf 100644
--- a/example/src/shader.rs
+++ b/example/src/shader.rs
@@ -259,6 +259,7 @@ pub fn create_shader_module(device: &wgpu::Device) -> wgpu::ShaderModule {
         source: wgpu::ShaderSource::Wgsl(source),
     })
 }
+pub const PUSH_CONSTANT_STAGES: wgpu::ShaderStages = wgpu::ShaderStages::FRAGMENT;
 pub fn create_pipeline_layout(device: &wgpu::Device) -> wgpu::PipelineLayout {
     device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
         label: None,
@@ -267,7 +268,7 @@ pub fn create_pipeline_layout(device: &wgpu::Device) -> wgpu::PipelineLayout {
             &bind_groups::BindGroup1::get_bind_group_layout(device),
         ],
         push_constant_ranges: &[wgpu::PushConstantRange {
-            stages: wgpu::ShaderStages::FRAGMENT,
+            stages: PUSH_CONSTANT_STAGES,
             range: 0..64,
         }],
     })
diff --git a/wgsl_to_wgpu/src/lib.rs b/wgsl_to_wgpu/src/lib.rs
index c69152a..2b2a6c9 100644
--- a/wgsl_to_wgpu/src/lib.rs
+++ b/wgsl_to_wgpu/src/lib.rs
@@ -248,7 +248,8 @@ fn create_shader_module_inner(
         })
         .collect();
 
-    let push_constant_range = push_constant_range(&module, &global_stages, entry_stages);
+    let (push_constant_range, push_constant_stages) =
+        push_constant_range_stages(&module, &global_stages, entry_stages).unzip();
 
     let create_pipeline_layout = quote! {
         pub fn create_pipeline_layout(device: &wgpu::Device) -> wgpu::PipelineLayout {
@@ -264,6 +265,12 @@ fn create_shader_module_inner(
 
     let override_constants = pipeline_overridable_constants(&module);
 
+    let push_constant_stages = push_constant_stages.map(|stages| {
+        quote! {
+            pub const PUSH_CONSTANT_STAGES: wgpu::ShaderStages = #stages;
+        }
+    });
+
     let output = quote! {
         #structs
         #(#consts)*
@@ -275,6 +282,7 @@ fn create_shader_module_inner(
         #vertex_states
         #fragment_states
         #create_shader_module
+        #push_constant_stages
         #create_pipeline_layout
     };
 
@@ -285,11 +293,11 @@ fn create_shader_module_inner(
     }
 }
 
-fn push_constant_range(
+fn push_constant_range_stages(
     module: &naga::Module,
     global_stages: &BTreeMap<String, wgpu::ShaderStages>,
     entry_stages: wgpu::ShaderStages,
-) -> Option<TokenStream> {
+) -> Option<(TokenStream, TokenStream)> {
     // Assume only one variable is used with var<push_constant> in WGSL.
     let (_, global) = module
         .global_variables
@@ -311,12 +319,15 @@ fn push_constant_range(
     // Use a single push constant range for all shader stages.
     // This allows easily setting push constants in a single call with offset 0.
     let size = Literal::usize_unsuffixed(push_constant_size as usize);
-    Some(quote! {
-        wgpu::PushConstantRange {
-            stages: #stages,
-            range: 0..#size
-        }
-    })
+    Some((
+        quote! {
+            wgpu::PushConstantRange {
+                stages: PUSH_CONSTANT_STAGES,
+                range: 0..#size
+            }
+        },
+        stages,
+    ))
 }
 
 fn pretty_print(output: TokenStream) -> String {
@@ -511,6 +522,7 @@ mod test {
                             source: wgpu::ShaderSource::Wgsl(source),
                         })
                 }
+                pub const PUSH_CONSTANT_STAGES: wgpu::ShaderStages = wgpu::ShaderStages::FRAGMENT;
                 pub fn create_pipeline_layout(device: &wgpu::Device) -> wgpu::PipelineLayout {
                     device
                         .create_pipeline_layout(
@@ -519,7 +531,7 @@ mod test {
                                 bind_group_layouts: &[],
                                 push_constant_ranges: &[
                                     wgpu::PushConstantRange {
-                                        stages: wgpu::ShaderStages::FRAGMENT,
+                                        stages: PUSH_CONSTANT_STAGES,
                                         range: 0..16,
                                     },
                                 ],