implement parsing of the flags

Author: inbelic

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -54,6 +54,15 @@ bool RootSignatureParser::parse() {
                               /*param of=*/TokenKind::kw_RootSignature);
 }
 
+template <typename FlagType>
+static FlagType maybeOrFlag(std::optional<FlagType> Flags, FlagType Flag) {
+  if (!Flags.has_value())
+    return Flag;
+
+  return static_cast<FlagType>(llvm::to_underlying(Flags.value()) |
+                               llvm::to_underlying(Flag));
+}
+
 std::optional<RootFlags> RootSignatureParser::parseRootFlags() {
   assert(CurToken.TokKind == TokenKind::kw_RootFlags &&
          "Expects to only be invoked starting at given keyword");
@@ -62,7 +71,36 @@ std::optional<RootFlags> RootSignatureParser::parseRootFlags() {
                            CurToken.TokKind))
     return std::nullopt;
 
-  RootFlags Flags = RootFlags::None;
+  std::optional<RootFlags> Flags = RootFlags::None;
+
+  // Handle the edge-case of '0' to specify no flags set
+  if (tryConsumeExpectedToken(TokenKind::int_literal)) {
+    if (!verifyZeroFlag()) {
+      getDiags().Report(CurToken.TokLoc, diag::err_expected) << "'0'";
+      return std::nullopt;
+    }
+  } else {
+    // Otherwise, parse as many flags as possible
+    TokenKind Expected[] = {
+  #define ROOT_FLAG_ENUM(NAME, LIT) TokenKind::en_##NAME,
+  #include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+    };
+
+    do {
+      if (tryConsumeExpectedToken(Expected)) {
+        switch (CurToken.TokKind) {
+  #define ROOT_FLAG_ENUM(NAME, LIT) \
+    case TokenKind::en_##NAME:                                                   \
+      Flags =                                                                    \
+          maybeOrFlag<RootFlags>(Flags, RootFlags::NAME);  \
+      break;
+  #include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+        default:
+          llvm_unreachable("Switch for consumed enum token was not provided");
+        }
+      }
+    } while (tryConsumeExpectedToken(TokenKind::pu_or));
+  }
 
   if (consumeExpectedToken(TokenKind::pu_r_paren,
                            diag::err_hlsl_unexpected_end_of_params,
@@ -492,15 +530,6 @@ RootSignatureParser::parseShaderVisibility() {
   return std::nullopt;
 }
 
-template <typename FlagType>
-static FlagType maybeOrFlag(std::optional<FlagType> Flags, FlagType Flag) {
-  if (!Flags.has_value())
-    return Flag;
-
-  return static_cast<FlagType>(llvm::to_underlying(Flags.value()) |
-                               llvm::to_underlying(Flag));
-}
-
 std::optional<llvm::hlsl::rootsig::DescriptorRangeFlags>
 RootSignatureParser::parseDescriptorRangeFlags() {
   assert(CurToken.TokKind == TokenKind::pu_equal &&

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -296,7 +296,22 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseRootConsantsTest) {
 
 TEST_F(ParseHLSLRootSignatureTest, ValidParseRootFlagsTest) {
   const llvm::StringLiteral Source = R"cc(
-    RootFlags()
+    RootFlags(),
+    RootFlags(0),
+    RootFlags(
+      deny_domain_shader_root_access |
+      deny_pixel_shader_root_access |
+      local_root_signature |
+      cbv_srv_uav_heap_directly_indexed |
+      deny_amplification_shader_root_access |
+      deny_geometry_shader_root_access |
+      deny_hull_shader_root_access |
+      deny_mesh_shader_root_access |
+      allow_stream_output |
+      sampler_heap_directly_indexed |
+      allow_input_assembler_input_layout |
+      deny_vertex_shader_root_access
+    )
   )cc";
 
   TrivialModuleLoader ModLoader;
@@ -312,12 +327,20 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseRootFlagsTest) {
 
   ASSERT_FALSE(Parser.parse());
 
-  ASSERT_EQ(Elements.size(), 1u);
+  ASSERT_EQ(Elements.size(), 3u);
 
   RootElement Elem = Elements[0];
   ASSERT_TRUE(std::holds_alternative<RootFlags>(Elem));
   ASSERT_EQ(std::get<RootFlags>(Elem), RootFlags::None);
 
+  Elem = Elements[1];
+  ASSERT_TRUE(std::holds_alternative<RootFlags>(Elem));
+  ASSERT_EQ(std::get<RootFlags>(Elem), RootFlags::None);
+
+  Elem = Elements[2];
+  ASSERT_TRUE(std::holds_alternative<RootFlags>(Elem));
+  ASSERT_EQ(std::get<RootFlags>(Elem), RootFlags::ValidFlags);
+
   ASSERT_TRUE(Consumer->isSatisfied());
 }