From fa6acdc2e879ed64c15903bd8588a007afe49c5a Mon Sep 17 00:00:00 2001
From: rttad <rttad@proton.me>
Date: Wed, 22 Jun 2022 21:45:08 +0200
Subject: [PATCH] Add Naga variant to ShaderSource

---
 wgpu/src/backend/direct.rs |  2 ++
 wgpu/src/backend/web.rs    | 16 ++++++++++++++++
 wgpu/src/lib.rs            |  4 ++++
 3 files changed, 22 insertions(+)

diff --git a/wgpu/src/backend/direct.rs b/wgpu/src/backend/direct.rs
index 22f04db41a2..98918655b18 100644
--- a/wgpu/src/backend/direct.rs
+++ b/wgpu/src/backend/direct.rs
@@ -1103,6 +1103,8 @@ impl crate::Context for Context {
                 wgc::pipeline::ShaderModuleSource::Naga(module)
             }
             ShaderSource::Wgsl(ref code) => wgc::pipeline::ShaderModuleSource::Wgsl(Borrowed(code)),
+            #[cfg(feature = "naga")]
+            ShaderSource::Naga(ref module) => wgc::pipeline::ShaderModuleSource::Naga(module.take()),
         };
         let (id, error) = wgc::gfx_select!(
             device.id => global.device_create_shader_module(device.id, &descriptor, source, PhantomData)
diff --git a/wgpu/src/backend/web.rs b/wgpu/src/backend/web.rs
index 7a2e99277a4..45f580ede4f 100644
--- a/wgpu/src/backend/web.rs
+++ b/wgpu/src/backend/web.rs
@@ -1265,6 +1265,22 @@ impl crate::Context for Context {
                 web_sys::GpuShaderModuleDescriptor::new(wgsl_text.as_str())
             }
             crate::ShaderSource::Wgsl(ref code) => web_sys::GpuShaderModuleDescriptor::new(code),
+            #[cfg(feature = "naga")]
+            crate::ShaderSource::Naga(ref cell) => {
+                use naga::{back, valid};
+
+                let module = cell.take();
+                let mut validator = valid::Validator::new(
+                    valid::ValidationFlags::all(),
+                    valid::Capabilities::all(),
+                );
+                let module_info = validator.validate(&module).unwrap();
+
+                let writer_flags = naga::back::wgsl::WriterFlags::empty();
+                let wgsl_text =
+                    back::wgsl::write_string(&module, &module_info, writer_flags).unwrap();
+                web_sys::GpuShaderModuleDescriptor::new(wgsl_text.as_str())
+            }
         };
         if let Some(label) = desc.label {
             descriptor.label(label);
diff --git a/wgpu/src/lib.rs b/wgpu/src/lib.rs
index 187ca59d8e4..06f31baf49d 100644
--- a/wgpu/src/lib.rs
+++ b/wgpu/src/lib.rs
@@ -812,6 +812,10 @@ pub enum ShaderSource<'a> {
     },
     /// WGSL module as a string slice.
     Wgsl(Cow<'a, str>),
+    /// Naga module.
+    #[cfg(feature = "naga")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "naga")))]
+    Naga(std::cell::Cell<naga::Module>),
 }
 
 /// Descriptor for use with [`Device::create_shader_module`].