WIP: feat(wgsl-in)!: allow opt-out of wgpu-specific enable-extensions

TODO: `CHANGELOG`
TODO: document
TODO: Explain new error surface on opt-out

Author: ErichDonGubler

naga/src/front/wgsl/parse/directive/enable_extension.rs

@@ -17,33 +17,70 @@ pub struct EnableExtensions {
     /// Whether `enable clip_distances;` was written earlier in the shader module.
     clip_distances: bool,
     /// Native-only extensions.
-    wgpu: WgpuEnableExtensions,
+    wgpu: NativeOnly<WgpuEnableExtensions>,
 }
 
 impl EnableExtensions {
-    pub(crate) const fn empty() -> Self {
+    pub(crate) const fn empty(native_only: NativeOnly<()>) -> Self {
+        // TODO: It'd be nice to use a `const` `map` helper here, once
+        // <https://github.com/rust-lang/rust/issues/73255> has landed.
+        let wgpu = match native_only {
+            NativeOnly::NotAllowed {
+                can_suggest_in_diagnostics,
+            } => NativeOnly::NotAllowed {
+                can_suggest_in_diagnostics,
+            },
+            NativeOnly::Allowed(()) => NativeOnly::Allowed(WgpuEnableExtensions::empty()),
+        };
         Self {
-            wgpu: WgpuEnableExtensions::empty(),
+            wgpu,
             f16: false,
             dual_source_blending: false,
             clip_distances: false,
         }
     }
 
     /// Add an enable-extension to the set requested by a module.
-    pub(crate) fn add(&mut self, ext: ImplementedEnableExtension) {
+    pub(crate) fn enable(
+        &mut self,
+        span: Span,
+        ext: ImplementedEnableExtension,
+    ) -> core::result::Result<(), EnableExtensionNotAvailableError> {
+        macro_rules! check_wgpu {
+            ($f:expr) => {{
+                let f: fn(&mut WgpuEnableExtensions) -> &mut bool = $f;
+                match &mut self.wgpu {
+                    &mut NativeOnly::Allowed(ref mut wgpu) => Ok(f(wgpu)),
+                    &mut NativeOnly::NotAllowed {
+                        can_suggest_in_diagnostics,
+                    } => Err(EnableExtensionNotAvailableError {
+                        span,
+                        extension: ext,
+                        reason: EnableExtensionNotAvailableErrorReason::NotAllowed {
+                            can_suggest_in_diagnostics,
+                        },
+                    }),
+                }
+            }};
+        }
         let field = match ext {
-            ImplementedEnableExtension::WgpuMeshShader => &mut self.wgpu.mesh_shader,
-            ImplementedEnableExtension::WgpuRayQuery => &mut self.wgpu.ray_query,
+            ImplementedEnableExtension::WgpuMeshShader => {
+                check_wgpu!(|wgpu| &mut wgpu.mesh_shader)?
+            }
+            ImplementedEnableExtension::WgpuRayQuery => check_wgpu!(|wgpu| &mut wgpu.ray_query)?,
             ImplementedEnableExtension::WgpuRayQueryVertexReturn => {
-                &mut self.wgpu.ray_query_vertex_return
+                check_wgpu!(|wgpu| &mut wgpu.ray_query_vertex_return)?
             }
             ImplementedEnableExtension::DualSourceBlending => &mut self.dual_source_blending,
             ImplementedEnableExtension::F16 => &mut self.f16,
             ImplementedEnableExtension::ClipDistances => &mut self.clip_distances,
-            ImplementedEnableExtension::WgpuCooperativeMatrix => &mut self.wgpu.cooperative_matrix,
+            ImplementedEnableExtension::WgpuCooperativeMatrix => {
+                check_wgpu!(|wgpu| &mut wgpu.cooperative_matrix)?
+            }
         };
         *field = true;
+
+        Ok(())
     }
 
     /// Query whether an enable-extension tracked here has been requested.
@@ -52,16 +89,31 @@ impl EnableExtensions {
         span: Span,
         extension: ImplementedEnableExtension,
     ) -> core::result::Result<(), EnableExtensionNotAvailableError> {
+        let check_wgpu = |f: fn(&WgpuEnableExtensions) -> bool| match &self.wgpu {
+            &NativeOnly::Allowed(ref wgpu) => Ok(f(wgpu)),
+            &NativeOnly::NotAllowed {
+                can_suggest_in_diagnostics,
+            } => Err(EnableExtensionNotAvailableError {
+                span,
+                extension,
+                reason: EnableExtensionNotAvailableErrorReason::NotAllowed {
+                    can_suggest_in_diagnostics,
+                },
+            }),
+        };
+
         let is_enabled = match extension {
-            ImplementedEnableExtension::WgpuMeshShader => self.wgpu.mesh_shader,
-            ImplementedEnableExtension::WgpuRayQuery => self.wgpu.ray_query,
+            ImplementedEnableExtension::WgpuMeshShader => check_wgpu(|wgpu| wgpu.mesh_shader)?,
+            ImplementedEnableExtension::WgpuRayQuery => check_wgpu(|wgpu| wgpu.ray_query)?,
             ImplementedEnableExtension::WgpuRayQueryVertexReturn => {
-                self.wgpu.ray_query_vertex_return
+                check_wgpu(|wgpu| wgpu.ray_query_vertex_return)?
             }
             ImplementedEnableExtension::DualSourceBlending => self.dual_source_blending,
             ImplementedEnableExtension::F16 => self.f16,
             ImplementedEnableExtension::ClipDistances => self.clip_distances,
-            ImplementedEnableExtension::WgpuCooperativeMatrix => self.wgpu.cooperative_matrix,
+            ImplementedEnableExtension::WgpuCooperativeMatrix => {
+                check_wgpu(|wgpu| wgpu.cooperative_matrix)?
+            }
         };
 
         if is_enabled {
@@ -78,7 +130,45 @@ impl EnableExtensions {
 
 impl Default for EnableExtensions {
     fn default() -> Self {
-        Self::empty()
+        Self::empty(NativeOnly::default())
+    }
+}
+
+/// Native-only extensions under [`EnableExtensions::wgpu`].
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum NativeOnly<T> {
+    NotAllowed {
+        /// Whether the diagnostic of a validation check encountering this case may contain a hint
+        /// about enabling native-only extensions.
+        ///
+        /// Set this to `false` when you are running user-provided shaders in non-native
+        /// environments.
+        can_suggest_in_diagnostics: bool,
+    },
+    Allowed(T),
+}
+
+impl<T> NativeOnly<T> {
+    pub fn map<U, F>(self, f: F) -> NativeOnly<U>
+    where
+        F: FnOnce(T) -> U,
+    {
+        match self {
+            Self::NotAllowed {
+                can_suggest_in_diagnostics,
+            } => NativeOnly::NotAllowed {
+                can_suggest_in_diagnostics,
+            },
+            Self::Allowed(t) => NativeOnly::Allowed(f(t)),
+        }
+    }
+}
+
+impl<T> Default for NativeOnly<T> {
+    fn default() -> Self {
+        Self::NotAllowed {
+            can_suggest_in_diagnostics: false,
+        }
     }
 }
 
@@ -124,6 +214,7 @@ pub struct EnableExtensionNotAvailableError {
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum EnableExtensionNotAvailableErrorReason {
     NotEnabled,
+    NotAllowed { can_suggest_in_diagnostics: bool },
 }
 
 /// An enable-extension not guaranteed to be present in all environments.

naga/src/front/wgsl/parse/lexer.rs

@@ -1,6 +1,6 @@
 use super::{number::consume_number, Error, ExpectedToken, Result};
 use crate::front::wgsl::error::NumberError;
-use crate::front::wgsl::parse::directive::enable_extension::EnableExtensions;
+use crate::front::wgsl::parse::directive::enable_extension::{EnableExtensions, NativeOnly};
 use crate::front::wgsl::parse::{conv, Number};
 use crate::front::wgsl::Scalar;
 use crate::Span;
@@ -243,6 +243,17 @@ fn is_word_part(c: char) -> bool {
     unicode_ident::is_xid_continue(c)
 }
 
+// TODO: document
+#[derive(Clone, Eq, PartialEq)]
+pub struct LexerConfig {
+    /// Whether to enable native-only [enable-extensions](EnableExtensions).
+    native_only_enable_extensions: NativeOnly<()>,
+
+    /// Whether or not to ignore doc comments.
+    /// If `true`, doc comments are treated as [`Token::Trivia`].
+    ignore_doc_comments: bool,
+}
+
 #[derive(Clone)]
 pub(in crate::front::wgsl) struct Lexer<'a> {
     /// The remaining unconsumed input.
@@ -262,20 +273,23 @@ pub(in crate::front::wgsl) struct Lexer<'a> {
     /// statements.
     last_end_offset: usize,
 
-    /// Whether or not to ignore doc comments.
-    /// If `true`, doc comments are treated as [`Token::Trivia`].
+    /// See [`LexerConfig::ignore_doc_comments`].
     ignore_doc_comments: bool,
 
     pub(in crate::front::wgsl) enable_extensions: EnableExtensions,
 }
 
 impl<'a> Lexer<'a> {
-    pub(in crate::front::wgsl) const fn new(input: &'a str, ignore_doc_comments: bool) -> Self {
+    pub(in crate::front::wgsl) const fn new(input: &'a str, config: LexerConfig) -> Self {
+        let LexerConfig {
+            native_only_enable_extensions,
+            ignore_doc_comments,
+        } = config;
         Lexer {
             input,
             source: input,
             last_end_offset: 0,
-            enable_extensions: EnableExtensions::empty(),
+            enable_extensions: EnableExtensions::empty(native_only_enable_extensions),
             ignore_doc_comments,
         }
     }