tint: Add enable directive for extensions

In this CL the enable directive is implemented.
1. Add AST node for enable directive, assign a ExtensionKind (enum) for
each supported extension.
2. Use an unorder_set in ast::Module to record all required extensions'
kind.
3. Provide inspector methods for getting names of used extension, and
getting all used enable directives' extension names and location.
4. For different writer, the extension nodes are handled in different
ways. MSL and HLSL writers will just ignore the extension nodes, while
SPIRV and GLSL writers will emit corresponding code.
5. Implement unittests and end2end test for enable directive and
inspector, using a reserved extension name `InternalExtensionForTesting`.

Bug: tint:1472
Change-Id: I40cb4061554deb477bc2005d7e38c9718385f825
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/86623
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Zhaoming Jiang <zhaoming.jiang@intel.com>

src/tint/BUILD.gn

@@ -224,6 +224,8 @@ libtint_source_set("libtint_core_all_src") {
     "ast/discard_statement.h",
     "ast/else_statement.cc",
     "ast/else_statement.h",
+    "ast/enable.cc",
+    "ast/enable.h",
     "ast/expression.cc",
     "ast/expression.h",
     "ast/external_texture.cc",

src/tint/CMakeLists.txt

@@ -112,6 +112,8 @@ set(TINT_LIB_SRCS
   ast/discard_statement.h
   ast/else_statement.cc
   ast/else_statement.h
+  ast/enable.cc
+  ast/enable.h
   ast/expression.cc
   ast/expression.h
   ast/external_texture.cc
@@ -675,6 +677,7 @@ if(TINT_BUILD_TESTS)
     ast/depth_texture_test.cc
     ast/discard_statement_test.cc
     ast/else_statement_test.cc
+    ast/enable_test.cc
     ast/external_texture_test.cc
     ast/f32_test.cc
     ast/fallthrough_statement_test.cc
@@ -892,6 +895,7 @@ if(TINT_BUILD_TESTS)
       reader/wgsl/parser_impl_continue_stmt_test.cc
       reader/wgsl/parser_impl_continuing_stmt_test.cc
       reader/wgsl/parser_impl_depth_texture_type_test.cc
+      reader/wgsl/parser_impl_enable_directive_test.cc
       reader/wgsl/parser_impl_external_texture_type_test.cc
       reader/wgsl/parser_impl_elseif_stmt_test.cc
       reader/wgsl/parser_impl_equality_expression_test.cc
@@ -1006,6 +1010,7 @@ if(TINT_BUILD_TESTS)
       writer/wgsl/generator_impl_constructor_test.cc
       writer/wgsl/generator_impl_continue_test.cc
       writer/wgsl/generator_impl_discard_test.cc
+      writer/wgsl/generator_impl_enable_test.cc
       writer/wgsl/generator_impl_fallthrough_test.cc
       writer/wgsl/generator_impl_function_test.cc
       writer/wgsl/generator_impl_global_decl_test.cc

src/tint/ast/enable.cc

@@ -0,0 +1,58 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/ast/enable.h"
+
+#include "src/tint/program_builder.h"
+#include "src/tint/sem/variable.h"
+
+TINT_INSTANTIATE_TYPEINFO(tint::ast::Enable);
+
+namespace tint::ast {
+
+Enable::ExtensionKind Enable::NameToKind(const std::string& name) {
+  // The reserved internal extension name for testing
+  if (name == "InternalExtensionForTesting") {
+    return Enable::ExtensionKind::kInternalExtensionForTesting;
+  }
+
+  return Enable::ExtensionKind::kNotAnExtension;
+}
+
+std::string Enable::KindToName(ExtensionKind kind) {
+  switch (kind) {
+    // The reserved internal extension for testing
+    case ExtensionKind::kInternalExtensionForTesting:
+      return "InternalExtensionForTesting";
+    case ExtensionKind::kNotAnExtension:
+      // Return an empty string for kNotAnExtension
+      return {};
+      // No default case, as this switch must cover all ExtensionKind values.
+  }
+  // This return shall never get hit.
+  return {};
+}
+
+Enable::Enable(ProgramID pid, const Source& src, const std::string& ext_name)
+    : Base(pid, src), name(ext_name), kind(NameToKind(ext_name)) {}
+
+Enable::Enable(Enable&&) = default;
+
+Enable::~Enable() = default;
+
+const Enable* Enable::Clone(CloneContext* ctx) const {
+  auto src = ctx->Clone(source);
+  return ctx->dst->create<Enable>(src, name);
+}
+}  // namespace tint::ast

src/tint/ast/enable.h

@@ -0,0 +1,88 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SRC_TINT_AST_ENABLE_H_
+#define SRC_TINT_AST_ENABLE_H_
+
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include "src/tint/ast/access.h"
+#include "src/tint/ast/expression.h"
+
+namespace tint::ast {
+
+/// An instance of this class represents one extension mentioned in a
+/// "enable" derictive. Example:
+///       // Enable an extension named "f16"
+///       enable f16;
+class Enable : public Castable<Enable, Node> {
+ public:
+  ///  The enum class identifing each supported WGSL extension
+  enum class ExtensionKind {
+    /// An internal reserved extension for test, named
+    /// "InternalExtensionForTesting"
+    kInternalExtensionForTesting = -2,
+    kNotAnExtension = -1,
+  };
+
+  /// Convert a string of extension name into one of ExtensionKind enum value,
+  /// the result will be ExtensionKind::kNotAnExtension if the name is not a
+  /// known extension name. A extension node of kind kNotAnExtension must not
+  /// exist in the AST tree, and using a unknown extension name in WGSL code
+  /// should result in a shader-creation error.
+  /// @param name string of the extension name
+  /// @return the ExtensionKind enum value for the extension of given name, or
+  /// kNotAnExtension if no known extension has the given name
+  static ExtensionKind NameToKind(const std::string& name);
+
+  /// Convert the ExtensionKind enum value to corresponding extension name
+  /// string. If the given enum value is kNotAnExtension or don't have a known
+  /// name, return an empty string instead.
+  /// @param kind the ExtensionKind enum value
+  /// @return string of the extension name corresponding to the given kind, or
+  /// an empty string if the given enum value is kNotAnExtension or don't have a
+  /// known corresponding name
+  static std::string KindToName(ExtensionKind kind);
+
+  /// Create a extension
+  /// @param pid the identifier of the program that owns this node
+  /// @param src the source of this node
+  /// @param name the name of extension
+  Enable(ProgramID pid, const Source& src, const std::string& name);
+  /// Move constructor
+  Enable(Enable&&);
+
+  ~Enable() override;
+
+  /// Clones this node and all transitive child nodes using the `CloneContext`
+  /// `ctx`.
+  /// @param ctx the clone context
+  /// @return the newly cloned node
+  const Enable* Clone(CloneContext* ctx) const override;
+
+  /// The extension name
+  const std::string name;
+
+  /// The extension kind
+  const ExtensionKind kind;
+};
+
+///  A set of extension kinds
+using ExtensionSet = std::unordered_set<Enable::ExtensionKind>;
+
+}  // namespace tint::ast
+
+#endif  // SRC_TINT_AST_ENABLE_H_

src/tint/ast/enable_test.cc

@@ -0,0 +1,66 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/ast/enable.h"
+
+#include "src/tint/ast/test_helper.h"
+
+namespace tint::ast {
+namespace {
+
+using AstExtensionTest = TestHelper;
+
+TEST_F(AstExtensionTest, Creation) {
+  auto* ext = create<Enable>(
+      Source{Source::Range{Source::Location{20, 2}, Source::Location{20, 5}}},
+      "InternalExtensionForTesting");
+  EXPECT_EQ(ext->source.range.begin.line, 20u);
+  EXPECT_EQ(ext->source.range.begin.column, 2u);
+  EXPECT_EQ(ext->source.range.end.line, 20u);
+  EXPECT_EQ(ext->source.range.end.column, 5u);
+  EXPECT_EQ(ext->kind,
+            ast::Enable::ExtensionKind::kInternalExtensionForTesting);
+}
+
+TEST_F(AstExtensionTest, Creation_InvalidName) {
+  auto* ext = create<Enable>(
+      Source{Source::Range{Source::Location{20, 2}, Source::Location{20, 5}}},
+      std::string());
+  EXPECT_EQ(ext->source.range.begin.line, 20u);
+  EXPECT_EQ(ext->source.range.begin.column, 2u);
+  EXPECT_EQ(ext->source.range.end.line, 20u);
+  EXPECT_EQ(ext->source.range.end.column, 5u);
+  EXPECT_EQ(ext->kind, ast::Enable::ExtensionKind::kNotAnExtension);
+}
+
+TEST_F(AstExtensionTest, NameToKind_InvalidName) {
+  EXPECT_EQ(ast::Enable::NameToKind(std::string()),
+            ast::Enable::ExtensionKind::kNotAnExtension);
+  EXPECT_EQ(ast::Enable::NameToKind("__ImpossibleExtensionName"),
+            ast::Enable::ExtensionKind::kNotAnExtension);
+  EXPECT_EQ(ast::Enable::NameToKind("123"),
+            ast::Enable::ExtensionKind::kNotAnExtension);
+}
+
+TEST_F(AstExtensionTest, KindToName) {
+  EXPECT_EQ(ast::Enable::KindToName(
+                ast::Enable::ExtensionKind::kInternalExtensionForTesting),
+            "InternalExtensionForTesting");
+  EXPECT_EQ(
+      ast::Enable::KindToName(ast::Enable::ExtensionKind::kNotAnExtension),
+      std::string());
+}
+
+}  // namespace
+}  // namespace tint::ast

src/tint/ast/module.cc

@@ -71,11 +71,22 @@ void Module::BinGlobalDeclaration(const tint::ast::Node* decl,
         TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, var, program_id);
         global_variables_.push_back(var);
       },
+      [&](const Enable* ext) {
+        TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, ext, program_id);
+        extensions_.insert(ext->kind);
+      },
       [&](Default) {
         TINT_ICE(AST, diags) << "Unknown global declaration type";
       });
 }
 
+void Module::AddEnable(const ast::Enable* ext) {
+  TINT_ASSERT(AST, ext);
+  TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, ext, program_id);
+  global_declarations_.push_back(ext);
+  extensions_.insert(ext->kind);
+}
+
 void Module::AddGlobalVariable(const ast::Variable* var) {
   TINT_ASSERT(AST, var);
   TINT_ASSERT_PROGRAM_IDS_EQUAL_IF_VALID(AST, var, program_id);
@@ -111,6 +122,7 @@ void Module::Copy(CloneContext* ctx, const Module* src) {
   type_decls_.clear();
   functions_.clear();
   global_variables_.clear();
+  extensions_.clear();
 
   for (auto* decl : global_declarations_) {
     if (!decl) {

src/tint/ast/module.h

@@ -18,6 +18,7 @@
 #include <string>
 #include <vector>
 
+#include "src/tint/ast/enable.h"
 #include "src/tint/ast/function.h"
 #include "src/tint/ast/type.h"
 
@@ -51,6 +52,10 @@ class Module final : public Castable<Module, Node> {
     return global_declarations_;
   }
 
+  /// Add a enable directive to the Builder
+  /// @param ext the enable directive to add
+  void AddEnable(const Enable* ext);
+
   /// Add a global variable to the Builder
   /// @param var the variable to add
   void AddGlobalVariable(const Variable* var);
@@ -76,6 +81,9 @@ class Module final : public Castable<Module, Node> {
   /// @returns the global variables for the module
   VariableList& GlobalVariables() { return global_variables_; }
 
+  /// @returns the extension set for the module
+  const ExtensionSet& Extensions() const { return extensions_; }
+
   /// Adds a type declaration to the Builder.
   /// @param decl the type declaration to add
   void AddTypeDecl(const TypeDecl* decl);
@@ -116,6 +124,7 @@ class Module final : public Castable<Module, Node> {
   std::vector<const TypeDecl*> type_decls_;
   FunctionList functions_;
   VariableList global_variables_;
+  ExtensionSet extensions_;
 };
 
 }  // namespace tint::ast

src/tint/inspector/inspector.cc

@@ -568,6 +568,33 @@ uint32_t Inspector::GetWorkgroupStorageSize(const std::string& entry_point) {
   return total_size;
 }
 
+std::vector<std::string> Inspector::GetUsedExtensionNames() {
+  std::vector<std::string> result;
+
+  ast::ExtensionSet set = program_->AST().Extensions();
+  result.reserve(set.size());
+  for (auto kind : set) {
+    std::string name = ast::Enable::KindToName(kind);
+    result.push_back(name);
+  }
+
+  return result;
+}
+
+std::vector<std::pair<std::string, Source>> Inspector::GetEnableDirectives() {
+  std::vector<std::pair<std::string, Source>> result;
+
+  // Ast nodes for enable directive are stored within global declarations list
+  auto global_decls = program_->AST().GlobalDeclarations();
+  for (auto node : global_decls) {
+    if (auto ext = node->As<ast::Enable>()) {
+      result.push_back({ext->name, ext->source});
+    }
+  }
+
+  return result;
+}
+
 const ast::Function* Inspector::FindEntryPointByName(const std::string& name) {
   auto* func = program_->AST().Functions().Find(program_->Symbols().Get(name));
   if (!func) {

src/tint/inspector/inspector.h

@@ -20,6 +20,7 @@
 #include <string>
 #include <tuple>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "src/tint/inspector/entry_point.h"
@@ -143,6 +144,18 @@ class Inspector {
   /// referenced transitively by the entry point.
   uint32_t GetWorkgroupStorageSize(const std::string& entry_point);
 
+  /// @returns vector of all valid extension names used by the program. There
+  /// will be no duplicated names in the returned vector even if an extension
+  /// is enabled multiple times.
+  std::vector<std::string> GetUsedExtensionNames();
+
+  /// @returns vector of all enable directives used by the program, each
+  /// enable directive represented by a std::pair<std::string,
+  /// tint::Source::Range> for its extension name and its location of the
+  /// extension name. There may be multiple enable directives for a same
+  /// extension.
+  std::vector<std::pair<std::string, Source>> GetEnableDirectives();
+
  private:
   const Program* program_;
   diag::List diagnostics_;

src/tint/inspector/inspector_test.cc

@@ -145,6 +145,12 @@ class InspectorGetSamplerTextureUsesTest : public InspectorRunner,
 class InspectorGetWorkgroupStorageSizeTest : public InspectorBuilder,
                                              public testing::Test {};
 
+class InspectorGetUsedExtensionNamesTest : public InspectorRunner,
+                                           public testing::Test {};
+
+class InspectorGetEnableDirectivesTest : public InspectorRunner,
+                                         public testing::Test {};
+
 // This is a catch all for shaders that have demonstrated regressions/crashes in
 // the wild.
 class InspectorRegressionTest : public InspectorRunner, public testing::Test {};
@@ -3004,6 +3010,124 @@ TEST_F(InspectorGetWorkgroupStorageSizeTest, StructAlignment) {
   EXPECT_EQ(1024u, inspector.GetWorkgroupStorageSize("ep_func"));
 }
 
+// Test calling GetUsedExtensionNames on a empty shader.
+TEST_F(InspectorGetUsedExtensionNamesTest, Empty) {
+  std::string shader = "";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetUsedExtensionNames();
+  EXPECT_EQ(result.size(), 0u);
+}
+
+// Test calling GetUsedExtensionNames on a shader with no extension.
+TEST_F(InspectorGetUsedExtensionNamesTest, None) {
+  std::string shader = R"(
+@stage(fragment)
+fn main() {
+})";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetUsedExtensionNames();
+  EXPECT_EQ(result.size(), 0u);
+}
+
+// Test calling GetUsedExtensionNames on a shader with valid extension.
+TEST_F(InspectorGetUsedExtensionNamesTest, Simple) {
+  std::string shader = R"(
+enable InternalExtensionForTesting;
+
+@stage(fragment)
+fn main() {
+})";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetUsedExtensionNames();
+  EXPECT_EQ(result.size(), 1u);
+  EXPECT_EQ(result[0], "InternalExtensionForTesting");
+}
+
+// Test calling GetUsedExtensionNames on a shader with a extension enabled for
+// multiple times.
+TEST_F(InspectorGetUsedExtensionNamesTest, Duplicated) {
+  std::string shader = R"(
+enable InternalExtensionForTesting;
+enable InternalExtensionForTesting;
+
+@stage(fragment)
+fn main() {
+})";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetUsedExtensionNames();
+  EXPECT_EQ(result.size(), 1u);
+  EXPECT_EQ(result[0], "InternalExtensionForTesting");
+}
+
+// Test calling GetEnableDirectives on a empty shader.
+TEST_F(InspectorGetEnableDirectivesTest, Empty) {
+  std::string shader = "";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetEnableDirectives();
+  EXPECT_EQ(result.size(), 0u);
+}
+
+// Test calling GetEnableDirectives on a shader with no extension.
+TEST_F(InspectorGetEnableDirectivesTest, None) {
+  std::string shader = R"(
+@stage(fragment)
+fn main() {
+})";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetEnableDirectives();
+  EXPECT_EQ(result.size(), 0u);
+}
+
+// Test calling GetEnableDirectives on a shader with valid extension.
+TEST_F(InspectorGetEnableDirectivesTest, Simple) {
+  std::string shader = R"(
+enable InternalExtensionForTesting;
+
+@stage(fragment)
+fn main() {
+})";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetEnableDirectives();
+  EXPECT_EQ(result.size(), 1u);
+  EXPECT_EQ(result[0].first, "InternalExtensionForTesting");
+  EXPECT_EQ(result[0].second.range, (Source::Range{{2, 8}, {2, 35}}));
+}
+
+// Test calling GetEnableDirectives on a shader with a extension enabled for
+// multiple times.
+TEST_F(InspectorGetEnableDirectivesTest, Duplicated) {
+  std::string shader = R"(
+enable InternalExtensionForTesting;
+
+enable InternalExtensionForTesting;
+@stage(fragment)
+fn main() {
+})";
+
+  Inspector& inspector = Initialize(shader);
+
+  auto result = inspector.GetEnableDirectives();
+  EXPECT_EQ(result.size(), 2u);
+  EXPECT_EQ(result[0].first, "InternalExtensionForTesting");
+  EXPECT_EQ(result[0].second.range, (Source::Range{{2, 8}, {2, 35}}));
+  EXPECT_EQ(result[1].first, "InternalExtensionForTesting");
+  EXPECT_EQ(result[1].second.range, (Source::Range{{4, 8}, {4, 35}}));
+}
+
 // Crash was occuring in ::GenerateSamplerTargets, when
 // ::GetSamplerTextureUses was called.
 TEST_F(InspectorRegressionTest, tint967) {

src/tint/program_builder.h

@@ -38,6 +38,7 @@
 #include "src/tint/ast/depth_texture.h"
 #include "src/tint/ast/disable_validation_attribute.h"
 #include "src/tint/ast/discard_statement.h"
+#include "src/tint/ast/enable.h"
 #include "src/tint/ast/external_texture.h"
 #include "src/tint/ast/f32.h"
 #include "src/tint/ast/fallthrough_statement.h"

src/tint/reader/wgsl/lexer.cc

@@ -1010,6 +1010,8 @@ Token Lexer::check_keyword(const Source& source, std::string_view str) {
     return {Token::Type::kDefault, source, "default"};
   if (str == "else")
     return {Token::Type::kElse, source, "else"};
+  if (str == "enable")
+    return {Token::Type::kEnable, source, "enable"};
   if (str == "f32")
     return {Token::Type::kF32, source, "f32"};
   if (str == "fallthrough")

src/tint/reader/wgsl/parser_impl.cc

@@ -305,14 +305,33 @@ bool ParserImpl::Parse() {
 }
 
 // translation_unit
-//  : global_decl* EOF
+//  : enable_directive* global_decl* EOF
 void ParserImpl::translation_unit() {
+  bool after_global_decl = false;
   while (continue_parsing()) {
     auto p = peek();
     if (p.IsEof()) {
       break;
     }
-    expect_global_decl();
+
+    auto ed = enable_directive();
+    if (ed.matched) {
+      if (after_global_decl) {
+        add_error(p,
+                  "enable directives must come before all global declarations");
+      }
+    } else {
+      auto gd = global_decl();
+
+      if (gd.matched) {
+        after_global_decl = true;
+      }
+
+      if (!gd.matched && !gd.errored) {
+        add_error(p, "unexpected token");
+      }
+    }
+
     if (builder_.Diagnostics().error_count() >= max_errors_) {
       add_error(Source{{}, p.source().file},
                 "stopping after " + std::to_string(max_errors_) + " errors");
@@ -321,6 +340,58 @@ void ParserImpl::translation_unit() {
   }
 }
 
+// enable_directive
+//  : enable name SEMICLON
+Maybe<bool> ParserImpl::enable_directive() {
+  auto decl = sync(Token::Type::kSemicolon, [&]() -> Maybe<bool> {
+    if (!match(Token::Type::kEnable)) {
+      return Failure::kNoMatch;
+    }
+
+    // Match the extension name.
+    Expect<std::string> name = {""};
+    auto t = peek();
+    if (t.IsIdentifier()) {
+      synchronized_ = true;
+      next();
+      name = {t.to_str(), t.source()};
+    } else if (handle_error(t)) {
+      // The token might itself be an error.
+      return Failure::kErrored;
+    } else {
+      // Failed to match an extension name.
+      synchronized_ = false;
+      return add_error(t.source(), "invalid extension name");
+    }
+
+    if (!expect("enable directive", Token::Type::kSemicolon)) {
+      return Failure::kErrored;
+    }
+
+    if (ast::Enable::NameToKind(name.value) !=
+        ast::Enable::ExtensionKind::kNotAnExtension) {
+      const ast::Enable* extension =
+          create<ast::Enable>(name.source, name.value);
+      builder_.AST().AddEnable(extension);
+    } else {
+      // Error if an unknown extension is used
+      return add_error(name.source,
+                       "unsupported extension: '" + name.value + "'");
+    }
+
+    return true;
+  });
+
+  if (decl.errored) {
+    return Failure::kErrored;
+  }
+  if (decl.matched) {
+    return true;
+  }
+
+  return Failure::kNoMatch;
+}
+
 // global_decl
 //  : SEMICOLON
 //  | global_variable_decl SEMICLON
@@ -328,7 +399,7 @@ void ParserImpl::translation_unit() {
 //  | type_alias SEMICOLON
 //  | struct_decl
 //  | function_decl
-Expect<bool> ParserImpl::expect_global_decl() {
+Maybe<bool> ParserImpl::global_decl() {
   if (match(Token::Type::kSemicolon) || match(Token::Type::kEOF))
     return true;
 
@@ -436,9 +507,9 @@ Expect<bool> ParserImpl::expect_global_decl() {
   }
 
   // Exhausted all attempts to make sense of where we're at.
-  // Spew a generic error.
+  // Return a no-match
 
-  return add_error(t, "unexpected token");
+  return Failure::kNoMatch;
 }
 
 // global_variable_decl

src/tint/reader/wgsl/parser_impl.h

@@ -380,9 +380,12 @@ class ParserImpl {
   void deprecated(const Source& source, const std::string& msg);
   /// Parses the `translation_unit` grammar element
   void translation_unit();
+  /// Parses the `enable_directive` grammar element, erroring on parse failure.
+  /// @return true on parse success, otherwise an error or no-match.
+  Maybe<bool> enable_directive();
   /// Parses the `global_decl` grammar element, erroring on parse failure.
-  /// @return true on parse success, otherwise an error.
-  Expect<bool> expect_global_decl();
+  /// @return true on parse success, otherwise an error or no-match.
+  Maybe<bool> global_decl();
   /// Parses a `global_variable_decl` grammar element with the initial
   /// `variable_attribute_list*` provided as `attrs`
   /// @returns the variable parsed or nullptr

src/tint/reader/wgsl/parser_impl_enable_directive_test.cc

@@ -0,0 +1,177 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/reader/wgsl/parser_impl_test_helper.h"
+
+#include "src/tint/ast/enable.h"
+
+namespace tint::reader::wgsl {
+namespace {
+
+using EnableDirectiveTest = ParserImplTest;
+
+// Test a valid enable directive.
+TEST_F(EnableDirectiveTest, Valid) {
+  auto p = parser("enable InternalExtensionForTesting;");
+  p->enable_directive();
+  EXPECT_FALSE(p->has_error()) << p->error();
+  auto program = p->program();
+  auto& ast = program.AST();
+  EXPECT_EQ(ast.Extensions(),
+            ast::ExtensionSet{
+                ast::Enable::ExtensionKind::kInternalExtensionForTesting});
+  EXPECT_EQ(ast.GlobalDeclarations().size(), 1u);
+  auto node = ast.GlobalDeclarations()[0]->As<ast::Enable>();
+  EXPECT_TRUE(node != nullptr);
+  EXPECT_EQ(node->name, "InternalExtensionForTesting");
+  EXPECT_EQ(node->kind,
+            ast::Enable::ExtensionKind::kInternalExtensionForTesting);
+}
+
+// Test multiple enable directives for a same extension.
+TEST_F(EnableDirectiveTest, EnableMultipleTime) {
+  auto p = parser(R"(
+enable InternalExtensionForTesting;
+enable InternalExtensionForTesting;
+)");
+  p->translation_unit();
+  EXPECT_FALSE(p->has_error()) << p->error();
+  auto program = p->program();
+  auto& ast = program.AST();
+  EXPECT_EQ(ast.Extensions(),
+            ast::ExtensionSet{
+                ast::Enable::ExtensionKind::kInternalExtensionForTesting});
+  EXPECT_EQ(ast.GlobalDeclarations().size(), 2u);
+  auto node1 = ast.GlobalDeclarations()[0]->As<ast::Enable>();
+  EXPECT_TRUE(node1 != nullptr);
+  EXPECT_EQ(node1->name, "InternalExtensionForTesting");
+  EXPECT_EQ(node1->kind,
+            ast::Enable::ExtensionKind::kInternalExtensionForTesting);
+  auto node2 = ast.GlobalDeclarations()[1]->As<ast::Enable>();
+  EXPECT_TRUE(node2 != nullptr);
+  EXPECT_EQ(node2->name, "InternalExtensionForTesting");
+  EXPECT_EQ(node2->kind,
+            ast::Enable::ExtensionKind::kInternalExtensionForTesting);
+}
+
+// Test an unknown extension identifier.
+TEST_F(EnableDirectiveTest, InvalidIdentifier) {
+  auto p = parser("enable NotAValidExtensionName;");
+  p->enable_directive();
+  // Error when unknown extension found
+  EXPECT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:8: unsupported extension: 'NotAValidExtensionName'");
+  auto program = p->program();
+  auto& ast = program.AST();
+  EXPECT_EQ(ast.Extensions().size(), 0u);
+  EXPECT_EQ(ast.GlobalDeclarations().size(), 0u);
+}
+
+// Test an enable directive missing ending semiclon.
+TEST_F(EnableDirectiveTest, MissingEndingSemiclon) {
+  auto p = parser("enable InternalExtensionForTesting");
+  p->translation_unit();
+  EXPECT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(), "1:35: expected ';' for enable directive");
+  auto program = p->program();
+  auto& ast = program.AST();
+  EXPECT_EQ(ast.Extensions().size(), 0u);
+  EXPECT_EQ(ast.GlobalDeclarations().size(), 0u);
+}
+
+// Test using invalid tokens in an enable directive.
+TEST_F(EnableDirectiveTest, InvalidTokens) {
+  {
+    auto p = parser("enable InternalExtensionForTesting<;");
+    p->translation_unit();
+    EXPECT_TRUE(p->has_error());
+    EXPECT_EQ(p->error(), "1:35: expected ';' for enable directive");
+    auto program = p->program();
+    auto& ast = program.AST();
+    EXPECT_EQ(ast.Extensions().size(), 0u);
+    EXPECT_EQ(ast.GlobalDeclarations().size(), 0u);
+  }
+  {
+    auto p = parser("enable <InternalExtensionForTesting;");
+    p->translation_unit();
+    EXPECT_TRUE(p->has_error());
+    EXPECT_EQ(p->error(), "1:8: invalid extension name");
+    auto program = p->program();
+    auto& ast = program.AST();
+    EXPECT_EQ(ast.Extensions().size(), 0u);
+    EXPECT_EQ(ast.GlobalDeclarations().size(), 0u);
+  }
+  {
+    auto p = parser("enable =;");
+    p->translation_unit();
+    EXPECT_TRUE(p->has_error());
+    EXPECT_EQ(p->error(), "1:8: invalid extension name");
+    auto program = p->program();
+    auto& ast = program.AST();
+    EXPECT_EQ(ast.Extensions().size(), 0u);
+    EXPECT_EQ(ast.GlobalDeclarations().size(), 0u);
+  }
+  {
+    auto p = parser("enable vec4;");
+    p->translation_unit();
+    EXPECT_TRUE(p->has_error());
+    EXPECT_EQ(p->error(), "1:8: invalid extension name");
+    auto program = p->program();
+    auto& ast = program.AST();
+    EXPECT_EQ(ast.Extensions().size(), 0u);
+    EXPECT_EQ(ast.GlobalDeclarations().size(), 0u);
+  }
+}
+
+// Test an enable directive go after other global declarations.
+TEST_F(EnableDirectiveTest, FollowingOtherGlobalDecl) {
+  auto p = parser(R"(
+var<private> t: f32 = 0f;
+enable InternalExtensionForTesting;
+)");
+  p->translation_unit();
+  EXPECT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(),
+            "3:1: enable directives must come before all global declarations");
+  auto program = p->program();
+  auto& ast = program.AST();
+  // Accept the enable directive although it cause an error
+  EXPECT_EQ(ast.Extensions(),
+            ast::ExtensionSet{
+                ast::Enable::ExtensionKind::kInternalExtensionForTesting});
+  EXPECT_EQ(ast.GlobalDeclarations().size(), 2u);
+}
+
+// Test an enable directive go after an empty semiclon.
+TEST_F(EnableDirectiveTest, FollowingEmptySemiclon) {
+  auto p = parser(R"(
+;
+enable InternalExtensionForTesting;
+)");
+  p->translation_unit();
+  // An empty semiclon is treated as a global declaration
+  EXPECT_TRUE(p->has_error());
+  EXPECT_EQ(p->error(),
+            "3:1: enable directives must come before all global declarations");
+  auto program = p->program();
+  auto& ast = program.AST();
+  // Accept the enable directive although it cause an error
+  EXPECT_EQ(ast.Extensions(),
+            ast::ExtensionSet{
+                ast::Enable::ExtensionKind::kInternalExtensionForTesting});
+  EXPECT_EQ(ast.GlobalDeclarations().size(), 1u);
+}
+
+}  // namespace
+}  // namespace tint::reader::wgsl

src/tint/reader/wgsl/parser_impl_global_decl_test.cc

@@ -19,13 +19,13 @@ namespace {
 
 TEST_F(ParserImplTest, GlobalDecl_Semicolon) {
   auto p = parser(";");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 }
 
 TEST_F(ParserImplTest, GlobalDecl_GlobalVariable) {
   auto p = parser("var<private> a : vec2<i32> = vec2<i32>(1, 2);");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -37,21 +37,21 @@ TEST_F(ParserImplTest, GlobalDecl_GlobalVariable) {
 
 TEST_F(ParserImplTest, GlobalDecl_GlobalVariable_Inferred_Invalid) {
   auto p = parser("var<private> a = vec2<i32>(1, 2);");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_TRUE(p->has_error());
   EXPECT_EQ(p->error(), "1:16: expected ':' for variable declaration");
 }
 
 TEST_F(ParserImplTest, GlobalDecl_GlobalVariable_MissingSemicolon) {
   auto p = parser("var<private> a : vec2<i32>");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_TRUE(p->has_error());
   EXPECT_EQ(p->error(), "1:27: expected ';' for variable declaration");
 }
 
 TEST_F(ParserImplTest, GlobalDecl_GlobalConstant) {
   auto p = parser("let a : i32 = 2;");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -63,21 +63,21 @@ TEST_F(ParserImplTest, GlobalDecl_GlobalConstant) {
 
 TEST_F(ParserImplTest, GlobalDecl_GlobalConstant_Invalid) {
   auto p = parser("let a : vec2<i32> 1.0;");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_TRUE(p->has_error());
   EXPECT_EQ(p->error(), "1:19: expected ';' for let declaration");
 }
 
 TEST_F(ParserImplTest, GlobalDecl_GlobalConstant_MissingSemicolon) {
   auto p = parser("let a : vec2<i32> = vec2<i32>(1, 2)");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_TRUE(p->has_error());
   EXPECT_EQ(p->error(), "1:36: expected ';' for let declaration");
 }
 
 TEST_F(ParserImplTest, GlobalDecl_TypeAlias) {
   auto p = parser("type A = i32;");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -93,8 +93,8 @@ TEST_F(ParserImplTest, GlobalDecl_TypeAlias_StructIdent) {
   a : f32,
 }
 type B = A;)");
-  p->expect_global_decl();
-  p->expect_global_decl();
+  p->global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -113,14 +113,14 @@ type B = A;)");
 
 TEST_F(ParserImplTest, GlobalDecl_TypeAlias_MissingSemicolon) {
   auto p = parser("type A = i32");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_TRUE(p->has_error());
   EXPECT_EQ(p->error(), "1:13: expected ';' for type alias");
 }
 
 TEST_F(ParserImplTest, GlobalDecl_Function) {
   auto p = parser("fn main() { return; }");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -131,7 +131,7 @@ TEST_F(ParserImplTest, GlobalDecl_Function) {
 
 TEST_F(ParserImplTest, GlobalDecl_Function_WithAttribute) {
   auto p = parser("@workgroup_size(2) fn main() { return; }");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -142,14 +142,14 @@ TEST_F(ParserImplTest, GlobalDecl_Function_WithAttribute) {
 
 TEST_F(ParserImplTest, GlobalDecl_Function_Invalid) {
   auto p = parser("fn main() -> { return; }");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_TRUE(p->has_error());
   EXPECT_EQ(p->error(), "1:14: unable to determine function return type");
 }
 
 TEST_F(ParserImplTest, GlobalDecl_ParsesStruct) {
   auto p = parser("struct A { b: i32, c: f32}");
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto program = p->program();
@@ -165,10 +165,20 @@ TEST_F(ParserImplTest, GlobalDecl_ParsesStruct) {
 }
 
 TEST_F(ParserImplTest, GlobalDecl_Struct_Invalid) {
+  {
     auto p = parser("A {}");
-  p->expect_global_decl();
+    auto decl = p->global_decl();
+    // global_decl will result in a no match.
+    ASSERT_FALSE(p->has_error()) << p->error();
+    ASSERT_TRUE(!decl.matched && !decl.errored);
+  }
+  {
+    auto p = parser("A {}");
+    p->translation_unit();
+    // translation_unit will result in a general error.
     ASSERT_TRUE(p->has_error());
     EXPECT_EQ(p->error(), "1:1: unexpected token");
+  }
 }
 
 }  // namespace

src/tint/reader/wgsl/parser_impl_primary_expression_test.cc

@@ -117,7 +117,7 @@ TEST_F(ParserImplTest, PrimaryExpression_TypeDecl_StructConstructor_Empty) {
   S()
   )");
 
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto e = p->primary_expression();
@@ -141,7 +141,7 @@ TEST_F(ParserImplTest, PrimaryExpression_TypeDecl_StructConstructor_NotEmpty) {
   S(1u, 2.0)
   )");
 
-  p->expect_global_decl();
+  p->global_decl();
   ASSERT_FALSE(p->has_error()) << p->error();
 
   auto e = p->primary_expression();

src/tint/reader/wgsl/token.cc

@@ -145,6 +145,8 @@ std::string_view Token::TypeToName(Type type) {
       return "default";
     case Token::Type::kElse:
       return "else";
+    case Token::Type::kEnable:
+      return "enable";
     case Token::Type::kF32:
       return "f32";
     case Token::Type::kFallthrough:

src/tint/reader/wgsl/token.h

@@ -156,6 +156,8 @@ class Token {
     kDefault,
     /// A 'else'
     kElse,
+    /// A 'enable'
+    kEnable,
     /// A 'f32'
     kF32,
     /// A 'fallthrough'

src/tint/resolver/dependency_graph.cc

@@ -162,6 +162,9 @@ class DependencyScanner {
             TraverseExpression(var->constructor);
           }
         },
+        [&](const ast::Enable*) {
+          // Enable directives do not effect the dependency graph.
+        },
         [&](Default) { UnhandledNode(diagnostics_, global->node); });
   }
 
@@ -523,7 +526,10 @@ struct DependencyAnalysis {
   void GatherGlobals(const ast::Module& module) {
     for (auto* node : module.GlobalDeclarations()) {
       auto* global = allocator_.Create(node);
+      // Enable directives do not form a symbol. Skip them.
+      if (!node->Is<ast::Enable>()) {
         globals_.emplace(SymbolOf(node), global);
+      }
       declaration_order_.emplace_back(global);
     }
   }

src/tint/resolver/resolver.cc

@@ -121,6 +121,10 @@ bool Resolver::ResolveInternal() {
   // Process all module-scope declarations in dependency order.
   for (auto* decl : dependencies_.ordered_globals) {
     Mark(decl);
+    // Enable directives don't have sem node.
+    if (decl->Is<ast::Enable>()) {
+      continue;
+    }
     if (!Switch(
             decl,  //
             [&](const ast::TypeDecl* td) { return TypeDecl(td); },

src/tint/transform/single_entry_point.cc

@@ -93,6 +93,8 @@ void SingleEntryPoint::Run(CloneContext& ctx,
       if (sem.Get(func)->HasAncestorEntryPoint(entry_point->symbol)) {
         ctx.dst->AST().AddFunction(ctx.Clone(func));
       }
+    } else if (auto* ext = decl->As<ast::Enable>()) {
+      ctx.dst->AST().AddEnable(ctx.Clone(ext));
     } else {
       TINT_UNREACHABLE(Transform, ctx.dst->Diagnostics())
           << "unhandled global declaration: " << decl->TypeInfo().name;

src/tint/writer/glsl/generator_impl.cc

@@ -282,6 +282,11 @@ bool GeneratorImpl::Generate() {
           return false;
         }
       }
+    } else if (auto* ext = decl->As<ast::Enable>()) {
+      // Record the required extension for generating extension directive later
+      if (!RecordExtension(ext)) {
+        return false;
+      }
     } else {
       TINT_ICE(Writer, diagnostics_)
           << "unhandled module-scope declaration: " << decl->TypeInfo().name;
@@ -316,6 +321,21 @@ bool GeneratorImpl::Generate() {
   return true;
 }
 
+bool GeneratorImpl::RecordExtension(const ast::Enable*) {
+  /*
+  Deal with extension node here, recording it within the generator for
+  later emition.
+  For example:
+  ```
+    if (ext->kind == ast::Enable::ExtensionKind::kF16) {
+    require_fp16_ = true;
+    }
+  ```
+  */
+
+  return true;
+}
+
 bool GeneratorImpl::EmitIndexAccessor(
     std::ostream& out,
     const ast::IndexAccessorExpression* expr) {

src/tint/writer/glsl/generator_impl.h

@@ -83,6 +83,10 @@ class GeneratorImpl : public TextGenerator {
   /// @returns true on successful generation; false otherwise
   bool Generate();
 
+  /// Record an extension directive within the generator
+  /// @param ext the extension to record
+  /// @returns true if the extension directive was recorded successfully
+  bool RecordExtension(const ast::Enable* ext);
   /// Handles an index accessor expression
   /// @param out the output of the expression stream
   /// @param expr the expression to emit

src/tint/writer/hlsl/generator_impl.cc

@@ -288,6 +288,11 @@ bool GeneratorImpl::Generate() {
           }
           return EmitFunction(func);
         },
+        [&](const ast::Enable*) {
+          // Currently we don't have to do anything for using a extension in
+          // HLSL
+          return true;
+        },
         [&](Default) {
           TINT_ICE(Writer, diagnostics_)
               << "unhandled module-scope declaration: "

src/tint/writer/msl/generator_impl.cc

@@ -250,6 +250,10 @@ bool GeneratorImpl::Generate() {
           }
           return EmitFunction(func);
         },
+        [&](const ast::Enable*) {
+          // Do nothing for enabling extension in MSL
+          return true;
+        },
         [&](Default) {
           // These are pushed into the entry point by sanitizer transforms.
           TINT_ICE(Writer, diagnostics_)

src/tint/writer/spirv/builder.cc

@@ -257,6 +257,10 @@ bool Builder::Build() {
                     {Operand::Int(SpvAddressingModelLogical),
                      Operand::Int(SpvMemoryModelGLSL450)});
 
+  for (auto ext : builder_.AST().Extensions()) {
+    GenerateExtension(ext);
+  }
+
   for (auto* var : builder_.AST().GlobalVariables()) {
     if (!GenerateGlobalVariable(var)) {
       return false;
@@ -340,6 +344,21 @@ void Builder::push_capability(uint32_t cap) {
   }
 }
 
+bool Builder::GenerateExtension(ast::Enable::ExtensionKind) {
+  /*
+  For each supported extension, push corresponding capability into the builder.
+  For example:
+    if (kind == ast::Extension::Kind::kF16) {
+      push_capability(SpvCapabilityFloat16);
+      push_capability(SpvCapabilityUniformAndStorageBuffer16BitAccess);
+      push_capability(SpvCapabilityStorageBuffer16BitAccess);
+      push_capability(SpvCapabilityStorageInputOutput16);
+    }
+  */
+
+  return true;
+}
+
 bool Builder::GenerateLabel(uint32_t id) {
   if (!push_function_inst(spv::Op::OpLabel, {Operand::Int(id)})) {
     return false;

src/tint/writer/spirv/builder.h

@@ -225,6 +225,11 @@ class Builder {
                                    ast::InterpolationType type,
                                    ast::InterpolationSampling sampling);
 
+  /// Generates a extension for the given extension kind. Emits an error and
+  /// returns false if the extension kind is not supported.
+  /// @param kind ExtensionKind of the extension to generate
+  /// @returns true on success.
+  bool GenerateExtension(ast::Enable::ExtensionKind kind);
   /// Generates a label for the given id. Emits an error and returns false if
   /// we're currently outside a function.
   /// @param id the id to use for the label

src/tint/writer/wgsl/generator_impl.cc

@@ -63,8 +63,20 @@ GeneratorImpl::GeneratorImpl(const Program* program) : TextGenerator(program) {}
 GeneratorImpl::~GeneratorImpl() = default;
 
 bool GeneratorImpl::Generate() {
+  // Generate enable directives before any other global declarations.
+  for (auto ext : program_->AST().Extensions()) {
+    if (!EmitEnableDirective(ext)) {
+      return false;
+    }
+  }
+  if (!program_->AST().Extensions().empty()) {
+    line();
+  }
   // Generate global declarations in the order they appear in the module.
   for (auto* decl : program_->AST().GlobalDeclarations()) {
+    if (decl->Is<ast::Enable>()) {
+      continue;
+    }
     if (!Switch(
             decl,  //
             [&](const ast::TypeDecl* td) { return EmitTypeDecl(td); },
@@ -84,6 +96,16 @@ bool GeneratorImpl::Generate() {
   return true;
 }
 
+bool GeneratorImpl::EmitEnableDirective(const ast::Enable::ExtensionKind ext) {
+  auto out = line();
+  auto extension = ast::Enable::KindToName(ext);
+  if (extension == "") {
+    return false;
+  }
+  out << "enable " << extension << ";";
+  return true;
+}
+
 bool GeneratorImpl::EmitTypeDecl(const ast::TypeDecl* ty) {
   return Switch(
       ty,

src/tint/writer/wgsl/generator_impl.h

@@ -52,6 +52,10 @@ class GeneratorImpl : public TextGenerator {
   /// @returns true on successful generation; false otherwise
   bool Generate();
 
+  /// Handles generating a enable directive
+  /// @param ext the extension kind in the enable directive to generate
+  /// @returns true if the enable directive was emitted
+  bool EmitEnableDirective(const ast::Enable::ExtensionKind ext);
   /// Handles generating a declared type
   /// @param ty the declared type to generate
   /// @returns true if the declared type was emitted

src/tint/writer/wgsl/generator_impl_enable_test.cc

@@ -0,0 +1,32 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/tint/writer/wgsl/test_helper.h"
+
+namespace tint::writer::wgsl {
+namespace {
+
+using WgslGeneratorImplTest = TestHelper;
+
+TEST_F(WgslGeneratorImplTest, Emit_Enable) {
+  GeneratorImpl& gen = Build();
+
+  ASSERT_TRUE(gen.EmitEnableDirective(
+      ast::Enable::ExtensionKind::kInternalExtensionForTesting));
+  EXPECT_EQ(gen.result(), R"(enable InternalExtensionForTesting;
+)");
+}
+
+}  // namespace
+}  // namespace tint::writer::wgsl

test/tint/BUILD.gn

@@ -164,6 +164,7 @@ tint_unittests_source_set("tint_unittests_ast_src") {
     "../../src/tint/ast/depth_texture_test.cc",
     "../../src/tint/ast/discard_statement_test.cc",
     "../../src/tint/ast/else_statement_test.cc",
+    "../../src/tint/ast/enable_test.cc",
     "../../src/tint/ast/external_texture_test.cc",
     "../../src/tint/ast/f32_test.cc",
     "../../src/tint/ast/fallthrough_statement_test.cc",
@@ -486,6 +487,7 @@ tint_unittests_source_set("tint_unittests_wgsl_reader_src") {
     "../../src/tint/reader/wgsl/parser_impl_continuing_stmt_test.cc",
     "../../src/tint/reader/wgsl/parser_impl_depth_texture_type_test.cc",
     "../../src/tint/reader/wgsl/parser_impl_elseif_stmt_test.cc",
+    "../../src/tint/reader/wgsl/parser_impl_enable_directive_test.cc",
     "../../src/tint/reader/wgsl/parser_impl_equality_expression_test.cc",
     "../../src/tint/reader/wgsl/parser_impl_error_msg_test.cc",
     "../../src/tint/reader/wgsl/parser_impl_error_resync_test.cc",
@@ -564,6 +566,7 @@ tint_unittests_source_set("tint_unittests_wgsl_writer_src") {
     "../../src/tint/writer/wgsl/generator_impl_constructor_test.cc",
     "../../src/tint/writer/wgsl/generator_impl_continue_test.cc",
     "../../src/tint/writer/wgsl/generator_impl_discard_test.cc",
+    "../../src/tint/writer/wgsl/generator_impl_enable_test.cc",
     "../../src/tint/writer/wgsl/generator_impl_fallthrough_test.cc",
     "../../src/tint/writer/wgsl/generator_impl_function_test.cc",
     "../../src/tint/writer/wgsl/generator_impl_global_decl_test.cc",

test/tint/extensions/InternalExtensionForTesting/simple_with_InternalExtensionForTesting.wgsl

@@ -0,0 +1,26 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Enable a void internal extension
+enable InternalExtensionForTesting;
+
+fn bar() {
+}
+
+@stage(fragment)
+fn main() -> @location(0) vec4<f32> {
+    var a : vec2<f32> = vec2<f32>();
+    bar();
+    return vec4<f32>(0.4, 0.4, 0.8, 1.0);
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_InternalExtensionForTesting.wgsl.expected.glsl

@@ -0,0 +1,18 @@
+#version 310 es
+precision mediump float;
+
+layout(location = 0) out vec4 value;
+void bar() {
+}
+
+vec4 tint_symbol() {
+  vec2 a = vec2(0.0f, 0.0f);
+  bar();
+  return vec4(0.400000006f, 0.400000006f, 0.800000012f, 1.0f);
+}
+
+void main() {
+  vec4 inner_result = tint_symbol();
+  value = inner_result;
+  return;
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_InternalExtensionForTesting.wgsl.expected.hlsl

@@ -0,0 +1,19 @@
+void bar() {
+}
+
+struct tint_symbol {
+  float4 value : SV_Target0;
+};
+
+float4 main_inner() {
+  float2 a = float2(0.0f, 0.0f);
+  bar();
+  return float4(0.400000006f, 0.400000006f, 0.800000012f, 1.0f);
+}
+
+tint_symbol main() {
+  const float4 inner_result = main_inner();
+  tint_symbol wrapper_result = (tint_symbol)0;
+  wrapper_result.value = inner_result;
+  return wrapper_result;
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_InternalExtensionForTesting.wgsl.expected.msl

@@ -0,0 +1,23 @@
+#include <metal_stdlib>
+
+using namespace metal;
+void bar() {
+}
+
+struct tint_symbol_1 {
+  float4 value [[color(0)]];
+};
+
+float4 tint_symbol_inner() {
+  float2 a = float2();
+  bar();
+  return float4(0.400000006f, 0.400000006f, 0.800000012f, 1.0f);
+}
+
+fragment tint_symbol_1 tint_symbol() {
+  float4 const inner_result = tint_symbol_inner();
+  tint_symbol_1 wrapper_result = {};
+  wrapper_result.value = inner_result;
+  return wrapper_result;
+}
+

test/tint/extensions/InternalExtensionForTesting/simple_with_InternalExtensionForTesting.wgsl.expected.spvasm

@@ -0,0 +1,47 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 25
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main" %value
+               OpExecutionMode %main OriginUpperLeft
+               OpName %value "value"
+               OpName %bar "bar"
+               OpName %main_inner "main_inner"
+               OpName %a "a"
+               OpName %main "main"
+               OpDecorate %value Location 0
+      %float = OpTypeFloat 32
+    %v4float = OpTypeVector %float 4
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+          %5 = OpConstantNull %v4float
+      %value = OpVariable %_ptr_Output_v4float Output %5
+       %void = OpTypeVoid
+          %6 = OpTypeFunction %void
+         %10 = OpTypeFunction %v4float
+    %v2float = OpTypeVector %float 2
+         %14 = OpConstantNull %v2float
+%_ptr_Function_v2float = OpTypePointer Function %v2float
+%float_0_400000006 = OpConstant %float 0.400000006
+%float_0_800000012 = OpConstant %float 0.800000012
+    %float_1 = OpConstant %float 1
+         %21 = OpConstantComposite %v4float %float_0_400000006 %float_0_400000006 %float_0_800000012 %float_1
+        %bar = OpFunction %void None %6
+          %9 = OpLabel
+               OpReturn
+               OpFunctionEnd
+ %main_inner = OpFunction %v4float None %10
+         %12 = OpLabel
+          %a = OpVariable %_ptr_Function_v2float Function %14
+               OpStore %a %14
+         %17 = OpFunctionCall %void %bar
+               OpReturnValue %21
+               OpFunctionEnd
+       %main = OpFunction %void None %6
+         %23 = OpLabel
+         %24 = OpFunctionCall %v4float %main_inner
+               OpStore %value %24
+               OpReturn
+               OpFunctionEnd

test/tint/extensions/InternalExtensionForTesting/simple_with_InternalExtensionForTesting.wgsl.expected.wgsl

@@ -0,0 +1,11 @@
+enable InternalExtensionForTesting;
+
+fn bar() {
+}
+
+@stage(fragment)
+fn main() -> @location(0) vec4<f32> {
+  var a : vec2<f32> = vec2<f32>();
+  bar();
+  return vec4<f32>(0.400000006, 0.400000006, 0.800000012, 1.0);
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_duplicated_InternalExtensionForTesting.wgsl

@@ -0,0 +1,28 @@
+// Copyright 2022 The Tint Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Enable a void internal extension for multiple times
+enable InternalExtensionForTesting;
+enable InternalExtensionForTesting;
+enable InternalExtensionForTesting;
+
+fn bar() {
+}
+
+@stage(fragment)
+fn main() -> @location(0) vec4<f32> {
+    var a : vec2<f32> = vec2<f32>();
+    bar();
+    return vec4<f32>(0.4, 0.4, 0.8, 1.0);
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_duplicated_InternalExtensionForTesting.wgsl.expected.glsl

@@ -0,0 +1,18 @@
+#version 310 es
+precision mediump float;
+
+layout(location = 0) out vec4 value;
+void bar() {
+}
+
+vec4 tint_symbol() {
+  vec2 a = vec2(0.0f, 0.0f);
+  bar();
+  return vec4(0.400000006f, 0.400000006f, 0.800000012f, 1.0f);
+}
+
+void main() {
+  vec4 inner_result = tint_symbol();
+  value = inner_result;
+  return;
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_duplicated_InternalExtensionForTesting.wgsl.expected.hlsl

@@ -0,0 +1,19 @@
+void bar() {
+}
+
+struct tint_symbol {
+  float4 value : SV_Target0;
+};
+
+float4 main_inner() {
+  float2 a = float2(0.0f, 0.0f);
+  bar();
+  return float4(0.400000006f, 0.400000006f, 0.800000012f, 1.0f);
+}
+
+tint_symbol main() {
+  const float4 inner_result = main_inner();
+  tint_symbol wrapper_result = (tint_symbol)0;
+  wrapper_result.value = inner_result;
+  return wrapper_result;
+}

test/tint/extensions/InternalExtensionForTesting/simple_with_duplicated_InternalExtensionForTesting.wgsl.expected.msl

@@ -0,0 +1,23 @@
+#include <metal_stdlib>
+
+using namespace metal;
+void bar() {
+}
+
+struct tint_symbol_1 {
+  float4 value [[color(0)]];
+};
+
+float4 tint_symbol_inner() {
+  float2 a = float2();
+  bar();
+  return float4(0.400000006f, 0.400000006f, 0.800000012f, 1.0f);
+}
+
+fragment tint_symbol_1 tint_symbol() {
+  float4 const inner_result = tint_symbol_inner();
+  tint_symbol_1 wrapper_result = {};
+  wrapper_result.value = inner_result;
+  return wrapper_result;
+}
+

test/tint/extensions/InternalExtensionForTesting/simple_with_duplicated_InternalExtensionForTesting.wgsl.expected.spvasm

@@ -0,0 +1,47 @@
+; SPIR-V
+; Version: 1.3
+; Generator: Google Tint Compiler; 0
+; Bound: 25
+; Schema: 0
+               OpCapability Shader
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint Fragment %main "main" %value
+               OpExecutionMode %main OriginUpperLeft
+               OpName %value "value"
+               OpName %bar "bar"
+               OpName %main_inner "main_inner"
+               OpName %a "a"
+               OpName %main "main"
+               OpDecorate %value Location 0
+      %float = OpTypeFloat 32
+    %v4float = OpTypeVector %float 4
+%_ptr_Output_v4float = OpTypePointer Output %v4float
+          %5 = OpConstantNull %v4float
+      %value = OpVariable %_ptr_Output_v4float Output %5
+       %void = OpTypeVoid
+          %6 = OpTypeFunction %void
+         %10 = OpTypeFunction %v4float
+    %v2float = OpTypeVector %float 2
+         %14 = OpConstantNull %v2float
+%_ptr_Function_v2float = OpTypePointer Function %v2float
+%float_0_400000006 = OpConstant %float 0.400000006
+%float_0_800000012 = OpConstant %float 0.800000012
+    %float_1 = OpConstant %float 1
+         %21 = OpConstantComposite %v4float %float_0_400000006 %float_0_400000006 %float_0_800000012 %float_1
+        %bar = OpFunction %void None %6
+          %9 = OpLabel
+               OpReturn
+               OpFunctionEnd
+ %main_inner = OpFunction %v4float None %10
+         %12 = OpLabel
+          %a = OpVariable %_ptr_Function_v2float Function %14
+               OpStore %a %14
+         %17 = OpFunctionCall %void %bar
+               OpReturnValue %21
+               OpFunctionEnd
+       %main = OpFunction %void None %6
+         %23 = OpLabel
+         %24 = OpFunctionCall %v4float %main_inner
+               OpStore %value %24
+               OpReturn
+               OpFunctionEnd

test/tint/extensions/InternalExtensionForTesting/simple_with_duplicated_InternalExtensionForTesting.wgsl.expected.wgsl

@@ -0,0 +1,11 @@
+enable InternalExtensionForTesting;
+
+fn bar() {
+}
+
+@stage(fragment)
+fn main() -> @location(0) vec4<f32> {
+  var a : vec2<f32> = vec2<f32>();
+  bar();
+  return vec4<f32>(0.400000006, 0.400000006, 0.800000012, 1.0);
+}