[shaders] Support "enable" via post-process directive

The WGSL source currently cannot contain the "enable" directive as naga
doesn't support it (see gfx-rs/wgpu#5476). This patch works around this
limitation by introducing the "#enable" post-process directive.

Lines that start with "#enable" get turned into an inline comment during
the pre-process stage and converted to a standard "enable" directive
following the intermediate module compilation step. This allows us to
pass the WGSL shaders with enable directives on to other WebGPU
implementations.

Author: armansito

crates/shaders/src/compile/mod.rs

@@ -120,7 +120,7 @@ impl ShaderInfo {
         bindings.sort_by_key(|res| res.location);
         let workgroup_size = entry.workgroup_size;
         Ok(ShaderInfo {
-            source,
+            source: postprocess(&source),
             module,
             module_info,
             workgroup_size,
@@ -182,3 +182,19 @@ impl ShaderInfo {
         info
     }
 }
+
+// TODO: This is a workaround for gfx-rs/wgpu#5476. Since naga can't handle the `enable` directive,
+// we allow its use in other WGSL compilers using our own "#enable" post-process directive. Remove
+// this mechanism once naga supports the directive.
+fn postprocess(wgsl: &str) -> String {
+    let mut output = String::with_capacity(wgsl.len());
+    for line in wgsl.lines() {
+        if line.starts_with("//__#enable") {
+            output.push_str(&line["//__#".len()..]);
+        } else {
+            output.push_str(line);
+        }
+        output.push('\n');
+    }
+    output
+}

crates/shaders/src/compile/preprocess.rs

@@ -65,9 +65,11 @@ pub fn preprocess(
             let directive_is_at_start = line.trim_start().starts_with('#');
 
             match directive {
-                if_item @ ("ifdef" | "ifndef" | "else" | "endif") if !directive_is_at_start => {
+                item @ ("ifdef" | "ifndef" | "else" | "endif" | "enable")
+                    if !directive_is_at_start =>
+                {
                     eprintln!(
-                        "#{if_item} directives must be the first non_whitespace items on \
+                        "#{item} directives must be the first non_whitespace items on \
                                their line, ignoring (line {line_number})"
                     );
                     break;
@@ -111,7 +113,7 @@ pub fn preprocess(
                         eprintln!("Mismatched endif (line {line_number})");
                     }
                     let remainder = directive_start[directive_len..].trim();
-                    if !remainder.is_empty() {
+                    if !remainder.is_empty() && !remainder.starts_with("//") {
                         eprintln!(
                             "#endif directives don't take an argument. `{remainder}` will \
                                    not be in output (line {line_number})"
@@ -151,6 +153,16 @@ pub fn preprocess(
                     }
                     continue;
                 }
+                "enable" => {
+                    // Turn this directive into a comment. It will be handled as part in
+                    // postprocess.
+                    if stack.iter().all(|item| item.active) {
+                        output.push_str("//__");
+                        output.push_str(line);
+                        output.push('\n');
+                    }
+                    continue 'all_lines;
+                }
                 val => {
                     eprintln!("Unknown preprocessor directive `{val}` (line {line_number})");
                 }

shader/fine.wgsl

@@ -7,6 +7,12 @@
 // To enable multisampled rendering, turn on both the msaa ifdef and one of msaa8
 // or msaa16.
 
+#ifdef r8
+// The R8 variant is only available via an internal extension in Dawn
+// (see https://dawn.googlesource.com/dawn/+/refs/heads/main/docs/tint/extensions/chromium_internal_graphite.md).
+#enable chromium_internal_graphite;
+#endif
+
 struct Tile {
     backdrop: i32,
     segments: u32,
@@ -34,10 +40,6 @@ var<storage> info: array<u32>;
 
 @group(0) @binding(4)
 #ifdef r8
-// The R8 variant is available via a non-standard extension in Dawn. We can't
-// optionally declare that extension here since naga doesn't understand the
-// `enable` directive (see https://github.com/gfx-rs/wgpu/issues/5476). The
-// directive must be injected by the client via some other means.
 var output: texture_storage_2d<r8unorm, write>;
 #else
 var output: texture_storage_2d<rgba8unorm, write>;
@@ -816,6 +818,8 @@ fn extend_mode(t: f32, mode: u32) -> f32 {
 
 let PIXELS_PER_THREAD = 4u;
 
+#ifndef msaa
+
 // Analytic area antialiasing.
 //
 // This is currently dead code if msaa is enabled, but it would be fairly straightforward
@@ -878,6 +882,8 @@ fn fill_path(fill: CmdFill, xy: vec2<f32>, result: ptr<function, array<f32, PIXE
     *result = area;
 }
 
+#endif
+
 // The X size should be 16 / PIXELS_PER_THREAD
 @compute @workgroup_size(4, 16)
 fn main(

src/shaders/preprocess.rs

@@ -69,7 +69,9 @@ pub fn preprocess(input: &str, defines: &HashSet<String>, imports: &HashMap<&str
             let directive_is_at_start = line.trim_start().starts_with('#');
 
             match directive {
-                if_item @ ("ifdef" | "ifndef" | "else" | "endif") if !directive_is_at_start => {
+                if_item @ ("ifdef" | "ifndef" | "else" | "endif" | "enable")
+                    if !directive_is_at_start =>
+                {
                     eprintln!("#{if_item} directives must be the first non_whitespace items on their line, ignoring (line {line_number})");
                     break;
                 }
@@ -142,6 +144,11 @@ pub fn preprocess(input: &str, defines: &HashSet<String>, imports: &HashMap<&str
                     }
                     continue;
                 }
+                "enable" => {
+                    // Ignore post-process directive. This is only supported by the shaders crate
+                    // (see #467)
+                    continue 'all_lines;
+                }
                 val => {
                     eprintln!("Unknown preprocessor directive `{val}` (line {line_number})");
                 }