Add support for wgsl directives.

src/compose/error.rs

@@ -134,6 +134,12 @@ pub enum ComposerErrorInner {
     },
     #[error("#define statements are only allowed at the start of the top-level shaders")]
     DefineInModule(usize),
+    #[error("Invalid WGSL directive '{directive}' at line {position}: {reason}")]
+    InvalidWgslDirective {
+        directive: String,
+        position: usize,
+        reason: String,
+    },
 }
 
 struct ErrorSources<'a> {
@@ -239,7 +245,8 @@ impl ComposerError {
             | ComposerErrorInner::OverrideNotVirtual { pos, .. }
             | ComposerErrorInner::GlslInvalidVersion(pos)
             | ComposerErrorInner::DefineInModule(pos)
-            | ComposerErrorInner::InvalidShaderDefDefinitionValue { pos, .. } => {
+            | ComposerErrorInner::InvalidShaderDefDefinitionValue { pos, .. }
+            | ComposerErrorInner::InvalidWgslDirective { position: pos, .. } => {
                 (vec![Label::primary((), *pos..*pos)], vec![])
             }
             ComposerErrorInner::WgslBackError(e) => {

src/compose/mod.rs

@@ -140,13 +140,17 @@ use crate::{
 
 pub use self::error::{ComposerError, ComposerErrorInner, ErrSource};
 use self::preprocess::Preprocessor;
+pub use self::wgsl_directives::{
+    DiagnosticDirective, EnableDirective, RequiresDirective, WgslDirectives,
+};
 
 pub mod comment_strip_iter;
 pub mod error;
 pub mod parse_imports;
 pub mod preprocess;
 mod test;
 pub mod tokenizer;
+pub mod wgsl_directives;
 
 #[derive(Hash, PartialEq, Eq, Clone, Copy, Debug, Default)]
 pub enum ShaderLanguage {
@@ -270,8 +274,6 @@ pub struct ComposableModuleDefinition {
     modules: HashMap<ModuleKey, ComposableModule>,
     // used in spans when this module is included
     module_index: usize,
-    // preprocessor meta data
-    // metadata: PreprocessorMetaData,
 }
 
 impl ComposableModuleDefinition {
@@ -580,11 +582,17 @@ impl Composer {
         language: ShaderLanguage,
         imports: &[ImportDefinition],
         shader_defs: &HashMap<String, ShaderDefValue>,
+        wgsl_directives: Option<WgslDirectives>,
     ) -> Result<IrBuildResult, ComposerError> {
         debug!("creating IR for {} with defs: {:?}", name, shader_defs);
 
+        let mut wgsl_string = String::new();
+        if let Some(wgsl_directives) = &wgsl_directives {
+            trace!("adding WGSL directives for {}", name);
+            wgsl_string = wgsl_directives.to_wgsl_string();
+        }
         let mut module_string = match language {
-            ShaderLanguage::Wgsl => String::new(),
+            ShaderLanguage::Wgsl => wgsl_string,
             #[cfg(feature = "glsl")]
             ShaderLanguage::Glsl => String::from("#version 450\n"),
         };
@@ -842,6 +850,7 @@ impl Composer {
         demote_entrypoints: bool,
         source: &str,
         imports: Vec<ImportDefWithOffset>,
+        wgsl_directives: Option<WgslDirectives>,
     ) -> Result<ComposableModule, ComposerError> {
         let mut imports: Vec<_> = imports
             .into_iter()
@@ -975,6 +984,7 @@ impl Composer {
             module_definition.language,
             &imports,
             shader_defs,
+            wgsl_directives,
         )?;
 
         // from here on errors need to be reported using the modified source with start_offset
@@ -1376,6 +1386,7 @@ impl Composer {
             true,
             &preprocessed_source,
             imports,
+            None,
         )
         .map_err(|err| err.into())
     }
@@ -1519,6 +1530,7 @@ impl Composer {
             name: module_name,
             mut imports,
             mut effective_defs,
+            cleaned_source,
             ..
         } = self
             .preprocessor
@@ -1595,7 +1607,7 @@ impl Composer {
 
         let module_set = ComposableModuleDefinition {
             name: module_name.clone(),
-            sanitized_source: substituted_source,
+            sanitized_source: cleaned_source,
             file_path: file_path.to_owned(),
             language,
             effective_defs: effective_defs.into_iter().collect(),
@@ -1650,7 +1662,13 @@ impl Composer {
 
         let sanitized_source = self.sanitize_and_set_auto_bindings(source);
 
-        let PreprocessorMetaData { name, defines, .. } = self
+        let PreprocessorMetaData {
+            name,
+            defines,
+            wgsl_directives,
+            cleaned_source,
+            ..
+        } = self
             .preprocessor
             .get_preprocessor_metadata(&sanitized_source, true)
             .map_err(|inner| ComposerError {
@@ -1667,7 +1685,7 @@ impl Composer {
 
         let PreprocessOutput { imports, .. } = self
             .preprocessor
-            .preprocess(&sanitized_source, &shader_defs)
+            .preprocess(&cleaned_source, &shader_defs)
             .map_err(|inner| ComposerError {
                 inner,
                 source: ErrSource::Constructing {
@@ -1734,7 +1752,7 @@ impl Composer {
 
         let definition = ComposableModuleDefinition {
             name,
-            sanitized_source: sanitized_source.clone(),
+            sanitized_source: cleaned_source.clone(),
             language: shader_type.into(),
             file_path: file_path.to_owned(),
             module_index: 0,
@@ -1751,7 +1769,7 @@ impl Composer {
             imports,
         } = self
             .preprocessor
-            .preprocess(&sanitized_source, &shader_defs)
+            .preprocess(&cleaned_source, &shader_defs)
             .map_err(|inner| ComposerError {
                 inner,
                 source: ErrSource::Constructing {
@@ -1770,6 +1788,7 @@ impl Composer {
                 false,
                 &preprocessed_source,
                 imports,
+                Some(wgsl_directives),
             )
             .map_err(|e| ComposerError {
                 inner: e.inner,

src/compose/preprocess.rs

@@ -6,6 +6,7 @@ use regex::Regex;
 use super::{
     comment_strip_iter::CommentReplaceExt,
     parse_imports::{parse_imports, substitute_identifiers},
+    wgsl_directives::{DiagnosticDirective, EnableDirective, RequiresDirective, WgslDirectives},
     ComposerErrorInner, ImportDefWithOffset, ShaderDefValue,
 };
 
@@ -22,6 +23,9 @@ pub struct Preprocessor {
     import_regex: Regex,
     define_import_path_regex: Regex,
     define_shader_def_regex: Regex,
+    enable_regex: Regex,
+    requires_regex: Regex,
+    diagnostic_regex: Regex,
 }
 
 impl Default for Preprocessor {
@@ -42,6 +46,10 @@ impl Default for Preprocessor {
             define_import_path_regex: Regex::new(r"^\s*#\s*define_import_path\s+([^\s]+)").unwrap(),
             define_shader_def_regex: Regex::new(r"^\s*#\s*define\s+([\w|\d|_]+)\s*([-\w|\d]+)?")
                 .unwrap(),
+            enable_regex: Regex::new(r"^\s*enable\s+([^;]+)\s*;").unwrap(),
+            requires_regex: Regex::new(r"^\s*requires\s+([^;]+)\s*;").unwrap(),
+            diagnostic_regex: Regex::new(r"^\s*diagnostic\s*\(\s*([^,]+)\s*,\s*([^)]+)\s*\)\s*;")
+                .unwrap(),
         }
     }
 }
@@ -52,6 +60,8 @@ pub struct PreprocessorMetaData {
     pub imports: Vec<ImportDefWithOffset>,
     pub defines: HashMap<String, ShaderDefValue>,
     pub effective_defs: HashSet<String>,
+    pub wgsl_directives: WgslDirectives,
+    pub cleaned_source: String,
 }
 
 enum ScopeLevel {
@@ -133,6 +143,145 @@ pub struct PreprocessOutput {
 }
 
 impl Preprocessor {
+    fn parse_enable_directive(
+        &self,
+        line: &str,
+        line_idx: usize,
+    ) -> Result<Option<EnableDirective>, ComposerErrorInner> {
+        if let Some(cap) = self.enable_regex.captures(line) {
+            let extensions_str = cap.get(1).unwrap().as_str().trim();
+            let extensions: Vec<String> = extensions_str
+                .split(',')
+                .map(|s| s.trim().to_string())
+                .filter(|s| !s.is_empty())
+                .collect();
+
+            if extensions.is_empty() {
+                return Err(ComposerErrorInner::InvalidWgslDirective {
+                    directive: line.to_string(),
+                    position: line_idx,
+                    reason: "No extensions specified".to_string(),
+                });
+            }
+
+            Ok(Some(EnableDirective {
+                extensions,
+                source_location: line_idx,
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn parse_requires_directive(
+        &self,
+        line: &str,
+        line_idx: usize,
+    ) -> Result<Option<RequiresDirective>, ComposerErrorInner> {
+        if let Some(cap) = self.requires_regex.captures(line) {
+            let extensions_str = cap.get(1).unwrap().as_str().trim();
+            let extensions: Vec<String> = extensions_str
+                .split(',')
+                .map(|s| s.trim().to_string())
+                .filter(|s| !s.is_empty())
+                .collect();
+
+            if extensions.is_empty() {
+                return Err(ComposerErrorInner::InvalidWgslDirective {
+                    directive: line.to_string(),
+                    position: line_idx,
+                    reason: "No extensions specified".to_string(),
+                });
+            }
+
+            Ok(Some(RequiresDirective {
+                extensions,
+                source_location: line_idx,
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn parse_diagnostic_directive(
+        &self,
+        line: &str,
+        line_idx: usize,
+    ) -> Result<Option<DiagnosticDirective>, ComposerErrorInner> {
+        if let Some(cap) = self.diagnostic_regex.captures(line) {
+            let severity = cap.get(1).unwrap().as_str().trim().to_string();
+            let rule = cap.get(2).unwrap().as_str().trim().to_string();
+
+            match severity.as_str() {
+                "off" | "info" | "warn" | "error" => {}
+                _ => {
+                    return Err(ComposerErrorInner::InvalidWgslDirective {
+                        directive: line.to_string(),
+                        position: line_idx,
+                        reason: format!(
+                            "Invalid severity '{}'. Must be one of: off, info, warn, error",
+                            severity
+                        ),
+                    });
+                }
+            }
+
+            Ok(Some(DiagnosticDirective {
+                severity,
+                rule,
+                source_location: line_idx,
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
+    pub fn extract_wgsl_directives(
+        &self,
+        source: &str,
+    ) -> Result<(String, WgslDirectives), ComposerErrorInner> {
+        let mut directives = WgslDirectives::default();
+        let mut cleaned_lines = Vec::new();
+        let mut in_directive_section = true;
+
+        for (line_idx, line) in source.lines().enumerate() {
+            let trimmed = line.trim();
+
+            if trimmed.is_empty() || trimmed.starts_with("//") {
+                cleaned_lines.push(line);
+                continue;
+            }
+
+            if in_directive_section {
+                if let Some(enable) = self.parse_enable_directive(trimmed, line_idx)? {
+                    cleaned_lines.push("");
+                    directives.enables.push(enable);
+                    continue;
+                } else if let Some(requires) = self.parse_requires_directive(trimmed, line_idx)? {
+                    cleaned_lines.push("");
+                    directives.requires.push(requires);
+                    continue;
+                } else if let Some(diagnostic) =
+                    self.parse_diagnostic_directive(trimmed, line_idx)?
+                {
+                    cleaned_lines.push("");
+                    directives.diagnostics.push(diagnostic);
+                    continue;
+                } else if !trimmed.starts_with("enable")
+                    && !trimmed.starts_with("requires")
+                    && !trimmed.starts_with("diagnostic")
+                {
+                    in_directive_section = false;
+                }
+            }
+
+            cleaned_lines.push(line);
+        }
+
+        let cleaned_source = cleaned_lines.join("\n");
+        Ok((cleaned_source, directives))
+    }
+
     fn check_scope<'a>(
         &self,
         shader_defs: &HashMap<String, ShaderDefValue>,
@@ -379,6 +528,8 @@ impl Preprocessor {
         shader_str: &str,
         allow_defines: bool,
     ) -> Result<PreprocessorMetaData, ComposerErrorInner> {
+        let (shader_str, wgsl_directives) = self.extract_wgsl_directives(shader_str)?;
+
         let mut declared_imports = IndexMap::default();
         let mut used_imports = IndexMap::default();
         let mut name = None;
@@ -477,6 +628,8 @@ impl Preprocessor {
             imports: used_imports.into_values().collect(),
             defines,
             effective_defs,
+            wgsl_directives,
+            cleaned_source: shader_str.to_string(),
         })
     }
 }