tint/ast: Add global diagnostic controls to Module

Bug: tint:1809
Change-Id: I363abc60597a6c340d5a4b4d147fa7f5747c2d9b
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/117563
Reviewed-by: Ben Clayton <bclayton@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/ast/module.cc

@@ -71,6 +71,10 @@ void Module::BinGlobalDeclaration(const tint::ast::Node* decl, diag::List& diags
             TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, var, program_id);
             global_variables_.Push(var);
         },
+        [&](const DiagnosticControl* diag_control) {
+            TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, diag_control, program_id);
+            diagnostic_controls_.Push(diag_control);
+        },
         [&](const Enable* enable) {
             TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, enable, program_id);
             enables_.Push(enable);
@@ -82,6 +86,13 @@ void Module::BinGlobalDeclaration(const tint::ast::Node* decl, diag::List& diags
         [&](Default) { TINT_ICE(AST, diags) << "Unknown global declaration type"; });
 }
 
+void Module::AddDiagnosticControl(const ast::DiagnosticControl* control) {
+    TINT_ASSERT(AST, control);
+    TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, control, program_id);
+    global_declarations_.Push(control);
+    diagnostic_controls_.Push(control);
+}
+
 void Module::AddEnable(const ast::Enable* enable) {
     TINT_ASSERT(AST, enable);
     TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, enable, program_id);

src/tint/ast/module.h

@@ -18,6 +18,7 @@
 #include <string>
 
 #include "src/tint/ast/const_assert.h"
+#include "src/tint/ast/diagnostic_control.h"
 #include "src/tint/ast/enable.h"
 #include "src/tint/ast/function.h"
 #include "src/tint/ast/type.h"
@@ -92,10 +93,17 @@ class Module final : public Castable<Module, Node> {
         return out;
     }
 
+    /// Add a global diagnostic control to the module
+    /// @param control the diagnostic control to add
+    void AddDiagnosticControl(const DiagnosticControl* control);
+
     /// Add a enable directive to the module
     /// @param ext the enable directive to add
     void AddEnable(const Enable* ext);
 
+    /// @returns the global diagnostic controls for the module
+    const auto& DiagnosticControls() const { return diagnostic_controls_; }
+
     /// @returns the extension set for the module
     const auto& Enables() const { return enables_; }
 
@@ -146,6 +154,7 @@ class Module final : public Castable<Module, Node> {
     utils::Vector<const TypeDecl*, 16> type_decls_;
     FunctionList functions_;
     utils::Vector<const Variable*, 32> global_variables_;
+    utils::Vector<const DiagnosticControl*, 8> diagnostic_controls_;
     utils::Vector<const Enable*, 8> enables_;
     utils::Vector<const ConstAssert*, 8> const_asserts_;
 };

src/tint/ast/module_test.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "gmock/gmock.h"
 #include "gtest/gtest-spi.h"
 #include "src/tint/ast/test_helper.h"
 #include "src/tint/clone_context.h"
@@ -130,5 +131,29 @@ TEST_F(ModuleTest, CloneOrder) {
     ASSERT_EQ(cloned.Symbols().NameFor(decls[4]->As<ast::Alias>()->name), "inserted_before_V");
 }
 
+TEST_F(ModuleTest, Directives) {
+    auto* enable_1 = Enable(ast::Extension::kF16);
+    auto* diagnostic_1 = DiagnosticDirective(DiagnosticSeverity::kWarning, Expr("foo"));
+    auto* enable_2 = Enable(ast::Extension::kChromiumExperimentalFullPtrParameters);
+    auto* diagnostic_2 = DiagnosticDirective(DiagnosticSeverity::kOff, Expr("bar"));
+
+    this->SetResolveOnBuild(false);
+    Program program(std::move(*this));
+    EXPECT_THAT(program.AST().GlobalDeclarations(), ::testing::ContainerEq(utils::Vector{
+                                                        enable_1,
+                                                        diagnostic_1,
+                                                        enable_2,
+                                                        diagnostic_2,
+                                                    }));
+    EXPECT_THAT(program.AST().Enables(), ::testing::ContainerEq(utils::Vector{
+                                             enable_1,
+                                             enable_2,
+                                         }));
+    EXPECT_THAT(program.AST().DiagnosticControls(), ::testing::ContainerEq(utils::Vector{
+                                                        diagnostic_1,
+                                                        diagnostic_2,
+                                                    }));
+}
+
 }  // namespace
 }  // namespace tint::ast

src/tint/program_builder.h

@@ -3266,6 +3266,30 @@ class ProgramBuilder {
         return create<ast::DiagnosticControl>(source_, severity, rule_name);
     }
 
+    /// Add a global diagnostic control to the module.
+    /// @param source the source information
+    /// @param severity the diagnostic severity control
+    /// @param rule_name the diagnostic rule name
+    /// @returns the diagnostic control pointer
+    const ast::DiagnosticControl* DiagnosticDirective(const Source& source,
+                                                      ast::DiagnosticSeverity severity,
+                                                      const ast::IdentifierExpression* rule_name) {
+        auto* control = DiagnosticControl(source, severity, rule_name);
+        AST().AddDiagnosticControl(control);
+        return control;
+    }
+
+    /// Add a global diagnostic control to the module.
+    /// @param severity the diagnostic severity control
+    /// @param rule_name the diagnostic rule name
+    /// @returns the diagnostic control pointer
+    const ast::DiagnosticControl* DiagnosticDirective(ast::DiagnosticSeverity severity,
+                                                      const ast::IdentifierExpression* rule_name) {
+        auto* control = DiagnosticControl(source_, severity, rule_name);
+        AST().AddDiagnosticControl(control);
+        return control;
+    }
+
     /// Sets the current builder source to `src`
     /// @param src the Source used for future create() calls
     void SetSource(const Source& src) {