From b4ffb884946101466681a8dd0eb81c14788bc285 Mon Sep 17 00:00:00 2001
From: Arman Uguray <armansito@google.com>
Date: Wed, 10 May 2023 14:03:50 -0700
Subject: [PATCH] [shaders] Revise access to backend-agnostic metadata

Previously the generated shader data structures were rooted in
backend-specific top-level mods (`mod wgsl`, `mod msl`, etc). This made
access to per-shader information that is common to all backends (e.g.
workgroup sizes, shader name etc) awkward to access from backend
agnostic code, especially when feature-gated conditional compilation is
used on the client side.

The data structures have been rearranged such that there is a top-level
`ComputeShader` declaration for each stage under a `gen` mod. The
`ComputeShader` struct declares feature-gated fields for backend shader
sources, such that backend specific data is now a leaf node in the
structure rather than the root. This has some additional benefits:

1. Common data doesn't have to be redeclared, saving on code size when
   multiple backends are enabled.

2. The backend specific source code was previously encoded as a `[u8]`.
   We can now use types that more closely match the expected format, for
   example `&str` for WGSL and MSL, `[u32]` for SPIR-V, etc.

3. If we ever need to expose additional backend-specific metadata in the
   future, we can bundle them alongside the source code in a
   backend-specific data structure at this level of the tree.
---
 crates/shaders/build.rs   | 33 ++++++++++++---------------------
 crates/shaders/src/lib.rs |  9 ++++++++-
 2 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/crates/shaders/build.rs b/crates/shaders/build.rs
index 7e5a952..8c948c1 100644
--- a/crates/shaders/build.rs
+++ b/crates/shaders/build.rs
@@ -21,18 +21,7 @@ fn main() {
     shaders.sort_by(|x, y| x.0.cmp(&y.0));
     let mut buf = String::default();
     write_types(&mut buf, &shaders).unwrap();
-    if cfg!(feature = "wgsl") {
-        write_shaders(&mut buf, "wgsl", &shaders, |info| {
-            info.source.as_bytes().to_owned()
-        })
-        .unwrap();
-    }
-    if cfg!(feature = "msl") {
-        write_shaders(&mut buf, "msl", &shaders, |info| {
-            compile::msl::translate(info).unwrap().as_bytes().to_owned()
-        })
-        .unwrap();
-    }
+    write_shaders(&mut buf, &shaders).unwrap();
     std::fs::write(&dest_path, &buf).unwrap();
     println!("cargo:rerun-if-changed=../shader");
 }
@@ -65,11 +54,9 @@ fn write_types(buf: &mut String, shaders: &[(String, ShaderInfo)]) -> Result<(),
 
 fn write_shaders(
     buf: &mut String,
-    mod_name: &str,
     shaders: &[(String, ShaderInfo)],
-    translate: impl Fn(&ShaderInfo) -> Vec<u8>,
 ) -> Result<(), std::fmt::Error> {
-    writeln!(buf, "pub mod {mod_name} {{")?;
+    writeln!(buf, "mod gen {{")?;
     writeln!(buf, "    use super::*;")?;
     writeln!(buf, "    use BindType::*;")?;
     writeln!(buf, "    pub const SHADERS: Shaders<'static> = Shaders {{")?;
@@ -80,14 +67,8 @@ fn write_shaders(
             .map(|binding| binding.ty)
             .collect::<Vec<_>>();
         let wg_bufs = &info.workgroup_buffers;
-        let source = translate(info);
         writeln!(buf, "        {name}: ComputeShader {{")?;
         writeln!(buf, "            name: Cow::Borrowed({:?}),", name)?;
-        writeln!(
-            buf,
-            "            code: Cow::Borrowed(&{:?}),",
-            source.as_slice()
-        )?;
         writeln!(
             buf,
             "            workgroup_size: {:?},",
@@ -99,6 +80,16 @@ fn write_shaders(
             "            workgroup_buffers: Cow::Borrowed(&{:?}),",
             wg_bufs
         )?;
+        if cfg!(feature = "wgsl") {
+            writeln!(buf, "            wgsl: Cow::Borrowed(&{:?}),", info.source)?;
+        }
+        if cfg!(feature = "msl") {
+            writeln!(
+                buf,
+                "            msl: Cow::Borrowed(&{:?}),",
+                compile::msl::translate(info).unwrap()
+            )?;
+        }
         writeln!(buf, "        }},")?;
     }
     writeln!(buf, "    }};")?;
diff --git a/crates/shaders/src/lib.rs b/crates/shaders/src/lib.rs
index 66ac937..73b621d 100644
--- a/crates/shaders/src/lib.rs
+++ b/crates/shaders/src/lib.rs
@@ -13,10 +13,15 @@ use std::borrow::Cow;
 #[derive(Clone, Debug)]
 pub struct ComputeShader<'a> {
     pub name: Cow<'a, str>,
-    pub code: Cow<'a, [u8]>,
     pub workgroup_size: [u32; 3],
     pub bindings: Cow<'a, [BindType]>,
     pub workgroup_buffers: Cow<'a, [WorkgroupBufferInfo]>,
+
+    #[cfg(feature = "wgsl")]
+    pub wgsl: Cow<'a, str>,
+
+    #[cfg(feature = "msl")]
+    pub msl: Cow<'a, str>,
 }
 
 pub trait PipelineHost {
@@ -32,3 +37,5 @@ pub trait PipelineHost {
 }
 
 include!(concat!(env!("OUT_DIR"), "/shaders.rs"));
+
+pub use gen::SHADERS;