From fd9851af18a9800b00c8c51e3e90286409ad9f0d Mon Sep 17 00:00:00 2001
From: Erich Gubler <erichdongubler@gmail.com>
Date: Fri, 23 Aug 2024 13:46:24 -0400
Subject: [PATCH] WIP: feat(wgsl-in): parse `diagnostic` attrs. on `fn`s

---
 naga/src/front/wgsl/parse/ast.rs |  1 +
 naga/src/front/wgsl/parse/mod.rs | 19 +++++++++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/naga/src/front/wgsl/parse/ast.rs b/naga/src/front/wgsl/parse/ast.rs
index 9385108934..9b0fd8a7aa 100644
--- a/naga/src/front/wgsl/parse/ast.rs
+++ b/naga/src/front/wgsl/parse/ast.rs
@@ -133,6 +133,7 @@ pub struct Function<'a> {
     pub arguments: Vec<FunctionArgument<'a>>,
     pub result: Option<FunctionResult<'a>>,
     pub body: Block<'a>,
+    pub diagnostic_filter_head: Option<Handle<DiagnosticFilterNode>>,
 }
 
 #[derive(Debug)]
diff --git a/naga/src/front/wgsl/parse/mod.rs b/naga/src/front/wgsl/parse/mod.rs
index e4e48ad4a4..49ab8d512a 100644
--- a/naga/src/front/wgsl/parse/mod.rs
+++ b/naga/src/front/wgsl/parse/mod.rs
@@ -2207,6 +2207,7 @@ impl Parser {
     fn function_decl<'a>(
         &mut self,
         lexer: &mut Lexer<'a>,
+        diagnostic_filter_head: Option<Handle<DiagnosticFilterNode>>,
         out: &mut ast::TranslationUnit<'a>,
         dependencies: &mut FastIndexSet<ast::Dependency<'a>>,
     ) -> Result<ast::Function<'a>, Error<'a>> {
@@ -2279,6 +2280,7 @@ impl Parser {
             arguments,
             result,
             body,
+            diagnostic_filter_head,
         };
 
         // done
@@ -2339,17 +2341,16 @@ impl Parser {
             types: &mut out.types,
             unresolved: &mut dependencies,
         };
+        let mut diagnostic_filters = DiagnosticFilterMap::new();
 
         self.push_rule_span(Rule::Attribute, lexer);
         while lexer.skip(Token::Attribute) {
             let (name, name_span) = lexer.next_ident_with_span()?;
             if let Some(DirectiveKind::Diagnostic) = DirectiveKind::from_ident(name) {
-                if let Some(_filter) = self.diagnostic_filter(lexer)? {
+                if let Some(filter) = self.diagnostic_filter(lexer)? {
                     let span = self.peek_rule_span(lexer);
-                    return Err(Error::DiagnosticNotYetImplementedAtParseSite {
-                        span,
-                        site_name_plural: "functions",
-                    });
+                    diagnostic_filters.add(filter, span)?;
+                    // TODO: ensure that it makes sense to apply this attribute
                 }
             }
             match name {
@@ -2491,7 +2492,13 @@ impl Parser {
                 Some(ast::GlobalDeclKind::Var(var))
             }
             (Token::Word("fn"), _) => {
-                let function = self.function_decl(lexer, out, &mut dependencies)?;
+                let diagnostic_filter_head = Self::write_diagnostic_filters(
+                    &mut out.diagnostic_filters,
+                    diagnostic_filters,
+                    out.diagnostic_filter_leaf,
+                );
+                let function =
+                    self.function_decl(lexer, diagnostic_filter_head, out, &mut dependencies)?;
                 Some(ast::GlobalDeclKind::Fn(ast::Function {
                     entry_point: if let Some(stage) = stage.value {
                         if stage == ShaderStage::Compute && workgroup_size.value.is_none() {