Reflect shader stage for bind groups.

crates/bevy_render/src/pipeline/pipeline_layout.rs

@@ -16,6 +16,7 @@ impl PipelineLayout {
     }
 
     pub fn from_shader_layouts(shader_layouts: &mut [ShaderLayout]) -> Self {
+        dbg!(&shader_layouts);
         let mut bind_groups = HashMap::<u32, BindGroupDescriptor>::new();
         let mut vertex_buffer_descriptors = Vec::new();
         for shader_layout in shader_layouts.iter_mut() {
@@ -25,10 +26,15 @@ impl PipelineLayout {
                         for shader_binding in shader_bind_group.bindings.iter() {
                             if let Some(binding) = bind_group
                                 .bindings
-                                .iter()
+                                .iter_mut()
                                 .find(|binding| binding.index == shader_binding.index)
                             {
-                                if binding != shader_binding {
+                                binding.shader_stage |= shader_binding.shader_stage;
+                                // Not sure we need to panic anymore here..
+                                if binding.bind_type != shader_binding.bind_type
+                                    || binding.name != shader_binding.name
+                                    || binding.index != shader_binding.index
+                                {
                                     panic!("Binding {} in BindGroup {} does not match across all shader types: {:?} {:?}", binding.index, bind_group.index, binding, shader_binding);
                                 }
                             } else {
@@ -42,6 +48,7 @@ impl PipelineLayout {
                 }
             }
         }
+        dbg!(&bind_groups);
 
         for vertex_buffer_descriptor in shader_layouts[0].vertex_buffer_descriptors.iter() {
             vertex_buffer_descriptors.push(vertex_buffer_descriptor.clone());
@@ -56,6 +63,8 @@ impl PipelineLayout {
         // TODO: try removing this
         bind_groups_result.sort_by(|a, b| a.index.partial_cmp(&b.index).unwrap());
 
+        dbg!(&bind_groups_result);
+
         PipelineLayout {
             bind_groups: bind_groups_result,
             vertex_buffer_descriptors,

crates/bevy_render/src/shader/shader_reflect.rs

@@ -9,7 +9,7 @@ use bevy_core::AsBytes;
 use spirv_reflect::{
     types::{
         ReflectDescriptorBinding, ReflectDescriptorSet, ReflectDescriptorType, ReflectDimension,
-        ReflectInterfaceVariable, ReflectTypeDescription, ReflectTypeFlags,
+        ReflectInterfaceVariable, ReflectTypeDescription, ReflectTypeFlags, ReflectShaderStageFlags
     },
     ShaderModule,
 };
@@ -30,9 +30,10 @@ impl ShaderLayout {
         match ShaderModule::load_u8_data(spirv_data.as_bytes()) {
             Ok(ref mut module) => {
                 let entry_point_name = module.get_entry_point_name();
+                let shader_stage = module.get_shader_stage();
                 let mut bind_groups = Vec::new();
                 for descriptor_set in module.enumerate_descriptor_sets(None).unwrap() {
-                    let bind_group = reflect_bind_group(&descriptor_set);
+                    let bind_group = reflect_bind_group(&descriptor_set, shader_stage);
                     bind_groups.push(bind_group);
                 }
 
@@ -150,10 +151,10 @@ fn reflect_vertex_attribute_descriptor(
     }
 }
 
-fn reflect_bind_group(descriptor_set: &ReflectDescriptorSet) -> BindGroupDescriptor {
+fn reflect_bind_group(descriptor_set: &ReflectDescriptorSet, shader_stage: ReflectShaderStageFlags) -> BindGroupDescriptor {
     let mut bindings = Vec::new();
     for descriptor_binding in descriptor_set.bindings.iter() {
-        let binding = reflect_binding(descriptor_binding);
+        let binding = reflect_binding(descriptor_binding, shader_stage);
         bindings.push(binding);
     }
 
@@ -170,7 +171,7 @@ fn reflect_dimension(type_description: &ReflectTypeDescription) -> TextureViewDi
     }
 }
 
-fn reflect_binding(binding: &ReflectDescriptorBinding) -> BindingDescriptor {
+fn reflect_binding(binding: &ReflectDescriptorBinding, shader_stage: ReflectShaderStageFlags) -> BindingDescriptor {
     let type_description = binding.type_description.as_ref().unwrap();
     let (name, bind_type) = match binding.descriptor_type {
         ReflectDescriptorType::UniformBuffer => (
@@ -200,12 +201,18 @@ fn reflect_binding(binding: &ReflectDescriptorBinding) -> BindingDescriptor {
         _ => panic!("unsupported bind type {:?}", binding.descriptor_type),
     };
 
+    let shader_stage = match shader_stage {
+        ReflectShaderStageFlags::COMPUTE => BindingShaderStage::COMPUTE,
+        ReflectShaderStageFlags::VERTEX => BindingShaderStage::VERTEX,
+        ReflectShaderStageFlags::FRAGMENT => BindingShaderStage::FRAGMENT,
+        _ => panic!("Only one specified shader stage is support.")
+    };
+
     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,
+        shader_stage,
     }
 }
 
@@ -416,7 +423,7 @@ mod tests {
                                     UniformProperty::Mat4
                                 ])],
                             },
-                            shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
+                            shader_stage: BindingShaderStage::VERTEX,
                         }]
                     ),
                     BindGroupDescriptor::new(
@@ -429,7 +436,7 @@ mod tests {
                                 dimension: TextureViewDimension::D2,
                                 component_type: TextureComponentType::Float,
                             },
-                            shader_stage: BindingShaderStage::VERTEX | BindingShaderStage::FRAGMENT,
+                            shader_stage: BindingShaderStage::VERTEX,
                         }]
                     ),
                 ]

crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs

@@ -133,6 +133,7 @@ impl WgpuRenderResourceContext {
             bindings: bind_group_layout_binding.as_slice(),
             label: None,
         };
+        dbg!(&wgpu_descriptor);
         let bind_group_layout = self.device.create_bind_group_layout(&wgpu_descriptor);
         bind_group_layouts.insert(descriptor.id, bind_group_layout);
     }