Merge pull request #173 from StarArawn/add-shader-stage-visibility-to-bindings

Added BindingShaderStage in order to pass the correct shader stage to wgpu.

crates/bevy_render/src/pipeline/binding.rs

@@ -1,12 +1,20 @@
 use super::UniformProperty;
 use crate::texture::{TextureComponentType, TextureFormat, TextureViewDimension};
 
+bitflags::bitflags! {
+    pub struct BindingShaderStage: u32 {
+        const VERTEX = 1;
+        const FRAGMENT = 2;
+        const COMPUTE = 4;
+    }
+}
+
 #[derive(Hash, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
 pub struct BindingDescriptor {
     pub name: String,
     pub index: u32,
     pub bind_type: BindType,
-    // TODO: ADD SHADER STAGE VISIBILITY
+    pub shader_stage: BindingShaderStage,
 }
 
 #[derive(Hash, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]

crates/bevy_render/src/render_graph/nodes/pass_node.rs

@@ -4,6 +4,7 @@ use crate::{
     pass::{ClearColor, LoadOp, PassDescriptor, TextureAttachment},
     pipeline::{
         BindGroupDescriptor, BindType, BindingDescriptor, PipelineDescriptor, UniformProperty,
+        BindingShaderStage,
     },
     render_graph::{Node, ResourceSlotInfo, ResourceSlots},
     renderer::{
@@ -78,6 +79,7 @@ impl<Q: HecsQuery> PassNode<Q> {
                     dynamic: false,
                     properties: vec![UniformProperty::Struct(vec![UniformProperty::Mat4])],
                 },
+                shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
             }],
         );
 

crates/bevy_render/src/renderer/render_resource/render_resource_bindings.rs

@@ -261,7 +261,7 @@ impl Default for RenderResourceBindingsId {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::pipeline::{BindType, BindingDescriptor, UniformProperty};
+    use crate::pipeline::{BindType, BindingDescriptor, UniformProperty, BindingShaderStage};
 
     #[test]
     fn test_bind_groups() {
@@ -275,6 +275,7 @@ mod tests {
                         dynamic: false,
                         properties: vec![UniformProperty::Struct(vec![UniformProperty::Mat4])],
                     },
+                    shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
                 },
                 BindingDescriptor {
                     index: 1,
@@ -283,6 +284,7 @@ mod tests {
                         dynamic: false,
                         properties: vec![UniformProperty::Float],
                     },
+                    shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
                 },
             ],
         );

crates/bevy_render/src/shader/shader_reflect.rs

@@ -1,7 +1,7 @@
 use crate::{
     pipeline::{
         BindGroupDescriptor, BindType, BindingDescriptor, InputStepMode, UniformProperty,
-        VertexAttributeDescriptor, VertexBufferDescriptor, VertexFormat,
+        VertexAttributeDescriptor, VertexBufferDescriptor, VertexFormat, BindingShaderStage,
     },
     texture::{TextureComponentType, TextureViewDimension},
 };
@@ -204,6 +204,8 @@ fn reflect_binding(binding: &ReflectDescriptorBinding) -> BindingDescriptor {
         index: binding.binding,
         bind_type,
         name: name.to_string(),
+        // TODO: We should be able to detect which shader program the binding is being used in..
+        shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
     }
 }
 
@@ -414,6 +416,7 @@ mod tests {
                                     UniformProperty::Mat4
                                 ])],
                             },
+                            shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
                         }]
                     ),
                     BindGroupDescriptor::new(
@@ -426,6 +429,7 @@ mod tests {
                                 dimension: TextureViewDimension::D2,
                                 component_type: TextureComponentType::Float,
                             },
+                            shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
                         }]
                     ),
                 ]

crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs

@@ -5,7 +5,7 @@ use crate::{
 
 use bevy_asset::{Assets, Handle, HandleUntyped};
 use bevy_render::{
-    pipeline::{BindGroupDescriptor, BindGroupDescriptorId, PipelineDescriptor},
+    pipeline::{BindGroupDescriptor, BindGroupDescriptorId, BindingShaderStage, PipelineDescriptor},
     renderer::{
         BindGroup, BufferId, BufferInfo, RenderResourceBinding, RenderResourceContext,
         RenderResourceId, SamplerId, TextureId,
@@ -113,9 +113,18 @@ impl WgpuRenderResourceContext {
             .bindings
             .iter()
             .map(|binding| {
+                let shader_stage = if binding.shader_stage == BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT {
+                    wgpu::ShaderStage::VERTEX | wgpu::ShaderStage::FRAGMENT
+                } else if binding.shader_stage == BindingShaderStage::VERTEX {
+                    wgpu::ShaderStage::VERTEX
+                } else if binding.shader_stage == BindingShaderStage::FRAGMENT {
+                    wgpu::ShaderStage::FRAGMENT
+                } else {
+                    panic!("Invalid binding shader stage.")
+                };
                 wgpu::BindGroupLayoutEntry::new(
                     binding.index,
-                    wgpu::ShaderStage::VERTEX | wgpu::ShaderStage::FRAGMENT,
+                    shader_stage,
                     (&binding.bind_type).wgpu_into(),
                 )
             })