add support for optional parameters

- use numDescriptors as an example

Author: inbelic

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -106,7 +106,22 @@ class RootSignatureParser {
   bool ParseDescriptorTableClause();
 
   // Helper dispatch method
+  //
+  // These will switch on the Variant kind to dispatch to the respective Parse
+  // method and store the parsed value back into Ref.
+  //
+  // It is helpful to have a generalized dispatch method so that when we need
+  // to parse multiple optional parameters in any order, we can invoke this
+  // method
+  bool ParseParam(rs::ParamType Ref);
+
+  // Parse as many optional parameters as possible in any order
+  bool
+  ParseOptionalParams(llvm::SmallDenseMap<TokenKind, rs::ParamType> RefMap);
+
+  // Common parsing helpers
   bool ParseRegister(rs::Register *Reg);
+  bool ParseUInt(uint32_t *X);
 
   // Increment the token iterator if we have not reached the end.
   // Return value denotes if we were already at the last token.

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -325,13 +325,71 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
   if (ParseRegister(&Clause.Register))
     return true;
 
+  // Parse optional paramaters
+  llvm::SmallDenseMap<TokenKind, rs::ParamType> RefMap = {
+      {TokenKind::kw_numDescriptors, &Clause.NumDescriptors},
+  };
+  if (ParseOptionalParams({RefMap}))
+    return true;
+
   if (ConsumeExpectedToken(TokenKind::pu_r_paren))
     return true;
 
   Elements.push_back(Clause);
   return false;
 }
 
+// Helper struct so that we can use the overloaded notation of std::visit
+template <class... Ts> struct OverloadedMethods : Ts... {
+  using Ts::operator()...;
+};
+template <class... Ts> OverloadedMethods(Ts...) -> OverloadedMethods<Ts...>;
+
+bool RootSignatureParser::ParseParam(ParamType Ref) {
+  if (ConsumeExpectedToken(TokenKind::pu_equal))
+    return true;
+
+  bool Error;
+  std::visit(OverloadedMethods{[&](uint32_t *X) { Error = ParseUInt(X); },
+  }, Ref);
+
+  return Error;
+}
+
+bool RootSignatureParser::ParseOptionalParams(
+    llvm::SmallDenseMap<TokenKind, rs::ParamType> RefMap) {
+  SmallVector<TokenKind> ParamKeywords;
+  for (auto RefPair : RefMap)
+    ParamKeywords.push_back(RefPair.first);
+
+  // Keep track of which keywords have been seen to report duplicates
+  llvm::SmallDenseSet<TokenKind> Seen;
+
+  while (!TryConsumeExpectedToken(TokenKind::pu_comma)) {
+    if (ConsumeExpectedToken(ParamKeywords))
+      return true;
+
+    TokenKind ParamKind = CurTok->Kind;
+    if (Seen.contains(ParamKind)) {
+      return true;
+    }
+    Seen.insert(ParamKind);
+
+    if (ParseParam(RefMap[ParamKind]))
+      return true;
+  }
+
+  return false;
+}
+
+bool RootSignatureParser::ParseUInt(uint32_t *X) {
+  if (ConsumeExpectedToken(TokenKind::int_literal))
+    return true;
+
+  *X = CurTok->NumLiteral.getInt().getExtValue();
+  return false;
+}
+
 bool RootSignatureParser::ParseRegister(Register *Register) {
   switch (CurTok->Kind) {
   case TokenKind::bReg:

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -309,7 +309,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
     DescriptorTable(
       CBV(b0),
-      SRV(t42),
+      SRV(t42, numDescriptors = 4),
       Sampler(s987),
       UAV(u987234)
     ),
@@ -338,6 +338,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.ViewType,
             RegisterType::BReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number, (uint32_t)0);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)1);
 
   Elem = Elements[1];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -346,6 +347,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
             RegisterType::TReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number,
             (uint32_t)42);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)4);
 
   Elem = Elements[2];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -354,6 +356,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
             RegisterType::SReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number,
             (uint32_t)987);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)1);
 
   Elem = Elements[3];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
@@ -362,6 +365,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
             RegisterType::UReg);
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number,
             (uint32_t)987234);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).NumDescriptors, (uint32_t)1);
 
   Elem = Elements[4];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));

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

@@ -40,11 +40,15 @@ using ClauseType = llvm::dxil::ResourceClass;
 struct DescriptorTableClause {
   ClauseType Type;
   Register Register;
+  uint32_t NumDescriptors = 1;
 };
 
 // Models RootElement : DescriptorTable | DescriptorTableClause
 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 *>;
 } // namespace root_signature
 } // namespace hlsl
 } // namespace llvm