tint: Support @diagnostic on if statements

Bug: tint:1809
Change-Id: I4137fc2c68c37193bd03e6a752edf0d80a4430c1
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/124060
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: James Price <jrprice@google.com>
diff --git a/src/tint/ast/if_statement.cc b/src/tint/ast/if_statement.cc
index fdea1da..6878c9f 100644
--- a/src/tint/ast/if_statement.cc
+++ b/src/tint/ast/if_statement.cc
@@ -25,8 +25,13 @@
                          const Source& src,
                          const Expression* cond,
                          const BlockStatement* b,
-                         const Statement* else_stmt)
-    : Base(pid, nid, src), condition(cond), body(b), else_statement(else_stmt) {
+                         const Statement* else_stmt,
+                         utils::VectorRef<const Attribute*> attrs)
+    : Base(pid, nid, src),
+      condition(cond),
+      body(b),
+      else_statement(else_stmt),
+      attributes(std::move(attrs)) {
     TINT_ASSERT(AST, condition);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, condition, program_id);
     TINT_ASSERT(AST, body);
@@ -35,6 +40,10 @@
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, else_statement, program_id);
         TINT_ASSERT(AST, (else_statement->IsAnyOf<IfStatement, BlockStatement>()));
     }
+    for (auto* attr : attributes) {
+        TINT_ASSERT(AST, attr);
+        TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, attr, program_id);
+    }
 }
 
 IfStatement::IfStatement(IfStatement&&) = default;
@@ -47,7 +56,8 @@
     auto* cond = ctx->Clone(condition);
     auto* b = ctx->Clone(body);
     auto* el = ctx->Clone(else_statement);
-    return ctx->dst->create<IfStatement>(src, cond, b, el);
+    auto attrs = ctx->Clone(attributes);
+    return ctx->dst->create<IfStatement>(src, cond, b, el, std::move(attrs));
 }
 
 }  // namespace tint::ast
diff --git a/src/tint/ast/if_statement.h b/src/tint/ast/if_statement.h
index 1c1abc3..8f291a8 100644
--- a/src/tint/ast/if_statement.h
+++ b/src/tint/ast/if_statement.h
@@ -32,12 +32,14 @@
     /// @param condition the if condition
     /// @param body the if body
     /// @param else_stmt the else statement, or nullptr
+    /// @param attributes the if statement attributes
     IfStatement(ProgramID pid,
                 NodeID nid,
                 const Source& src,
                 const Expression* condition,
                 const BlockStatement* body,
-                const Statement* else_stmt);
+                const Statement* else_stmt,
+                utils::VectorRef<const Attribute*> attributes);
     /// Move constructor
     IfStatement(IfStatement&&);
     ~IfStatement() override;
@@ -56,6 +58,9 @@
 
     /// The optional else statement, or nullptr
     const Statement* const else_statement;
+
+    /// The attribute list
+    const utils::Vector<const Attribute*, 1> attributes;
 };
 
 }  // namespace tint::ast
diff --git a/src/tint/ast/if_statement_test.cc b/src/tint/ast/if_statement_test.cc
index 9115cb7..b5b85fe 100644
--- a/src/tint/ast/if_statement_test.cc
+++ b/src/tint/ast/if_statement_test.cc
@@ -14,6 +14,7 @@
 
 #include "src/tint/ast/if_statement.h"
 
+#include "gmock/gmock.h"
 #include "gtest/gtest-spi.h"
 #include "src/tint/ast/discard_statement.h"
 #include "src/tint/ast/test_helper.h"
@@ -31,6 +32,15 @@
     EXPECT_EQ(src.range.begin.column, 2u);
 }
 
+TEST_F(IfStatementTest, Creation_WithAttributes) {
+    auto* attr1 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "foo");
+    auto* attr2 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "bar");
+    auto* cond = Expr("cond");
+    auto* stmt = If(cond, Block(), ElseStmt(), utils::Vector{attr1, attr2});
+
+    EXPECT_THAT(stmt->attributes, testing::ElementsAre(attr1, attr2));
+}
+
 TEST_F(IfStatementTest, IsIf) {
     auto* stmt = If(Expr(true), Block());
     EXPECT_TRUE(stmt->Is<IfStatement>());
diff --git a/src/tint/program_builder.h b/src/tint/program_builder.h
index c84a6e0..4c77d75 100644
--- a/src/tint/program_builder.h
+++ b/src/tint/program_builder.h
@@ -3130,14 +3130,16 @@
     /// @param condition the if statement condition expression
     /// @param body the if statement body
     /// @param else_stmt optional else statement
+    /// @param attributes optional attributes
     /// @returns the if statement pointer
     template <typename CONDITION>
     const ast::IfStatement* If(const Source& source,
                                CONDITION&& condition,
                                const ast::BlockStatement* body,
-                               const ElseStmt else_stmt = ElseStmt()) {
+                               const ElseStmt else_stmt = ElseStmt(),
+                               utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
         return create<ast::IfStatement>(source, Expr(std::forward<CONDITION>(condition)), body,
-                                        else_stmt.stmt);
+                                        else_stmt.stmt, std::move(attributes));
     }
 
     /// Creates a ast::IfStatement with input condition, body, and optional
@@ -3145,13 +3147,15 @@
     /// @param condition the if statement condition expression
     /// @param body the if statement body
     /// @param else_stmt optional else statement
+    /// @param attributes optional attributes
     /// @returns the if statement pointer
     template <typename CONDITION>
     const ast::IfStatement* If(CONDITION&& condition,
                                const ast::BlockStatement* body,
-                               const ElseStmt else_stmt = ElseStmt()) {
+                               const ElseStmt else_stmt = ElseStmt(),
+                               utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
         return create<ast::IfStatement>(Expr(std::forward<CONDITION>(condition)), body,
-                                        else_stmt.stmt);
+                                        else_stmt.stmt, std::move(attributes));
     }
 
     /// Creates an Else object.
diff --git a/src/tint/reader/spirv/function.cc b/src/tint/reader/spirv/function.cc
index 9318277..b78719c 100644
--- a/src/tint/reader/spirv/function.cc
+++ b/src/tint/reader/spirv/function.cc
@@ -726,7 +726,7 @@
     /// @param builder the program builder
     /// @returns the built ast::IfStatement
     const ast::IfStatement* Build(ProgramBuilder* builder) const override {
-        return builder->create<ast::IfStatement>(Source{}, cond, body, else_stmt);
+        return builder->create<ast::IfStatement>(Source{}, cond, body, else_stmt, utils::Empty);
     }
 
     /// If-statement condition
@@ -3325,7 +3325,8 @@
         else_block = create<ast::BlockStatement>(StatementList{else_stmt}, utils::Empty);
     }
 
-    auto* if_stmt = create<ast::IfStatement>(Source{}, condition, if_block, else_block);
+    auto* if_stmt =
+        create<ast::IfStatement>(Source{}, condition, if_block, else_block, utils::Empty);
 
     return if_stmt;
 }
diff --git a/src/tint/reader/wgsl/parser_impl.cc b/src/tint/reader/wgsl/parser_impl.cc
index 7649a0b..35a6d15 100644
--- a/src/tint/reader/wgsl/parser_impl.cc
+++ b/src/tint/reader/wgsl/parser_impl.cc
@@ -1264,7 +1264,7 @@
         return stmt;
     }
 
-    auto stmt_if = if_statement();
+    auto stmt_if = if_statement(attrs.value);
     if (stmt_if.errored) {
         return Failure::kErrored;
     }
@@ -1516,11 +1516,14 @@
 }
 
 // if_statement
-//   : IF expression compound_stmt ( ELSE else_stmt ) ?
-// else_stmt
-//  : compound_statement
-//  | if_statement
-Maybe<const ast::IfStatement*> ParserImpl::if_statement() {
+//   : attribute* if_clause else_if_clause* else_clause?
+// if_clause:
+//   : IF expression compound_stmt
+// else_if_clause:
+//   : ELSE IF expression compound_stmt
+// else_clause
+//   : ELSE compound_statement
+Maybe<const ast::IfStatement*> ParserImpl::if_statement(AttributeList& attrs) {
     // Parse if-else chains iteratively instead of recursively, to avoid
     // stack-overflow for long chains of if-else statements.
 
@@ -1528,6 +1531,7 @@
         Source source;
         const ast::Expression* condition;
         const ast::BlockStatement* body;
+        AttributeList attributes;
     };
 
     // Parse an if statement, capturing the source, condition, and body statement.
@@ -1550,7 +1554,8 @@
             return Failure::kErrored;
         }
 
-        return IfInfo{source, condition.value, body.value};
+        TINT_DEFER(attrs.Clear());
+        return IfInfo{source, condition.value, body.value, std::move(attrs)};
     };
 
     std::vector<IfInfo> statements;
@@ -1591,7 +1596,8 @@
 
     // Now walk back through the statements to create their AST nodes.
     for (auto itr = statements.rbegin(); itr != statements.rend(); itr++) {
-        last_stmt = create<ast::IfStatement>(itr->source, itr->condition, itr->body, last_stmt);
+        last_stmt = create<ast::IfStatement>(itr->source, itr->condition, itr->body, last_stmt,
+                                             std::move(itr->attributes));
     }
 
     return last_stmt->As<ast::IfStatement>();
diff --git a/src/tint/reader/wgsl/parser_impl.h b/src/tint/reader/wgsl/parser_impl.h
index bfc2f9a..8f039b7 100644
--- a/src/tint/reader/wgsl/parser_impl.h
+++ b/src/tint/reader/wgsl/parser_impl.h
@@ -488,9 +488,10 @@
     /// Parses a `variable_statement` grammar element
     /// @returns the parsed variable or nullptr
     Maybe<const ast::VariableDeclStatement*> variable_statement();
-    /// Parses a `if_statement` grammar element
+    /// Parses a `if_statement` grammar element, with the attribute list provided as `attrs`.
+    /// @param attrs the list of attributes for the statement
     /// @returns the parsed statement or nullptr
-    Maybe<const ast::IfStatement*> if_statement();
+    Maybe<const ast::IfStatement*> if_statement(AttributeList& attrs);
     /// Parses a `switch_statement` grammar element
     /// @returns the parsed statement or nullptr
     Maybe<const ast::SwitchStatement*> switch_statement();
diff --git a/src/tint/reader/wgsl/parser_impl_if_stmt_test.cc b/src/tint/reader/wgsl/parser_impl_if_stmt_test.cc
index 55effc7..2acb26c 100644
--- a/src/tint/reader/wgsl/parser_impl_if_stmt_test.cc
+++ b/src/tint/reader/wgsl/parser_impl_if_stmt_test.cc
@@ -19,7 +19,8 @@
 
 TEST_F(ParserImplTest, IfStmt) {
     auto p = parser("if a == 4 { a = b; c = d; }");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -34,7 +35,8 @@
 
 TEST_F(ParserImplTest, IfStmt_WithElse) {
     auto p = parser("if a == 4 { a = b; c = d; } else if(c) { d = 2; } else {}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -57,7 +59,8 @@
 
 TEST_F(ParserImplTest, IfStmt_WithElse_WithParens) {
     auto p = parser("if(a==4) { a = b; c = d; } else if(c) { d = 2; } else {}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_TRUE(e.matched);
     EXPECT_FALSE(e.errored);
     EXPECT_FALSE(p->has_error()) << p->error();
@@ -78,9 +81,25 @@
     EXPECT_EQ(el->statements.Length(), 0u);
 }
 
+TEST_F(ParserImplTest, IfStmt_WithAttributes) {
+    auto p = parser(R"(@diagnostic(off, derivative_uniformity) if true { })");
+    auto a = p->attribute_list();
+    auto e = p->if_statement(a.value);
+    EXPECT_TRUE(e.matched);
+    EXPECT_FALSE(e.errored);
+    EXPECT_FALSE(p->has_error()) << p->error();
+    ASSERT_NE(e.value, nullptr);
+    ASSERT_TRUE(e->Is<ast::IfStatement>());
+
+    EXPECT_TRUE(a->IsEmpty());
+    ASSERT_EQ(e->attributes.Length(), 1u);
+    EXPECT_TRUE(e->attributes[0]->Is<ast::DiagnosticAttribute>());
+}
+
 TEST_F(ParserImplTest, IfStmt_InvalidCondition) {
     auto p = parser("if a = 3 {}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -90,7 +109,8 @@
 
 TEST_F(ParserImplTest, IfStmt_MissingCondition) {
     auto p = parser("if {}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -100,7 +120,8 @@
 
 TEST_F(ParserImplTest, IfStmt_InvalidBody) {
     auto p = parser("if a { fn main() {}}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -110,7 +131,8 @@
 
 TEST_F(ParserImplTest, IfStmt_MissingBody) {
     auto p = parser("if a");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -120,7 +142,8 @@
 
 TEST_F(ParserImplTest, IfStmt_InvalidElseif) {
     auto p = parser("if a {} else if a { fn main() -> a{}}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
@@ -130,7 +153,8 @@
 
 TEST_F(ParserImplTest, IfStmt_InvalidElse) {
     auto p = parser("if a {} else { fn main() -> a{}}");
-    auto e = p->if_statement();
+    ParserImpl::AttributeList attrs;
+    auto e = p->if_statement(attrs);
     EXPECT_FALSE(e.matched);
     EXPECT_TRUE(e.errored);
     EXPECT_EQ(e.value, nullptr);
diff --git a/src/tint/reader/wgsl/parser_impl_statement_test.cc b/src/tint/reader/wgsl/parser_impl_statement_test.cc
index 634bfb4..dce5114 100644
--- a/src/tint/reader/wgsl/parser_impl_statement_test.cc
+++ b/src/tint/reader/wgsl/parser_impl_statement_test.cc
@@ -314,6 +314,30 @@
     EXPECT_EQ(sa->condition->source.range.end.column, 19u);
 }
 
+TEST_F(ParserImplTest, Statement_ConsumedAttributes_Block) {
+    auto p = parser("@diagnostic(off, derivative_uniformity) {}");
+    auto e = p->statement();
+    ASSERT_FALSE(p->has_error()) << p->error();
+    EXPECT_TRUE(e.matched);
+    EXPECT_FALSE(e.errored);
+
+    auto* s = As<ast::BlockStatement>(e.value);
+    ASSERT_NE(s, nullptr);
+    EXPECT_EQ(s->attributes.Length(), 1u);
+}
+
+TEST_F(ParserImplTest, Statement_ConsumedAttributes_If) {
+    auto p = parser("@diagnostic(off, derivative_uniformity) if true {}");
+    auto e = p->statement();
+    ASSERT_FALSE(p->has_error()) << p->error();
+    EXPECT_TRUE(e.matched);
+    EXPECT_FALSE(e.errored);
+
+    auto* s = As<ast::IfStatement>(e.value);
+    ASSERT_NE(s, nullptr);
+    EXPECT_EQ(s->attributes.Length(), 1u);
+}
+
 TEST_F(ParserImplTest, Statement_UnexpectedAttributes) {
     auto p = parser("@diagnostic(off, derivative_uniformity) return;");
     auto e = p->statement();
diff --git a/src/tint/resolver/attribute_validation_test.cc b/src/tint/resolver/attribute_validation_test.cc
index 5c09c2a..3430d9d 100644
--- a/src/tint/resolver/attribute_validation_test.cc
+++ b/src/tint/resolver/attribute_validation_test.cc
@@ -1045,6 +1045,39 @@
 12:34 note: first attribute declared here)");
 }
 
+using IfStatementAttributeTest = TestWithParams;
+TEST_P(IfStatementAttributeTest, IsValid) {
+    auto& params = GetParam();
+
+    WrapInFunction(If(Expr(true), Block(), ElseStmt(),
+                      createAttributes(Source{{12, 34}}, *this, params.kind)));
+
+    if (params.should_pass) {
+        EXPECT_TRUE(r()->Resolve()) << r()->error();
+    } else {
+        EXPECT_FALSE(r()->Resolve());
+        EXPECT_EQ(r()->error(), "12:34 error: attribute is not valid for if statements");
+    }
+}
+INSTANTIATE_TEST_SUITE_P(ResolverAttributeValidationTest,
+                         IfStatementAttributeTest,
+                         testing::Values(TestParams{AttributeKind::kAlign, false},
+                                         TestParams{AttributeKind::kBinding, false},
+                                         TestParams{AttributeKind::kBuiltin, false},
+                                         TestParams{AttributeKind::kDiagnostic, true},
+                                         TestParams{AttributeKind::kGroup, false},
+                                         TestParams{AttributeKind::kId, false},
+                                         TestParams{AttributeKind::kInterpolate, false},
+                                         TestParams{AttributeKind::kInvariant, false},
+                                         TestParams{AttributeKind::kLocation, false},
+                                         TestParams{AttributeKind::kMustUse, false},
+                                         TestParams{AttributeKind::kOffset, false},
+                                         TestParams{AttributeKind::kSize, false},
+                                         TestParams{AttributeKind::kStage, false},
+                                         TestParams{AttributeKind::kStride, false},
+                                         TestParams{AttributeKind::kWorkgroup, false},
+                                         TestParams{AttributeKind::kBindingAndGroup, false}));
+
 namespace BlockStatementTests {
 class BlockStatementTest : public TestWithParams {
   protected:
diff --git a/src/tint/resolver/resolver.cc b/src/tint/resolver/resolver.cc
index f6f1876..ff053c3 100644
--- a/src/tint/resolver/resolver.cc
+++ b/src/tint/resolver/resolver.cc
@@ -4274,6 +4274,7 @@
             [&](const ast::BlockStatement* block) {
                 return handle_attributes(block, sem, "block statements");
             },
+            [&](const ast::IfStatement* i) { return handle_attributes(i, sem, "if statements"); },
             [&](Default) { return true; })) {
         return nullptr;
     }
diff --git a/src/tint/resolver/uniformity_test.cc b/src/tint/resolver/uniformity_test.cc
index 863df84..7fa32f6 100644
--- a/src/tint/resolver/uniformity_test.cc
+++ b/src/tint/resolver/uniformity_test.cc
@@ -8376,6 +8376,82 @@
     }
 }
 
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnIfStatement_CallInCondition) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+fn foo() {
+  )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"(if (non_uniform == 42 && dpdx(1.0) > 0.0) {
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'dpdx' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnIfStatement_CallInBody) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+@group(0) @binding(1) var t : texture_2d<f32>;
+@group(0) @binding(2) var s : sampler;
+fn foo() {
+  )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"(if (non_uniform == 42) {
+    let color = textureSample(t, s, vec2(0, 0));
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'textureSample' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnIfStatement_CallInElse) {
+    auto& param = GetParam();
+    utils::StringStream ss;
+    ss << R"(
+@group(0) @binding(0) var<storage, read_write> non_uniform : i32;
+@group(0) @binding(1) var t : texture_2d<f32>;
+@group(0) @binding(2) var s : sampler;
+fn foo() {
+  )"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"(if (non_uniform == 42) {
+  } else {
+    let color = textureSample(t, s, vec2(0, 0));
+  }
+}
+)";
+
+    RunTest(ss.str(), param != builtin::DiagnosticSeverity::kError);
+    if (param == builtin::DiagnosticSeverity::kOff) {
+        EXPECT_TRUE(error_.empty());
+    } else {
+        utils::StringStream err;
+        err << ToStr(param) << ": 'textureSample' must only be called";
+        EXPECT_THAT(error_, ::testing::HasSubstr(err.str()));
+    }
+}
+
 INSTANTIATE_TEST_SUITE_P(UniformityAnalysisTest,
                          UniformityAnalysisDiagnosticFilterTest,
                          ::testing::Values(builtin::DiagnosticSeverity::kError,
diff --git a/src/tint/sem/diagnostic_severity_test.cc b/src/tint/sem/diagnostic_severity_test.cc
index fa57b55..b3b0b3a 100644
--- a/src/tint/sem/diagnostic_severity_test.cc
+++ b/src/tint/sem/diagnostic_severity_test.cc
@@ -31,45 +31,64 @@
         // @diagnostic(off, chromium_unreachable_code)
         // fn foo() {
         //   @diagnostic(info, chromium_unreachable_code) {
+        //     @diagnostic(error, chromium_unreachable_code)
         //     if (true) @diagnostic(warning, chromium_unreachable_code) {
         //       return;
+        //     } else if (false) {
+        //       return;
+        //     } else @diagnostic(info, chromium_unreachable_code) {
+        //       return;
         //     }
+        //     return;
         //   }
         // }
         //
         // fn bar() {
-        //   {
-        //     if (true) {
-        //       return;
-        //     }
-        //   }
+        //   return;
         // }
         auto rule = builtin::DiagnosticRule::kChromiumUnreachableCode;
         auto func_severity = builtin::DiagnosticSeverity::kOff;
         auto block_severity = builtin::DiagnosticSeverity::kInfo;
-        auto if_severity = builtin::DiagnosticSeverity::kInfo;
+        auto if_severity = builtin::DiagnosticSeverity::kError;
+        auto if_body_severity = builtin::DiagnosticSeverity::kWarning;
+        auto else_body_severity = builtin::DiagnosticSeverity::kInfo;
         auto attr = [&](auto severity) {
             return utils::Vector{DiagnosticAttribute(severity, "chromium_unreachable_code")};
         };
 
-        auto* return_1 = Return();
-        auto* if_1 = If(Expr(true), Block(utils::Vector{return_1}, attr(if_severity)));
-        auto* block_1 = Block(utils::Vector{if_1}, attr(block_severity));
+        auto* return_foo_if = Return();
+        auto* return_foo_elseif = Return();
+        auto* return_foo_else = Return();
+        auto* return_foo_block = Return();
+        auto* else_stmt = Block(utils::Vector{return_foo_else}, attr(else_body_severity));
+        auto* elseif = If(Expr(false), Block(return_foo_elseif), Else(else_stmt));
+        auto* if_foo = If(Expr(true), Block(utils::Vector{return_foo_if}, attr(if_body_severity)),
+                          Else(elseif), attr(if_severity));
+        auto* block_1 = Block(utils::Vector{if_foo, return_foo_block}, attr(block_severity));
         auto* func_attr = DiagnosticAttribute(func_severity, "chromium_unreachable_code");
         auto* foo = Func("foo", {}, ty.void_(), utils::Vector{block_1}, utils::Vector{func_attr});
 
-        auto* return_2 = Return();
-        auto* bar = Func("bar", {}, ty.void_(), utils::Vector{return_2});
+        auto* return_bar = Return();
+        auto* bar = Func("bar", {}, ty.void_(), utils::Vector{return_bar});
 
         auto p = Build();
         EXPECT_TRUE(p.IsValid()) << p.Diagnostics().str();
 
         EXPECT_EQ(p.Sem().DiagnosticSeverity(foo, rule), func_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(block_1, rule), block_severity);
-        EXPECT_EQ(p.Sem().DiagnosticSeverity(if_1, rule), block_severity);
-        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_1, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(if_foo, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(if_foo->condition, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(if_foo->body, rule), if_body_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_if, rule), if_body_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(elseif, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(elseif->condition, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(elseif->body, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_elseif, rule), if_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(else_stmt, rule), else_body_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_else, rule), else_body_severity);
+
         EXPECT_EQ(p.Sem().DiagnosticSeverity(bar, rule), global_severity);
-        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_2, rule), global_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_bar, rule), global_severity);
     }
 };
 
diff --git a/src/tint/writer/wgsl/generator_impl.cc b/src/tint/writer/wgsl/generator_impl.cc
index f98cf86..76075f1 100644
--- a/src/tint/writer/wgsl/generator_impl.cc
+++ b/src/tint/writer/wgsl/generator_impl.cc
@@ -942,6 +942,14 @@
 bool GeneratorImpl::EmitIf(const ast::IfStatement* stmt) {
     {
         auto out = line();
+
+        if (!stmt->attributes.IsEmpty()) {
+            if (!EmitAttributes(out, stmt->attributes)) {
+                return false;
+            }
+            out << " ";
+        }
+
         out << "if (";
         if (!EmitExpression(out, stmt->condition)) {
             return false;
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl
new file mode 100644
index 0000000..a1fbf92
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl
@@ -0,0 +1,10 @@
+@group(0) @binding(1) var t : texture_2d<f32>;
+@group(0) @binding(2) var s : sampler;
+
+@fragment
+fn main(@location(0) x : f32) {
+  @diagnostic(warning, derivative_uniformity)
+  if (x > 0) {
+  } else if (dpdx(1.0) > 0)  {
+  }
+}
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.dxc.hlsl b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.dxc.hlsl
new file mode 100644
index 0000000..ec72f1c
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.dxc.hlsl
@@ -0,0 +1,31 @@
+diagnostic_filtering/if_statement_attribute.wgsl:8:14 warning: 'dpdx' must only be called from uniform control flow
+  } else if (dpdx(1.0) > 0)  {
+             ^^^^^^^^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:3 note: control flow depends on possibly non-uniform value
+  if (x > 0) {
+  ^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:7 note: user-defined input 'x' of 'main' may be non-uniform
+  if (x > 0) {
+      ^
+
+Texture2D<float4> t : register(t1, space0);
+SamplerState s : register(s2, space0);
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  if ((x > 0.0f)) {
+  } else {
+    if ((ddx(1.0f) > 0.0f)) {
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.fxc.hlsl b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.fxc.hlsl
new file mode 100644
index 0000000..ec72f1c
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.fxc.hlsl
@@ -0,0 +1,31 @@
+diagnostic_filtering/if_statement_attribute.wgsl:8:14 warning: 'dpdx' must only be called from uniform control flow
+  } else if (dpdx(1.0) > 0)  {
+             ^^^^^^^^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:3 note: control flow depends on possibly non-uniform value
+  if (x > 0) {
+  ^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:7 note: user-defined input 'x' of 'main' may be non-uniform
+  if (x > 0) {
+      ^
+
+Texture2D<float4> t : register(t1, space0);
+SamplerState s : register(s2, space0);
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  if ((x > 0.0f)) {
+  } else {
+    if ((ddx(1.0f) > 0.0f)) {
+    }
+  }
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.glsl b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.glsl
new file mode 100644
index 0000000..79434b5
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.glsl
@@ -0,0 +1,28 @@
+diagnostic_filtering/if_statement_attribute.wgsl:8:14 warning: 'dpdx' must only be called from uniform control flow
+  } else if (dpdx(1.0) > 0)  {
+             ^^^^^^^^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:3 note: control flow depends on possibly non-uniform value
+  if (x > 0) {
+  ^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:7 note: user-defined input 'x' of 'main' may be non-uniform
+  if (x > 0) {
+      ^
+
+#version 310 es
+precision highp float;
+
+layout(location = 0) in float x_1;
+void tint_symbol(float x) {
+  if ((x > 0.0f)) {
+  } else {
+    if ((dFdx(1.0f) > 0.0f)) {
+    }
+  }
+}
+
+void main() {
+  tint_symbol(x_1);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.msl b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.msl
new file mode 100644
index 0000000..6b2d12b
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.msl
@@ -0,0 +1,32 @@
+diagnostic_filtering/if_statement_attribute.wgsl:8:14 warning: 'dpdx' must only be called from uniform control flow
+  } else if (dpdx(1.0) > 0)  {
+             ^^^^^^^^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:3 note: control flow depends on possibly non-uniform value
+  if (x > 0) {
+  ^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:7 note: user-defined input 'x' of 'main' may be non-uniform
+  if (x > 0) {
+      ^
+
+#include <metal_stdlib>
+
+using namespace metal;
+struct tint_symbol_2 {
+  float x [[user(locn0)]];
+};
+
+void tint_symbol_inner(float x) {
+  if ((x > 0.0f)) {
+  } else {
+    if ((dfdx(1.0f) > 0.0f)) {
+    }
+  }
+}
+
+fragment void tint_symbol(tint_symbol_2 tint_symbol_1 [[stage_in]]) {
+  tint_symbol_inner(tint_symbol_1.x);
+  return;
+}
+
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.spvasm b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.spvasm
new file mode 100644
index 0000000..d23951c
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.spvasm
@@ -0,0 +1,73 @@
+diagnostic_filtering/if_statement_attribute.wgsl:8:14 warning: 'dpdx' must only be called from uniform control flow
+  } else if (dpdx(1.0) > 0)  {
+             ^^^^^^^^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:3 note: control flow depends on possibly non-uniform value
+  if (x > 0) {
+  ^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:7 note: user-defined input 'x' of 'main' may be non-uniform
+  if (x > 0) {
+      ^
+
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 31
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main" %x_1
+               OpExecutionMode %main OriginUpperLeft
+               OpName %x_1 "x_1"
+               OpName %t "t"
+               OpName %s "s"
+               OpName %main_inner "main_inner"
+               OpName %x "x"
+               OpName %main "main"
+               OpDecorate %x_1 Location 0
+               OpDecorate %t DescriptorSet 0
+               OpDecorate %t Binding 1
+               OpDecorate %s DescriptorSet 0
+               OpDecorate %s Binding 2
+      %float = OpTypeFloat 32
+%_ptr_Input_float = OpTypePointer Input %float
+        %x_1 = OpVariable %_ptr_Input_float Input
+          %6 = OpTypeImage %float 2D 0 0 0 1 Unknown
+%_ptr_UniformConstant_6 = OpTypePointer UniformConstant %6
+          %t = OpVariable %_ptr_UniformConstant_6 UniformConstant
+          %9 = OpTypeSampler
+%_ptr_UniformConstant_9 = OpTypePointer UniformConstant %9
+          %s = OpVariable %_ptr_UniformConstant_9 UniformConstant
+       %void = OpTypeVoid
+         %10 = OpTypeFunction %void %float
+         %15 = OpConstantNull %float
+       %bool = OpTypeBool
+    %float_1 = OpConstant %float 1
+         %26 = OpTypeFunction %void
+ %main_inner = OpFunction %void None %10
+          %x = OpFunctionParameter %float
+         %14 = OpLabel
+         %16 = OpFOrdGreaterThan %bool %x %15
+               OpSelectionMerge %18 None
+               OpBranchConditional %16 %19 %20
+         %19 = OpLabel
+               OpBranch %18
+         %20 = OpLabel
+         %21 = OpDPdx %float %float_1
+         %23 = OpFOrdGreaterThan %bool %21 %15
+               OpSelectionMerge %24 None
+               OpBranchConditional %23 %25 %24
+         %25 = OpLabel
+               OpBranch %24
+         %24 = OpLabel
+               OpBranch %18
+         %18 = OpLabel
+               OpReturn
+               OpFunctionEnd
+       %main = OpFunction %void None %26
+         %28 = OpLabel
+         %30 = OpLoad %float %x_1
+         %29 = OpFunctionCall %void %main_inner %30
+               OpReturn
+               OpFunctionEnd
diff --git a/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.wgsl b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.wgsl
new file mode 100644
index 0000000..bc8f8dd
--- /dev/null
+++ b/test/tint/diagnostic_filtering/if_statement_attribute.wgsl.expected.wgsl
@@ -0,0 +1,22 @@
+diagnostic_filtering/if_statement_attribute.wgsl:8:14 warning: 'dpdx' must only be called from uniform control flow
+  } else if (dpdx(1.0) > 0)  {
+             ^^^^^^^^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:3 note: control flow depends on possibly non-uniform value
+  if (x > 0) {
+  ^^
+
+diagnostic_filtering/if_statement_attribute.wgsl:7:7 note: user-defined input 'x' of 'main' may be non-uniform
+  if (x > 0) {
+      ^
+
+@group(0) @binding(1) var t : texture_2d<f32>;
+
+@group(0) @binding(2) var s : sampler;
+
+@fragment
+fn main(@location(0) x : f32) {
+  @diagnostic(warning, derivative_uniformity) if ((x > 0)) {
+  } else if ((dpdx(1.0) > 0)) {
+  }
+}