From 20cc41a6394f515350cfdbc055c777e74b121940 Mon Sep 17 00:00:00 2001
From: Carter Anderson <mcanders1@gmail.com>
Date: Wed, 6 May 2020 19:26:55 -0700
Subject: [PATCH] create shader module from source

---
 .../headless_render_resource_context.rs       |  6 +++++
 .../src/renderer/render_resource_context.rs   |  5 ++++
 crates/bevy_render/src/shader/shader.rs       | 24 +++++++++++++++----
 .../renderer/wgpu_render_resource_context.rs  | 17 ++++++++-----
 4 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/crates/bevy_render/src/renderer/headless_render_resource_context.rs b/crates/bevy_render/src/renderer/headless_render_resource_context.rs
index c4f01ee6de..14616c7b6f 100644
--- a/crates/bevy_render/src/renderer/headless_render_resource_context.rs
+++ b/crates/bevy_render/src/renderer/headless_render_resource_context.rs
@@ -129,4 +129,10 @@ impl RenderResourceContext for HeadlessRenderResourceContext {
             None
         }
     }
+    fn create_shader_module_from_source(
+        &self,
+        _shader_handle: Handle<Shader>,
+        _shader: &Shader,
+    ) {
+    }
 }
diff --git a/crates/bevy_render/src/renderer/render_resource_context.rs b/crates/bevy_render/src/renderer/render_resource_context.rs
index 24b1571061..9cc7291bef 100644
--- a/crates/bevy_render/src/renderer/render_resource_context.rs
+++ b/crates/bevy_render/src/renderer/render_resource_context.rs
@@ -45,6 +45,11 @@ pub trait RenderResourceContext: Downcast + Send + Sync + 'static {
         shader_handle: Handle<Shader>,
         shader_storage: &AssetStorage<Shader>,
     );
+    fn create_shader_module_from_source(
+        &self,
+        shader_handle: Handle<Shader>,
+        shader: &Shader,
+    );
     fn remove_buffer(&self, resource: RenderResource);
     fn remove_texture(&self, resource: RenderResource);
     fn remove_sampler(&self, resource: RenderResource);
diff --git a/crates/bevy_render/src/shader/shader.rs b/crates/bevy_render/src/shader/shader.rs
index 5fee468133..5e7c324393 100644
--- a/crates/bevy_render/src/shader/shader.rs
+++ b/crates/bevy_render/src/shader/shader.rs
@@ -27,14 +27,18 @@ pub fn glsl_to_spirv(
     let mut output = compile(glsl_source, stage.into(), shader_defs).unwrap();
     let mut spv_bytes = Vec::new();
     output.read_to_end(&mut spv_bytes).unwrap();
+    bytes_to_words(&spv_bytes)
+}
 
-    let mut spv_words = Vec::new();
-    for bytes4 in spv_bytes.chunks(4) {
-        spv_words.push(u32::from_le_bytes([
+fn bytes_to_words(bytes: &[u8]) -> Vec<u32> {
+    let mut words = Vec::new();
+    for bytes4 in bytes.chunks(4) {
+        words.push(u32::from_le_bytes([
             bytes4[0], bytes4[1], bytes4[2], bytes4[3],
         ]));
     }
-    spv_words
+
+    words
 }
 
 #[derive(Clone, Debug, Hash, Eq, PartialEq)]
@@ -43,6 +47,12 @@ pub enum ShaderSource {
     Glsl(String),
 }
 
+impl ShaderSource {
+    pub fn spirv_from_bytes(bytes: &[u8]) -> ShaderSource {
+        ShaderSource::Spirv(bytes_to_words(bytes))
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct Shader {
     pub source: ShaderSource,
@@ -51,6 +61,12 @@ pub struct Shader {
 }
 
 impl Shader {
+    pub fn new(stage: ShaderStage, source: ShaderSource) -> Shader {
+        Shader {
+            stage,
+            source,
+        }
+    }
     pub fn from_glsl(stage: ShaderStage, glsl: &str) -> Shader {
         Shader {
             source: ShaderSource::Glsl(glsl.to_string()),
diff --git a/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs b/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs
index 9822514a1d..0eaa4b94fb 100644
--- a/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs
+++ b/crates/bevy_wgpu/src/renderer/wgpu_render_resource_context.rs
@@ -287,6 +287,16 @@ impl RenderResourceContext for WgpuRenderResourceContext {
         handle_info(info);
     }
 
+    fn create_shader_module_from_source(
+        &self,
+        shader_handle: Handle<Shader>,
+        shader: &Shader,
+    ) {
+        let mut shader_modules = self.resources.shader_modules.write().unwrap();
+        let shader_module = self.device.create_shader_module(&shader.get_spirv(None));
+        shader_modules.insert(shader_handle, shader_module);
+    }
+
     fn create_shader_module(
         &self,
         shader_handle: Handle<Shader>,
@@ -302,13 +312,8 @@ impl RenderResourceContext for WgpuRenderResourceContext {
         {
             return;
         }
-
-        let mut shader_modules = self.resources.shader_modules.write().unwrap();
-        // TODO: consider re-checking existence here
-
         let shader = shader_storage.get(&shader_handle).unwrap();
-        let shader_module = self.device.create_shader_module(&shader.get_spirv(None));
-        shader_modules.insert(shader_handle, shader_module);
+        self.create_shader_module_from_source(shader_handle, shader);
     }
 
     fn create_swap_chain(&self, window: &Window) {