wasm32: non-spirv shader specialization (#843)

wasm32: non-spirv shader specialization

crates/bevy_render/src/pipeline/pipeline_compiler.rs

@@ -70,6 +70,7 @@ pub struct PipelineCompiler {
 impl PipelineCompiler {
     fn compile_shader(
         &mut self,
+        render_resource_context: &dyn RenderResourceContext,
         shaders: &mut Assets<Shader>,
         shader_handle: &Handle<Shader>,
         shader_specialization: &ShaderSpecialization,
@@ -102,7 +103,8 @@ impl PipelineCompiler {
                 .iter()
                 .cloned()
                 .collect::<Vec<String>>();
-            let compiled_shader = shader.get_spirv_shader(Some(&shader_def_vec));
+            let compiled_shader =
+                render_resource_context.get_specialized_shader(shader, Some(&shader_def_vec));
             let specialized_handle = shaders.add(compiled_shader);
             let weak_specialized_handle = specialized_handle.clone_weak();
             specialized_shaders.push(SpecializedShader {
@@ -141,6 +143,7 @@ impl PipelineCompiler {
         let source_descriptor = pipelines.get(source_pipeline).unwrap();
         let mut specialized_descriptor = source_descriptor.clone();
         specialized_descriptor.shader_stages.vertex = self.compile_shader(
+            render_resource_context,
             shaders,
             &specialized_descriptor.shader_stages.vertex,
             &pipeline_specialization.shader_specialization,
@@ -151,6 +154,7 @@ impl PipelineCompiler {
             .as_ref()
             .map(|fragment| {
                 self.compile_shader(
+                    render_resource_context,
                     shaders,
                     fragment,
                     &pipeline_specialization.shader_specialization,

crates/bevy_render/src/renderer/headless_render_resource_context.rs

@@ -148,4 +148,8 @@ impl RenderResourceContext for HeadlessRenderResourceContext {
     fn get_aligned_texture_size(&self, size: usize) -> usize {
         size
     }
+
+    fn get_specialized_shader(&self, shader: &Shader, _macros: Option<&[String]>) -> Shader {
+        shader.clone()
+    }
 }

crates/bevy_render/src/renderer/render_resource_context.rs

@@ -29,6 +29,7 @@ pub trait RenderResourceContext: Downcast + Send + Sync + 'static {
     fn create_buffer_with_data(&self, buffer_info: BufferInfo, data: &[u8]) -> BufferId;
     fn create_shader_module(&self, shader_handle: &Handle<Shader>, shaders: &Assets<Shader>);
     fn create_shader_module_from_source(&self, shader_handle: &Handle<Shader>, shader: &Shader);
+    fn get_specialized_shader(&self, shader: &Shader, macros: Option<&[String]>) -> Shader;
     fn remove_buffer(&self, buffer: BufferId);
     fn remove_texture(&self, texture: TextureId);
     fn remove_sampler(&self, sampler: SamplerId);

crates/bevy_render/src/shader/mod.rs

@@ -5,12 +5,10 @@ mod shader_defs;
 #[cfg(not(target_arch = "wasm32"))]
 mod shader_reflect;
 
-#[cfg(target_arch = "wasm32")]
-#[path = "shader_reflect_wasm.rs"]
-mod shader_reflect;
-
 pub use shader::*;
 pub use shader_defs::*;
+
+#[cfg(not(target_arch = "wasm32"))]
 pub use shader_reflect::*;
 
 use crate::pipeline::{BindGroupDescriptor, VertexBufferDescriptor};

crates/bevy_render/src/shader/shader.rs

@@ -23,7 +23,7 @@ impl Into<bevy_glsl_to_spirv::ShaderType> for ShaderStage {
 }
 
 #[cfg(all(not(target_os = "ios"), not(target_arch = "wasm32")))]
-fn glsl_to_spirv(
+pub fn glsl_to_spirv(
     glsl_source: &str,
     stage: ShaderStage,
     shader_defs: Option<&[String]>,
@@ -43,7 +43,7 @@ impl Into<shaderc::ShaderKind> for ShaderStage {
 }
 
 #[cfg(target_os = "ios")]
-fn glsl_to_spirv(
+pub fn glsl_to_spirv(
     glsl_source: &str,
     stage: ShaderStage,
     shader_defs: Option<&[String]>,
@@ -121,17 +121,15 @@ impl Shader {
         }
     }
 
-    #[allow(unused_variables)]
+    #[cfg(not(target_arch = "wasm32"))]
     pub fn get_spirv_shader(&self, macros: Option<&[String]>) -> Shader {
         Shader {
-            #[cfg(not(target_arch = "wasm32"))]
             source: ShaderSource::Spirv(self.get_spirv(macros)),
-            #[cfg(target_arch = "wasm32")]
-            source: self.source.clone(),
             stage: self.stage,
         }
     }
 
+    #[cfg(not(target_arch = "wasm32"))]
     pub fn reflect_layout(&self, enforce_bevy_conventions: bool) -> Option<ShaderLayout> {
         if let ShaderSource::Spirv(ref spirv) = self.source {
             Some(ShaderLayout::from_spirv(
@@ -142,6 +140,11 @@ impl Shader {
             panic!("Cannot reflect layout of non-SpirV shader. Try compiling this shader to SpirV first using self.get_spirv_shader()");
         }
     }
+
+    #[cfg(target_arch = "wasm32")]
+    pub fn reflect_layout(&self, _enforce_bevy_conventions: bool) -> Option<ShaderLayout> {
+        panic!("Cannot reflect layout on wasm32");
+    }
 }
 
 /// All stages in a shader program

crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs

@@ -12,7 +12,7 @@ use bevy_render::{
         BindGroup, BufferId, BufferInfo, RenderResourceBinding, RenderResourceContext,
         RenderResourceId, SamplerId, TextureId,
     },
-    shader::Shader,
+    shader::{glsl_to_spirv, Shader, ShaderSource},
     texture::{Extent3d, SamplerDescriptor, TextureDescriptor},
 };
 use bevy_utils::tracing::trace;
@@ -569,4 +569,15 @@ impl RenderResourceContext for WgpuRenderResourceContext {
             size
         }
     }
+
+    fn get_specialized_shader(&self, shader: &Shader, macros: Option<&[String]>) -> Shader {
+        let spirv_data = match shader.source {
+            ShaderSource::Spirv(ref bytes) => bytes.clone(),
+            ShaderSource::Glsl(ref source) => glsl_to_spirv(&source, shader.stage, macros),
+        };
+        Shader {
+            source: ShaderSource::Spirv(spirv_data),
+            ..*shader
+        }
+    }
 }