[WebGPU] https://philogb.github.io/page/indraspearls/ fails to compil…

…e shader

https://bugs.webkit.org/show_bug.cgi?id=273070
rdar://126865140

Reviewed by Mike Wyrzykowski.

Add support for attributes before compound statements.

* Source/WebGPU/WGSL/AST/ASTCompoundStatement.h:
* Source/WebGPU/WGSL/AttributeValidator.cpp:
(WGSL::AttributeValidator::visit):
* Source/WebGPU/WGSL/GlobalVariableRewriter.cpp:
(WGSL::RewriteGlobalVariables::initializeVariables):
(WGSL::RewriteGlobalVariables::storeInitialValue):
* Source/WebGPU/WGSL/Parser.cpp:
(WGSL::Parser<Lexer>::parseAttribute):
(WGSL::Parser<Lexer>::parseCompoundStatement):
* Source/WebGPU/WGSL/tests/valid/for.wgsl:

Canonical link: https://commits.webkit.org/278682@main

Source/WebGPU/WGSL/AST/ASTCompoundStatement.h

@@ -35,15 +35,18 @@ class CompoundStatement final : public Statement {
     using Ref = std::reference_wrapper<CompoundStatement>;
 
     NodeKind kind() const override;
+    Attribute::List& attributes() { return m_attributes; }
     Statement::List& statements() { return m_statements; }
     const Statement::List& statements() const { return m_statements; }
 
 private:
-    CompoundStatement(SourceSpan span, Statement::List&& statements)
+    CompoundStatement(SourceSpan span, Attribute::List&& attributes, Statement::List&& statements)
         : Statement(span)
+        , m_attributes(WTFMove(attributes))
         , m_statements(WTFMove(statements))
     { }
 
+    Attribute::List m_attributes;
     Statement::List m_statements;
 };
 

Source/WebGPU/WGSL/AttributeValidator.cpp

@@ -51,6 +51,7 @@ class AttributeValidator : public AST::Visitor {
     void visit(AST::Variable&) override;
     void visit(AST::Structure&) override;
     void visit(AST::StructureMember&) override;
+    void visit(AST::CompoundStatement&) override;
 
 private:
     bool parseBuiltin(AST::Function*, std::optional<Builtin>&, AST::Attribute&);
@@ -385,6 +386,15 @@ void AttributeValidator::visit(AST::StructureMember& member)
     AST::Visitor::visit(member);
 }
 
+void AttributeValidator::visit(AST::CompoundStatement& statement)
+{
+    for (auto& attribute : statement.attributes()) {
+        if (!is<AST::DiagnosticAttribute>(attribute))
+            error(attribute.span(), "invalid attribute for compound statement"_s);
+    }
+    AST::Visitor::visit(statement);
+}
+
 bool AttributeValidator::parseBuiltin(AST::Function* function, std::optional<Builtin>& builtin, AST::Attribute& attribute)
 {
     auto* builtinAttribute = dynamicDowncast<AST::BuiltinAttribute>(attribute);

Source/WebGPU/WGSL/GlobalVariableRewriter.cpp

@@ -1857,6 +1857,7 @@ void RewriteGlobalVariables::initializeVariables(AST::Function& function, const
 
     auto& body = m_shaderModule.astBuilder().construct<AST::CompoundStatement>(
         SourceSpan::empty(),
+        AST::Attribute::List { },
         WTFMove(initializations)
     );
 
@@ -2044,6 +2045,7 @@ void RewriteGlobalVariables::storeInitialValue(AST::Expression& target, AST::Sta
 
         auto& forBody = m_shaderModule.astBuilder().construct<AST::CompoundStatement>(
             SourceSpan::empty(),
+            AST::Attribute::List { },
             WTFMove(forBodyStatements)
         );
 

Source/WebGPU/WGSL/Parser.cpp

@@ -628,6 +628,14 @@ Result<AST::Attribute::Ref> Parser<Lexer>::parseAttribute()
     START_PARSE();
 
     CONSUME_TYPE(Attribute);
+
+    if (current().type == TokenType::KeywordDiagnostic) {
+        consume();
+        PARSE(diagnostic, Diagnostic);
+        RETURN_ARENA_NODE(DiagnosticAttribute, WTFMove(diagnostic));
+    }
+
+
     CONSUME_TYPE_NAMED(ident, Identifier);
 
     if (ident.ident == "group"_s) {
@@ -770,11 +778,6 @@ Result<AST::Attribute::Ref> Parser<Lexer>::parseAttribute()
     if (ident.ident == "const"_s)
         RETURN_ARENA_NODE(ConstAttribute);
 
-    if (ident.ident == "diagnostic"_s) {
-        PARSE(diagnostic, Diagnostic);
-        RETURN_ARENA_NODE(DiagnosticAttribute, WTFMove(diagnostic));
-    }
-
     // https://gpuweb.github.io/gpuweb/wgsl/#pipeline-stage-attributes
     if (ident.ident == "vertex"_s)
         RETURN_ARENA_NODE(StageAttribute, ShaderStage::Vertex);
@@ -1209,6 +1212,8 @@ Result<AST::CompoundStatement::Ref> Parser<Lexer>::parseCompoundStatement()
 {
     START_PARSE();
 
+    PARSE(attributes, Attributes);
+
     CONSUME_TYPE(BraceLeft);
 
     AST::Statement::List statements;
@@ -1224,7 +1229,7 @@ Result<AST::CompoundStatement::Ref> Parser<Lexer>::parseCompoundStatement()
 
     CONSUME_TYPE(BraceRight);
 
-    RETURN_ARENA_NODE(CompoundStatement, WTFMove(statements));
+    RETURN_ARENA_NODE(CompoundStatement, WTFMove(attributes), WTFMove(statements));
 }
 
 template<typename Lexer>

Source/WebGPU/WGSL/tests/valid/for.wgsl

@@ -3,7 +3,9 @@
 @compute
 @workgroup_size(1, 1, 1)
 fn testForStatement() {
-    for (var i = -1; i <= 1; i++) {
+    for (var i = -1; i <= 1; i++)
+    @diagnostic(off,derivative_uniformity)
+    {
     }
 
     var i = 0;