From 0a51a26aba4d309bcc909b84b4b5b07922d9aec1 Mon Sep 17 00:00:00 2001
From: Jakob Hellermann <jakob.hellermann@protonmail.com>
Date: Thu, 24 Dec 2020 02:41:34 +0100
Subject: [PATCH] bevy_render: load .spv assets (#1104)

bevy_render: ShaderLoader can now load spv files
---
 .../src/pipeline/pipeline_compiler.rs         |  7 +-----
 crates/bevy_render/src/shader/shader.rs       | 24 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_render/src/pipeline/pipeline_compiler.rs b/crates/bevy_render/src/pipeline/pipeline_compiler.rs
index 1bb10a519a9b2..98c041160fa82 100644
--- a/crates/bevy_render/src/pipeline/pipeline_compiler.rs
+++ b/crates/bevy_render/src/pipeline/pipeline_compiler.rs
@@ -2,7 +2,7 @@ use super::{state_descriptors::PrimitiveTopology, IndexFormat, PipelineDescripto
 use crate::{
     pipeline::{BindType, InputStepMode, VertexBufferDescriptor},
     renderer::RenderResourceContext,
-    shader::{Shader, ShaderError, ShaderSource},
+    shader::{Shader, ShaderError},
 };
 use bevy_asset::{Assets, Handle};
 use bevy_reflect::Reflect;
@@ -79,11 +79,6 @@ impl PipelineCompiler {
 
         let shader = shaders.get(shader_handle).unwrap();
 
-        // don't produce new shader if the input source is already spirv
-        if let ShaderSource::Spirv(_) = shader.source {
-            return Ok(shader_handle.clone_weak());
-        }
-
         if let Some(specialized_shader) =
             specialized_shaders
                 .iter()
diff --git a/crates/bevy_render/src/shader/shader.rs b/crates/bevy_render/src/shader/shader.rs
index 19a89af632f0f..d1f2dfde57ad8 100644
--- a/crates/bevy_render/src/shader/shader.rs
+++ b/crates/bevy_render/src/shader/shader.rs
@@ -145,6 +145,24 @@ impl Shader {
         Shader { stage, source }
     }
 
+    #[cfg(not(target_arch = "wasm32"))]
+    pub fn from_spirv(spirv: &[u8]) -> Result<Shader, ShaderError> {
+        use spirv_reflect::{types::ReflectShaderStageFlags, ShaderModule};
+
+        let module = ShaderModule::load_u8_data(spirv)
+            .map_err(|msg| ShaderError::Compilation(msg.to_string()))?;
+        let stage = match module.get_shader_stage() {
+            ReflectShaderStageFlags::VERTEX => ShaderStage::Vertex,
+            ReflectShaderStageFlags::FRAGMENT => ShaderStage::Fragment,
+            other => panic!("cannot load {:?} shader", other),
+        };
+
+        Ok(Shader {
+            source: ShaderSource::spirv_from_bytes(spirv),
+            stage,
+        })
+    }
+
     pub fn from_glsl(stage: ShaderStage, glsl: &str) -> Shader {
         Shader {
             source: ShaderSource::Glsl(glsl.to_string()),
@@ -243,6 +261,10 @@ impl AssetLoader for ShaderLoader {
             let shader = match ext {
                 "vert" => Shader::from_glsl(ShaderStage::Vertex, std::str::from_utf8(bytes)?),
                 "frag" => Shader::from_glsl(ShaderStage::Fragment, std::str::from_utf8(bytes)?),
+                #[cfg(not(target_arch = "wasm32"))]
+                "spv" => Shader::from_spirv(bytes)?,
+                #[cfg(target_arch = "wasm32")]
+                "spv" => panic!("cannot load .spv file on wasm"),
                 _ => panic!("unhandled extension: {}", ext),
             };
 
@@ -252,7 +274,7 @@ impl AssetLoader for ShaderLoader {
     }
 
     fn extensions(&self) -> &[&str] {
-        &["vert", "frag"]
+        &["vert", "frag", "spv"]
     }
 }