add support for shader visibility

- introduces the ParseEnum function that will parse any of the ENUM
token definitions

Author: inbelic

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -135,6 +135,12 @@ class RootSignatureParser {
   bool
   ParseDescriptorRangeOffset(llvm::hlsl::rootsig::DescriptorRangeOffset *X);
 
+  // Various flags/enum parsing helpers
+  template <bool AllowZero = false, typename EnumType>
+  bool ParseEnum(llvm::SmallDenseMap<TokenKind, EnumType> &EnumMap,
+                 EnumType *Enum);
+  bool ParseShaderVisibility(llvm::hlsl::rootsig::ShaderVisibility *Enum);
+
   /// Invoke the lexer to consume a token and update CurToken with the result
   ///
   /// Return value denotes if we were already at the last token.

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -274,8 +274,23 @@ bool RootSignatureParser::ParseDescriptorTable() {
     return false;
   }
 
+  bool SeenVisibility = false;
   // Iterate through all the defined clauses
   do {
+    // Handle the visibility parameter
+    if (!TryConsumeExpectedToken(TokenKind::kw_visibility)) {
+      if (SeenVisibility) {
+        Diags.Report(CurToken.TokLoc, diag::err_hlsl_rootsig_repeat_param)
+            << FormatTokenKinds(CurToken.Kind);
+        return true;
+      }
+      SeenVisibility = true;
+      if (ParseParam(&Table.Visibility))
+        return true;
+      continue;
+    }
+
+    // Otherwise, we expect a clause
     if (ParseDescriptorTableClause())
       return true;
     Table.NumClauses++;
@@ -351,6 +366,9 @@ bool RootSignatureParser::ParseParam(ParamType Ref) {
                                [&](DescriptorRangeOffset *X) {
                                  Error = ParseDescriptorRangeOffset(X);
                                },
+                               [&](ShaderVisibility *Enum) {
+                                 Error = ParseShaderVisibility(Enum);
+                               },
   }, Ref);
 
   return Error;
@@ -432,6 +450,39 @@ bool RootSignatureParser::ParseRegister(Register *Register) {
   return false;
 }
 
+template <typename EnumType>
+bool RootSignatureParser::ParseEnum(
+    llvm::SmallDenseMap<TokenKind, EnumType> &EnumMap, EnumType *Enum) {
+  SmallVector<TokenKind> EnumToks;
+  for (auto EnumPair : EnumMap)
+    EnumToks.push_back(EnumPair.first);
+
+  // If invoked we expect to have an enum
+  if (ConsumeExpectedToken(EnumToks))
+    return true;
+
+  // Effectively a switch statement on the token kinds
+  for (auto EnumPair : EnumMap)
+    if (CurToken.Kind == EnumPair.first) {
+      *Enum = EnumPair.second;
+      return false;
+    }
+
+  llvm_unreachable("Switch for an expected token was not provided");
+  return true;
+}
+
+bool RootSignatureParser::ParseShaderVisibility(ShaderVisibility *Enum) {
+  // Define the possible flag kinds
+  llvm::SmallDenseMap<TokenKind, ShaderVisibility> EnumMap = {
+#define SHADER_VISIBILITY_ENUM(NAME, LIT)                                      \
+  {TokenKind::en_##NAME, ShaderVisibility::NAME},
+#include "clang/Parse/HLSLRootSignatureTokenKinds.def"
+  };
+
+  return ParseEnum(EnumMap, Enum);
+}
+
 // Is given token one of the expected kinds
 static bool IsExpectedToken(TokenKind Kind, ArrayRef<TokenKind> AnyExpected) {
   for (auto Expected : AnyExpected)

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -338,6 +338,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseEmptyTest) {
 TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
     DescriptorTable(
+      visibility = SHADER_VISIBILITY_PIXEL,
       CBV(b0),
       SRV(t42, space = 3, offset = 32, numDescriptors = +4),
       Sampler(s987, space = 2, offset = DESCRIPTOR_RANGE_OFFSET_APPEND),
@@ -408,11 +409,15 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   Elem = Elements[4];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));
   ASSERT_EQ(std::get<DescriptorTable>(Elem).NumClauses, (uint32_t)4);
+  ASSERT_EQ(std::get<DescriptorTable>(Elem).Visibility,
+            ShaderVisibility::Pixel);
 
   Elem = Elements[5];
   // Test generated DescriptorTable start has correct default values
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));
   ASSERT_EQ(std::get<DescriptorTable>(Elem).NumClauses, (uint32_t)0);
+  ASSERT_EQ(std::get<DescriptorTable>(Elem).Visibility, ShaderVisibility::All);
+
   ASSERT_TRUE(Consumer->IsSatisfied());
 }
 
@@ -501,4 +506,27 @@ TEST_F(ParseHLSLRootSignatureTest, InvalidParseRepeatedParamTest) {
   ASSERT_TRUE(Consumer->IsSatisfied());
 }
 
+TEST_F(ParseHLSLRootSignatureTest, InvalidParseRepeatedVisibilityTest) {
+  const llvm::StringLiteral Source = R"cc(
+    DescriptorTable(
+      visibility = SHADER_VISIBILITY_GEOMETRY,
+      visibility = SHADER_VISIBILITY_HULL
+    )
+  )cc";
+
+  TrivialModuleLoader ModLoader;
+  auto PP = CreatePP(Source, ModLoader);
+  auto TokLoc = SourceLocation();
+
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
+  SmallVector<RootElement> Elements;
+  hlsl::RootSignatureParser Parser(Elements, Lexer, Diags);
+
+  // Test correct diagnostic produced
+  Consumer->SetExpected(diag::err_hlsl_rootsig_repeat_param);
+  ASSERT_TRUE(Parser.Parse());
+
+  ASSERT_TRUE(Consumer->IsSatisfied());
+}
+
 } // anonymous namespace

llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h

@@ -25,6 +25,17 @@ namespace rootsig {
 
 enum class DescriptorRangeOffset : uint32_t;
 
+enum class ShaderVisibility {
+  All = 0,
+  Vertex = 1,
+  Hull = 2,
+  Domain = 3,
+  Geometry = 4,
+  Pixel = 5,
+  Amplification = 6,
+  Mesh = 7,
+};
+
 // Definitions of the in-memory data layout structures
 
 // Models the different registers: bReg | tReg | uReg | sReg
@@ -36,6 +47,7 @@ struct Register {
 
 // Models the end of a descriptor table and stores its visibility
 struct DescriptorTable {
+  ShaderVisibility Visibility = ShaderVisibility::All;
   uint32_t NumClauses = 0; // The number of clauses in the table
 };
 
@@ -56,7 +68,8 @@ using RootElement = std::variant<DescriptorTable, DescriptorTableClause>;
 
 // Models a reference to all assignment parameter types that any RootElement
 // may have. Things of the form: Keyword = Param
-using ParamType = std::variant<uint32_t *, DescriptorRangeOffset *>;
+using ParamType = std::variant<uint32_t *, DescriptorRangeOffset *,
+                               ShaderVisibility *>;
 
 } // namespace rootsig
 } // namespace hlsl