tint: Support @diagnostic on switch body

Bug: tint:1809
Change-Id: I33259080c143b0cf7afe01343a1f179f61078dd3
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/124520
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/switch_statement.cc b/src/tint/ast/switch_statement.cc
index 4e5d95d..fa6fe07 100644
--- a/src/tint/ast/switch_statement.cc
+++ b/src/tint/ast/switch_statement.cc
@@ -27,8 +27,13 @@
                                  const Source& src,
                                  const Expression* cond,
                                  utils::VectorRef<const CaseStatement*> b,
-                                 utils::VectorRef<const Attribute*> attrs)
-    : Base(pid, nid, src), condition(cond), body(std::move(b)), attributes(std::move(attrs)) {
+                                 utils::VectorRef<const Attribute*> stmt_attrs,
+                                 utils::VectorRef<const Attribute*> body_attrs)
+    : Base(pid, nid, src),
+      condition(cond),
+      body(std::move(b)),
+      attributes(std::move(stmt_attrs)),
+      body_attributes(std::move(body_attrs)) {
     TINT_ASSERT(AST, condition);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, condition, program_id);
     for (auto* stmt : body) {
@@ -39,6 +44,10 @@
         TINT_ASSERT(AST, attr);
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, attr, program_id);
     }
+    for (auto* attr : body_attributes) {
+        TINT_ASSERT(AST, attr);
+        TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, attr, program_id);
+    }
 }
 
 SwitchStatement::~SwitchStatement() = default;
@@ -49,7 +58,9 @@
     auto* cond = ctx->Clone(condition);
     auto b = ctx->Clone(body);
     auto attrs = ctx->Clone(attributes);
-    return ctx->dst->create<SwitchStatement>(src, cond, std::move(b), std::move(attrs));
+    auto body_attrs = ctx->Clone(body_attributes);
+    return ctx->dst->create<SwitchStatement>(src, cond, std::move(b), std::move(attrs),
+                                             std::move(body_attrs));
 }
 
 }  // namespace tint::ast
diff --git a/src/tint/ast/switch_statement.h b/src/tint/ast/switch_statement.h
index 354b6ad..f2b8201 100644
--- a/src/tint/ast/switch_statement.h
+++ b/src/tint/ast/switch_statement.h
@@ -29,13 +29,15 @@
     /// @param src the source of this node
     /// @param condition the switch condition
     /// @param body the switch body
-    /// @param attributes the switch statement attributes
+    /// @param stmt_attributes the switch statement attributes
+    /// @param body_attributes the switch body attributes
     SwitchStatement(ProgramID pid,
                     NodeID nid,
                     const Source& src,
                     const Expression* condition,
                     utils::VectorRef<const CaseStatement*> body,
-                    utils::VectorRef<const Attribute*> attributes);
+                    utils::VectorRef<const Attribute*> stmt_attributes,
+                    utils::VectorRef<const Attribute*> body_attributes);
 
     /// Destructor
     ~SwitchStatement() override;
@@ -53,8 +55,11 @@
     const utils::Vector<const CaseStatement*, 4> body;
     SwitchStatement(const SwitchStatement&) = delete;
 
-    /// The attribute list
+    /// The attribute list for the statement
     const utils::Vector<const Attribute*, 1> attributes;
+
+    /// The attribute list for the body
+    const utils::Vector<const Attribute*, 1> body_attributes;
 };
 
 }  // namespace tint::ast
diff --git a/src/tint/ast/switch_statement_test.cc b/src/tint/ast/switch_statement_test.cc
index aab2604..795220d 100644
--- a/src/tint/ast/switch_statement_test.cc
+++ b/src/tint/ast/switch_statement_test.cc
@@ -30,7 +30,7 @@
     auto* ident = Expr("ident");
     utils::Vector body{case_stmt};
 
-    auto* stmt = create<SwitchStatement>(ident, body, utils::Empty);
+    auto* stmt = create<SwitchStatement>(ident, body, utils::Empty, utils::Empty);
     EXPECT_EQ(stmt->condition, ident);
     ASSERT_EQ(stmt->body.Length(), 1u);
     EXPECT_EQ(stmt->body[0], case_stmt);
@@ -38,8 +38,8 @@
 
 TEST_F(SwitchStatementTest, Creation_WithSource) {
     auto* ident = Expr("ident");
-    auto* stmt =
-        create<SwitchStatement>(Source{Source::Location{20, 2}}, ident, utils::Empty, utils::Empty);
+    auto* stmt = create<SwitchStatement>(Source{Source::Location{20, 2}}, ident, utils::Empty,
+                                         utils::Empty, utils::Empty);
     auto src = stmt->source;
     EXPECT_EQ(src.range.begin.line, 20u);
     EXPECT_EQ(src.range.begin.column, 2u);
@@ -49,17 +49,28 @@
     auto* attr1 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "foo");
     auto* attr2 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "bar");
     auto* ident = Expr("ident");
-    auto* stmt = create<SwitchStatement>(ident, utils::Empty, utils::Vector{attr1, attr2});
+    auto* stmt =
+        create<SwitchStatement>(ident, utils::Empty, utils::Vector{attr1, attr2}, utils::Empty);
 
     EXPECT_THAT(stmt->attributes, testing::ElementsAre(attr1, attr2));
 }
 
+TEST_F(SwitchStatementTest, Creation_WithBodyAttributes) {
+    auto* attr1 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "foo");
+    auto* attr2 = DiagnosticAttribute(builtin::DiagnosticSeverity::kOff, "bar");
+    auto* ident = Expr("ident");
+    auto* stmt =
+        create<SwitchStatement>(ident, utils::Empty, utils::Empty, utils::Vector{attr1, attr2});
+
+    EXPECT_THAT(stmt->body_attributes, testing::ElementsAre(attr1, attr2));
+}
+
 TEST_F(SwitchStatementTest, IsSwitch) {
     utils::Vector lit{CaseSelector(2_i)};
     auto* ident = Expr("ident");
     utils::Vector body{create<CaseStatement>(lit, Block())};
 
-    auto* stmt = create<SwitchStatement>(ident, body, utils::Empty);
+    auto* stmt = create<SwitchStatement>(ident, body, utils::Empty, utils::Empty);
     EXPECT_TRUE(stmt->Is<SwitchStatement>());
 }
 
@@ -71,7 +82,7 @@
             CaseStatementList cases;
             cases.Push(
                 b.create<CaseStatement>(utils::Vector{b.CaseSelector(b.Expr(1_i))}, b.Block()));
-            b.create<SwitchStatement>(nullptr, cases, utils::Empty);
+            b.create<SwitchStatement>(nullptr, cases, utils::Empty, utils::Empty);
         },
         "internal compiler error");
 }
@@ -81,7 +92,8 @@
     EXPECT_FATAL_FAILURE(
         {
             ProgramBuilder b;
-            b.create<SwitchStatement>(b.Expr(true), CaseStatementList{nullptr}, utils::Empty);
+            b.create<SwitchStatement>(b.Expr(true), CaseStatementList{nullptr}, utils::Empty,
+                                      utils::Empty);
         },
         "internal compiler error");
 }
@@ -99,7 +111,7 @@
                                                },
                                                b1.Block()),
                                        },
-                                       utils::Empty);
+                                       utils::Empty, utils::Empty);
         },
         "internal compiler error");
 }
@@ -117,7 +129,7 @@
                                                },
                                                b2.Block()),
                                        },
-                                       utils::Empty);
+                                       utils::Empty, utils::Empty);
         },
         "internal compiler error");
 }
diff --git a/src/tint/program_builder.h b/src/tint/program_builder.h
index bcfaf6e..0c00452 100644
--- a/src/tint/program_builder.h
+++ b/src/tint/program_builder.h
@@ -3384,7 +3384,7 @@
             source, Expr(std::forward<ExpressionInit>(condition)),
             utils::Vector<const ast::CaseStatement*, sizeof...(cases)>{
                 std::forward<Cases>(cases)...},
-            utils::Empty);
+            utils::Empty, utils::Empty);
     }
 
     /// Creates a ast::SwitchStatement with input expression and cases
@@ -3400,37 +3400,42 @@
             Expr(std::forward<ExpressionInit>(condition)),
             utils::Vector<const ast::CaseStatement*, sizeof...(cases)>{
                 std::forward<Cases>(cases)...},
-            utils::Empty);
+            utils::Empty, utils::Empty);
     }
 
     /// Creates a ast::SwitchStatement with input expression, cases, and optional attributes
     /// @param source the source information
     /// @param condition the condition expression initializer
     /// @param cases case statements
-    /// @param attributes optional attributes
+    /// @param stmt_attributes optional statement attributes
+    /// @param body_attributes optional body attributes
     /// @returns the switch statement pointer
     template <typename ExpressionInit>
     const ast::SwitchStatement* Switch(
         const Source& source,
         ExpressionInit&& condition,
         utils::VectorRef<const ast::CaseStatement*> cases,
-        utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
+        utils::VectorRef<const ast::Attribute*> stmt_attributes = utils::Empty,
+        utils::VectorRef<const ast::Attribute*> body_attributes = utils::Empty) {
         return create<ast::SwitchStatement>(source, Expr(std::forward<ExpressionInit>(condition)),
-                                            cases, std::move(attributes));
+                                            cases, std::move(stmt_attributes),
+                                            std::move(body_attributes));
     }
 
     /// Creates a ast::SwitchStatement with input expression, cases, and optional attributes
     /// @param condition the condition expression initializer
     /// @param cases case statements
-    /// @param attributes optional attributes
+    /// @param stmt_attributes optional statement attributes
+    /// @param body_attributes optional body attributes
     /// @returns the switch statement pointer
     template <typename ExpressionInit, typename = DisableIfSource<ExpressionInit>>
     const ast::SwitchStatement* Switch(
         ExpressionInit&& condition,
         utils::VectorRef<const ast::CaseStatement*> cases,
-        utils::VectorRef<const ast::Attribute*> attributes = utils::Empty) {
+        utils::VectorRef<const ast::Attribute*> stmt_attributes = utils::Empty,
+        utils::VectorRef<const ast::Attribute*> body_attributes = utils::Empty) {
         return create<ast::SwitchStatement>(Expr(std::forward<ExpressionInit>(condition)), cases,
-                                            std::move(attributes));
+                                            std::move(stmt_attributes), std::move(body_attributes));
     }
 
     /// Creates a ast::CaseStatement with input list of selectors, and body
diff --git a/src/tint/reader/spirv/function.cc b/src/tint/reader/spirv/function.cc
index 413e58b..e729a17 100644
--- a/src/tint/reader/spirv/function.cc
+++ b/src/tint/reader/spirv/function.cc
@@ -706,8 +706,7 @@
         auto reversed_cases = cases;
         std::reverse(reversed_cases.begin(), reversed_cases.end());
 
-        return builder->create<ast::SwitchStatement>(Source{}, condition, std::move(reversed_cases),
-                                                     utils::Empty);
+        return builder->Switch(Source{}, condition, std::move(reversed_cases));
     }
 
     /// Switch statement condition
diff --git a/src/tint/reader/wgsl/parser_impl.cc b/src/tint/reader/wgsl/parser_impl.cc
index 3ec5416..753f60b 100644
--- a/src/tint/reader/wgsl/parser_impl.cc
+++ b/src/tint/reader/wgsl/parser_impl.cc
@@ -1620,6 +1620,11 @@
         return add_error(peek(), "unable to parse selector expression");
     }
 
+    auto body_attrs = attribute_list();
+    if (body_attrs.errored) {
+        return Failure::kErrored;
+    }
+
     auto body = expect_brace_block("switch statement", [&]() -> Expect<CaseStatementList> {
         bool errored = false;
         CaseStatementList list;
@@ -1645,7 +1650,8 @@
     }
 
     TINT_DEFER(attrs.Clear());
-    return create<ast::SwitchStatement>(source, condition.value, body.value, std::move(attrs));
+    return create<ast::SwitchStatement>(source, condition.value, body.value, std::move(attrs),
+                                        std::move(body_attrs.value));
 }
 
 // switch_body
diff --git a/src/tint/reader/wgsl/parser_impl_switch_stmt_test.cc b/src/tint/reader/wgsl/parser_impl_switch_stmt_test.cc
index ef0b1b4..89b4ba8 100644
--- a/src/tint/reader/wgsl/parser_impl_switch_stmt_test.cc
+++ b/src/tint/reader/wgsl/parser_impl_switch_stmt_test.cc
@@ -95,7 +95,7 @@
 }
 
 TEST_F(ParserImplTest, SwitchStmt_WithAttributes) {
-    auto p = parser(R"(@diagnostic(off, derivative_uniformity) switch a { default{} })");
+    auto p = parser("@diagnostic(off, derivative_uniformity) switch a { default{} }");
     auto a = p->attribute_list();
     auto e = p->switch_statement(a.value);
     EXPECT_TRUE(e.matched);
@@ -109,6 +109,21 @@
     EXPECT_TRUE(e->attributes[0]->Is<ast::DiagnosticAttribute>());
 }
 
+TEST_F(ParserImplTest, SwitchStmt_WithBodyAttributes) {
+    auto p = parser("switch a @diagnostic(off, derivative_uniformity) { default{} }");
+    ParserImpl::AttributeList attrs;
+    auto e = p->switch_statement(attrs);
+    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::SwitchStatement>());
+
+    EXPECT_TRUE(e->attributes.IsEmpty());
+    ASSERT_EQ(e->body_attributes.Length(), 1u);
+    EXPECT_TRUE(e->body_attributes[0]->Is<ast::DiagnosticAttribute>());
+}
+
 TEST_F(ParserImplTest, SwitchStmt_InvalidExpression) {
     auto p = parser("switch a=b {}");
     ParserImpl::AttributeList attrs;
diff --git a/src/tint/resolver/attribute_validation_test.cc b/src/tint/resolver/attribute_validation_test.cc
index 1e123f6..45dca0b 100644
--- a/src/tint/resolver/attribute_validation_test.cc
+++ b/src/tint/resolver/attribute_validation_test.cc
@@ -1078,6 +1078,39 @@
                                          TestParams{AttributeKind::kWorkgroup, false},
                                          TestParams{AttributeKind::kBindingAndGroup, false}));
 
+using SwitchBodyAttributeTest = TestWithParams;
+TEST_P(SwitchBodyAttributeTest, IsValid) {
+    auto& params = GetParam();
+
+    WrapInFunction(Switch(Expr(0_a), utils::Vector{DefaultCase()}, utils::Empty,
+                          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 switch body");
+    }
+}
+INSTANTIATE_TEST_SUITE_P(ResolverAttributeValidationTest,
+                         SwitchBodyAttributeTest,
+                         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}));
+
 using IfStatementAttributeTest = TestWithParams;
 TEST_P(IfStatementAttributeTest, IsValid) {
     auto& params = GetParam();
diff --git a/src/tint/resolver/resolver.cc b/src/tint/resolver/resolver.cc
index 5158aed..fb8af76 100644
--- a/src/tint/resolver/resolver.cc
+++ b/src/tint/resolver/resolver.cc
@@ -3975,6 +3975,22 @@
             return false;
         }
 
+        // Handle switch body attributes.
+        for (auto* attr : stmt->body_attributes) {
+            Mark(attr);
+            if (auto* dc = attr->As<ast::DiagnosticAttribute>()) {
+                if (!DiagnosticControl(dc->control)) {
+                    return false;
+                }
+            } else {
+                AddError("attribute is not valid for switch body", attr->source);
+                return false;
+            }
+        }
+        if (!validator_.NoDuplicateAttributes(stmt->body_attributes)) {
+            return false;
+        }
+
         utils::Vector<sem::CaseStatement*, 4> cases;
         cases.Reserve(stmt->body.Length());
         for (auto* case_stmt : stmt->body) {
@@ -3986,6 +4002,8 @@
             cases.Push(c);
             behaviors.Add(c->Behaviors());
             sem->Cases().emplace_back(c);
+
+            ApplyDiagnosticSeverities(c);
         }
 
         if (behaviors.Contains(sem::Behavior::kBreak)) {
diff --git a/src/tint/resolver/uniformity_test.cc b/src/tint/resolver/uniformity_test.cc
index 4cb95dd..0bf2436 100644
--- a/src/tint/resolver/uniformity_test.cc
+++ b/src/tint/resolver/uniformity_test.cc
@@ -8734,6 +8734,34 @@
     }
 }
 
+TEST_P(UniformityAnalysisDiagnosticFilterTest, AttributeOnSwitchBody_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() {
+  switch (non_uniform))"
+       << "@diagnostic(" << param << ", derivative_uniformity)"
+       << R"( {
+    default {
+      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, AttributeOnWhileStatement_CallInCondition) {
     auto& param = GetParam();
     utils::StringStream ss;
diff --git a/src/tint/sem/diagnostic_severity_test.cc b/src/tint/sem/diagnostic_severity_test.cc
index 06d5a7e..9d36ea2 100644
--- a/src/tint/sem/diagnostic_severity_test.cc
+++ b/src/tint/sem/diagnostic_severity_test.cc
@@ -42,7 +42,7 @@
         //     return;
         //
         //     @diagnostic(error, chromium_unreachable_code)
-        //     switch (42) {
+        //     switch (42) @diagnostic(off, chromium_unreachable_code) {
         //       case 0 @diagnostic(warning, chromium_unreachable_code) {
         //         return;
         //       }
@@ -81,6 +81,7 @@
         auto if_body_severity = builtin::DiagnosticSeverity::kWarning;
         auto else_body_severity = builtin::DiagnosticSeverity::kInfo;
         auto switch_severity = builtin::DiagnosticSeverity::kError;
+        auto switch_body_severity = builtin::DiagnosticSeverity::kOff;
         auto case_severity = builtin::DiagnosticSeverity::kWarning;
         auto for_severity = builtin::DiagnosticSeverity::kError;
         auto for_body_severity = builtin::DiagnosticSeverity::kWarning;
@@ -109,8 +110,9 @@
                           Else(elseif), attr(if_severity));
         auto* case_stmt =
             Case(CaseSelector(0_a), Block(utils::Vector{return_foo_case}, attr(case_severity)));
-        auto* swtch = Switch(42_a, utils::Vector{case_stmt, DefaultCase(Block(return_foo_default))},
-                             attr(switch_severity));
+        auto* default_stmt = DefaultCase(Block(return_foo_default));
+        auto* swtch = Switch(42_a, utils::Vector{case_stmt, default_stmt}, attr(switch_severity),
+                             attr(switch_body_severity));
         auto* fl =
             For(Decl(Var("i", ty.i32())), false, Increment("i"),
                 Block(utils::Vector{return_foo_for}, attr(for_body_severity)), attr(for_severity));
@@ -144,10 +146,11 @@
         EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_else, rule), else_body_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(swtch, rule), switch_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(swtch->condition, rule), switch_severity);
-        EXPECT_EQ(p.Sem().DiagnosticSeverity(case_stmt, rule), switch_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(case_stmt, rule), switch_body_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(case_stmt->body, rule), case_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_case, rule), case_severity);
-        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_default, rule), switch_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(default_stmt, rule), switch_body_severity);
+        EXPECT_EQ(p.Sem().DiagnosticSeverity(return_foo_default, rule), switch_body_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(fl, rule), while_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(fl->initializer, rule), for_severity);
         EXPECT_EQ(p.Sem().DiagnosticSeverity(fl->condition, rule), for_severity);
diff --git a/src/tint/writer/glsl/generator_impl_switch_test.cc b/src/tint/writer/glsl/generator_impl_switch_test.cc
index 8e2f739..ce978dc 100644
--- a/src/tint/writer/glsl/generator_impl_switch_test.cc
+++ b/src/tint/writer/glsl/generator_impl_switch_test.cc
@@ -31,7 +31,7 @@
     auto* case_stmt = create<ast::CaseStatement>(utils::Vector{CaseSelector(5_i)}, case_body);
 
     auto* cond = Expr("cond");
-    auto* s = create<ast::SwitchStatement>(cond, utils::Vector{case_stmt, def}, utils::Empty);
+    auto* s = Switch(cond, utils::Vector{case_stmt, def});
     WrapInFunction(s);
 
     GeneratorImpl& gen = Build();
@@ -58,7 +58,7 @@
                                            def_body);
 
     auto* cond = Expr("cond");
-    auto* s = create<ast::SwitchStatement>(cond, utils::Vector{def}, utils::Empty);
+    auto* s = Switch(cond, utils::Vector{def});
     WrapInFunction(s);
 
     GeneratorImpl& gen = Build();
diff --git a/src/tint/writer/wgsl/generator_impl.cc b/src/tint/writer/wgsl/generator_impl.cc
index 52df265..e1bf9fc 100644
--- a/src/tint/writer/wgsl/generator_impl.cc
+++ b/src/tint/writer/wgsl/generator_impl.cc
@@ -1229,7 +1229,16 @@
         if (!EmitExpression(out, stmt->condition)) {
             return false;
         }
-        out << ") {";
+        out << ") ";
+
+        if (!stmt->body_attributes.IsEmpty()) {
+            if (!EmitAttributes(out, stmt->body_attributes)) {
+                return false;
+            }
+            out << " ";
+        }
+
+        out << "{";
     }
 
     {
diff --git a/test/tint/diagnostic_filtering/switch_body_attribute.wgsl b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl
new file mode 100644
index 0000000..a012a9b
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl
@@ -0,0 +1,8 @@
+@fragment
+fn main(@location(0) x : f32) {
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+    default {
+      _ = dpdx(1.0);
+    }
+  }
+}
diff --git a/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.dxc.hlsl b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.dxc.hlsl
new file mode 100644
index 0000000..eb7b7d4
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.dxc.hlsl
@@ -0,0 +1,30 @@
+diagnostic_filtering/switch_body_attribute.wgsl:5:11 warning: 'dpdx' must only be called from uniform control flow
+      _ = dpdx(1.0);
+          ^^^^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:3 note: control flow depends on possibly non-uniform value
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+  ^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:15 note: user-defined input 'x' of 'main' may be non-uniform
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+              ^
+
+int tint_ftoi(float v) {
+  return ((v < 2147483520.0f) ? ((v < -2147483648.0f) ? -2147483648 : int(v)) : 2147483647);
+}
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  tint_ftoi(x);
+  do {
+  } while (false);
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.fxc.hlsl b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.fxc.hlsl
new file mode 100644
index 0000000..eb7b7d4
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.fxc.hlsl
@@ -0,0 +1,30 @@
+diagnostic_filtering/switch_body_attribute.wgsl:5:11 warning: 'dpdx' must only be called from uniform control flow
+      _ = dpdx(1.0);
+          ^^^^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:3 note: control flow depends on possibly non-uniform value
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+  ^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:15 note: user-defined input 'x' of 'main' may be non-uniform
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+              ^
+
+int tint_ftoi(float v) {
+  return ((v < 2147483520.0f) ? ((v < -2147483648.0f) ? -2147483648 : int(v)) : 2147483647);
+}
+
+struct tint_symbol_1 {
+  float x : TEXCOORD0;
+};
+
+void main_inner(float x) {
+  tint_ftoi(x);
+  do {
+  } while (false);
+}
+
+void main(tint_symbol_1 tint_symbol) {
+  main_inner(tint_symbol.x);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.glsl b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.glsl
new file mode 100644
index 0000000..06f82a0
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.glsl
@@ -0,0 +1,32 @@
+diagnostic_filtering/switch_body_attribute.wgsl:5:11 warning: 'dpdx' must only be called from uniform control flow
+      _ = dpdx(1.0);
+          ^^^^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:3 note: control flow depends on possibly non-uniform value
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+  ^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:15 note: user-defined input 'x' of 'main' may be non-uniform
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+              ^
+
+#version 310 es
+precision highp float;
+
+layout(location = 0) in float x_1;
+int tint_ftoi(float v) {
+  return ((v < 2147483520.0f) ? ((v < -2147483648.0f) ? (-2147483647 - 1) : int(v)) : 2147483647);
+}
+
+void tint_symbol(float x) {
+  switch(tint_ftoi(x)) {
+    default: {
+      break;
+    }
+  }
+}
+
+void main() {
+  tint_symbol(x_1);
+  return;
+}
diff --git a/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.msl b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.msl
new file mode 100644
index 0000000..e1b9131
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.msl
@@ -0,0 +1,36 @@
+diagnostic_filtering/switch_body_attribute.wgsl:5:11 warning: 'dpdx' must only be called from uniform control flow
+      _ = dpdx(1.0);
+          ^^^^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:3 note: control flow depends on possibly non-uniform value
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+  ^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:15 note: user-defined input 'x' of 'main' may be non-uniform
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+              ^
+
+#include <metal_stdlib>
+
+using namespace metal;
+int tint_ftoi(float v) {
+  return select(2147483647, select(int(v), (-2147483647 - 1), (v < -2147483648.0f)), (v < 2147483520.0f));
+}
+
+struct tint_symbol_2 {
+  float x [[user(locn0)]];
+};
+
+void tint_symbol_inner(float x) {
+  switch(tint_ftoi(x)) {
+    default: {
+      break;
+    }
+  }
+}
+
+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/switch_body_attribute.wgsl.expected.spvasm b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.spvasm
new file mode 100644
index 0000000..69b01be
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.spvasm
@@ -0,0 +1,68 @@
+diagnostic_filtering/switch_body_attribute.wgsl:5:11 warning: 'dpdx' must only be called from uniform control flow
+      _ = dpdx(1.0);
+          ^^^^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:3 note: control flow depends on possibly non-uniform value
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+  ^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:15 note: user-defined input 'x' of 'main' may be non-uniform
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+              ^
+
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 32
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main" %x_1
+               OpExecutionMode %main OriginUpperLeft
+               OpName %x_1 "x_1"
+               OpName %tint_ftoi "tint_ftoi"
+               OpName %v "v"
+               OpName %main_inner "main_inner"
+               OpName %x "x"
+               OpName %main "main"
+               OpDecorate %x_1 Location 0
+      %float = OpTypeFloat 32
+%_ptr_Input_float = OpTypePointer Input %float
+        %x_1 = OpVariable %_ptr_Input_float Input
+        %int = OpTypeInt 32 1
+          %4 = OpTypeFunction %int %float
+%float_2_14748352e_09 = OpConstant %float 2.14748352e+09
+       %bool = OpTypeBool
+%float_n2_14748365e_09 = OpConstant %float -2.14748365e+09
+%int_n2147483648 = OpConstant %int -2147483648
+%int_2147483647 = OpConstant %int 2147483647
+       %void = OpTypeVoid
+         %19 = OpTypeFunction %void %float
+         %27 = OpTypeFunction %void
+  %tint_ftoi = OpFunction %int None %4
+          %v = OpFunctionParameter %float
+          %8 = OpLabel
+         %11 = OpFOrdLessThan %bool %v %float_2_14748352e_09
+         %15 = OpFOrdLessThan %bool %v %float_n2_14748365e_09
+         %17 = OpConvertFToS %int %v
+         %13 = OpSelect %int %15 %int_n2147483648 %17
+          %9 = OpSelect %int %11 %13 %int_2147483647
+               OpReturnValue %9
+               OpFunctionEnd
+ %main_inner = OpFunction %void None %19
+          %x = OpFunctionParameter %float
+         %23 = OpLabel
+         %25 = OpFunctionCall %int %tint_ftoi %x
+               OpSelectionMerge %24 None
+               OpSwitch %25 %26
+         %26 = OpLabel
+               OpBranch %24
+         %24 = OpLabel
+               OpReturn
+               OpFunctionEnd
+       %main = OpFunction %void None %27
+         %29 = OpLabel
+         %31 = OpLoad %float %x_1
+         %30 = OpFunctionCall %void %main_inner %31
+               OpReturn
+               OpFunctionEnd
diff --git a/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.wgsl b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.wgsl
new file mode 100644
index 0000000..0c3a256
--- /dev/null
+++ b/test/tint/diagnostic_filtering/switch_body_attribute.wgsl.expected.wgsl
@@ -0,0 +1,20 @@
+diagnostic_filtering/switch_body_attribute.wgsl:5:11 warning: 'dpdx' must only be called from uniform control flow
+      _ = dpdx(1.0);
+          ^^^^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:3 note: control flow depends on possibly non-uniform value
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+  ^^^^^^
+
+diagnostic_filtering/switch_body_attribute.wgsl:3:15 note: user-defined input 'x' of 'main' may be non-uniform
+  switch (i32(x)) @diagnostic(warning, derivative_uniformity) {
+              ^
+
+@fragment
+fn main(@location(0) x : f32) {
+  switch(i32(x)) @diagnostic(warning, derivative_uniformity) {
+    default: {
+      _ = dpdx(1.0);
+    }
+  }
+}