From 1fcb7d0c2e7086f8c3162b8ea5f9db3ff5f829c9 Mon Sep 17 00:00:00 2001
From: Matthias Deiml <matthias@deiml.net>
Date: Fri, 3 Jun 2022 20:50:50 +0000
Subject: [PATCH] Set naga capabilities corresponding to wgpu features (#4824)

# Objective

At the moment all extra capabilities are disabled when validating shaders with naga:
https://github.com/bevyengine/bevy/blob/c7c08f95cb784afc366eb2dcedd21d9d40e72d32/crates/bevy_render/src/render_resource/shader.rs#L146-L149
This means these features can't be used even if the corresponding wgpu features are active.

## Solution

With these changes capabilities are now set corresponding to `RenderDevice::features`.

---

I have validated these changes for push constants with a project I am currently working on. Though bevy does not support creating pipelines with push constants yet, so I was only able to see that shaders are validated and compiled as expected.
---
 .../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 dd0699992a5a7..475654e5a13d2 100644
--- a/crates/bevy_render/src/render_resource/shader.rs
+++ b/crates/bevy_render/src/render_resource/shader.rs
@@ -2,6 +2,7 @@ use bevy_asset::{AssetLoader, Handle, LoadContext, LoadedAsset};
 use bevy_reflect::{TypeUuid, Uuid};
 use bevy_utils::{tracing::error, BoxedFuture, HashMap};
 use naga::back::wgsl::WriterFlags;
+use naga::valid::Capabilities;
 use naga::{valid::ModuleInfo, Module};
 use once_cell::sync::Lazy;
 use regex::Regex;
@@ -9,6 +10,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 +127,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 +145,23 @@ impl ProcessedShader {
                 },
             )?,
         };
-        let module_info = naga::valid::Validator::new(
-            naga::valid::ValidationFlags::default(),
-            naga::valid::Capabilities::default(),
-        )
-        .validate(&module)?;
+        const CAPABILITIES: &[(Features, Capabilities)] = &[
+            (Features::PUSH_CONSTANTS, Capabilities::PUSH_CONSTANT),
+            (Features::SHADER_FLOAT64, Capabilities::FLOAT64),
+            (
+                Features::SHADER_PRIMITIVE_INDEX,
+                Capabilities::PRIMITIVE_INDEX,
+            ),
+        ];
+        let mut capabilities = Capabilities::empty();
+        for (feature, capability) in CAPABILITIES {
+            if features.contains(*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()?;