From cdff7e5c1c25e595859cc2e26480bcb65e601c90 Mon Sep 17 00:00:00 2001
From: MDeiml <matthias@deiml.net>
Date: Sun, 22 May 2022 14:57:28 +0200
Subject: [PATCH] Set naga capabilities

---
 .../src/render_resource/pipeline_cache.rs     |  4 ++-
 .../bevy_render/src/render_resource/shader.rs | 35 ++++++++++++++-----
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/crates/bevy_render/src/render_resource/pipeline_cache.rs b/crates/bevy_render/src/render_resource/pipeline_cache.rs
index 559ff745fa314..02971c966c26c 100644
--- a/crates/bevy_render/src/render_resource/pipeline_cache.rs
+++ b/crates/bevy_render/src/render_resource/pipeline_cache.rs
@@ -123,7 +123,9 @@ impl ShaderCache {
                     &self.shaders,
                     &self.import_path_shaders,
                 )?;
-                let module_descriptor = match processed.get_module_descriptor() {
+                let module_descriptor = match processed
+                    .get_module_descriptor(render_device.features())
+                {
                     Ok(module_descriptor) => module_descriptor,
                     Err(err) => {
                         return Err(PipelineCacheError::AsModuleDescriptorError(err, processed));
diff --git a/crates/bevy_render/src/render_resource/shader.rs b/crates/bevy_render/src/render_resource/shader.rs
index 38b8843258805..77ba771b314cc 100644
--- a/crates/bevy_render/src/render_resource/shader.rs
+++ b/crates/bevy_render/src/render_resource/shader.rs
@@ -9,6 +9,7 @@ use std::{
     borrow::Cow, collections::HashSet, marker::Copy, ops::Deref, path::PathBuf, str::FromStr,
 };
 use thiserror::Error;
+use wgpu::Features;
 use wgpu::{util::make_spirv, ShaderModuleDescriptor, ShaderSource};
 
 #[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
@@ -125,7 +126,7 @@ impl ProcessedShader {
         }
     }
 
-    pub fn reflect(&self) -> Result<ShaderReflection, ShaderReflectError> {
+    pub fn reflect(&self, features: Features) -> Result<ShaderReflection, ShaderReflectError> {
         let module = match &self {
             // TODO: process macros here
             ProcessedShader::Wgsl(source) => naga::front::wgsl::parse_str(source)?,
@@ -143,11 +144,24 @@ impl ProcessedShader {
                 },
             )?,
         };
-        let module_info = naga::valid::Validator::new(
-            naga::valid::ValidationFlags::default(),
-            naga::valid::Capabilities::default(),
-        )
-        .validate(&module)?;
+        let mut capabilities = naga::valid::Capabilities::empty();
+        // iterate over all capabilities
+        let all_bits: u8 = naga::valid::Capabilities::all().bits();
+        for i in 0..8 - all_bits.leading_zeros() {
+            let capability = naga::valid::Capabilities::from_bits(1 << i).unwrap();
+            let wgpu_feature = match capability {
+                naga::valid::Capabilities::PUSH_CONSTANT => Features::PUSH_CONSTANTS,
+                naga::valid::Capabilities::FLOAT64 => Features::SHADER_FLOAT64,
+                naga::valid::Capabilities::PRIMITIVE_INDEX => Features::SHADER_PRIMITIVE_INDEX,
+                _ => unimplemented!(),
+            };
+            if features.contains(wgpu_feature) {
+                capabilities |= capability;
+            }
+        }
+        let module_info =
+            naga::valid::Validator::new(naga::valid::ValidationFlags::default(), capabilities)
+                .validate(&module)?;
 
         Ok(ShaderReflection {
             module,
@@ -155,7 +169,10 @@ impl ProcessedShader {
         })
     }
 
-    pub fn get_module_descriptor(&self) -> Result<ShaderModuleDescriptor, AsModuleDescriptorError> {
+    pub fn get_module_descriptor(
+        &self,
+        features: Features,
+    ) -> Result<ShaderModuleDescriptor, AsModuleDescriptorError> {
         Ok(ShaderModuleDescriptor {
             label: None,
             source: match self {
@@ -164,12 +181,12 @@ impl ProcessedShader {
                     // Parse and validate the shader early, so that (e.g. while hot reloading) we can
                     // display nicely formatted error messages instead of relying on just displaying the error string
                     // returned by wgpu upon creating the shader module.
-                    let _ = self.reflect()?;
+                    let _ = self.reflect(features)?;
 
                     ShaderSource::Wgsl(source.clone())
                 }
                 ProcessedShader::Glsl(_source, _stage) => {
-                    let reflection = self.reflect()?;
+                    let reflection = self.reflect(features)?;
                     // TODO: it probably makes more sense to convert this to spirv, but as of writing
                     // this comment, naga's spirv conversion is broken
                     let wgsl = reflection.get_wgsl()?;