add support for optional parameters

- use numDescriptors as an example

Author: inbelic

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -116,7 +116,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(llvm::hlsl::rootsig::ParamType Ref);
+
+  // Parse as many optional parameters as possible in any order
+  bool ParseOptionalParams(
+      llvm::SmallDenseMap<TokenKind, llvm::hlsl::rootsig::ParamType> &RefMap);
+
+  // Common parsing helpers
   bool ParseRegister(llvm::hlsl::rootsig::Register *Reg);
+  bool ParseUInt(uint32_t *X);
 
   /// Invoke the lexer to consume a token and update CurToken with the result
   ///

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -320,13 +320,73 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
   if (ParseRegister(&Clause.Register))
     return true;
 
+  // Parse optional paramaters
+  llvm::SmallDenseMap<TokenKind, 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, 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 = CurToken.Kind;
+    if (Seen.contains(ParamKind)) {
+      return true;
+    }
+    Seen.insert(ParamKind);
+
+    if (ParseParam(RefMap[ParamKind]))
+      return true;
+  }
+
+  return false;
+}
+
+bool RootSignatureParser::ParseUInt(uint32_t *X) {
+  // Treat a postively signed integer as though it is unsigned to match DXC
+  TryConsumeExpectedToken(TokenKind::pu_plus);
+  if (ConsumeExpectedToken(TokenKind::int_literal))
+    return true;
+
+  *X = CurToken.NumLiteral.getInt().getExtValue();
+  return false;
+}
+
 bool RootSignatureParser::ParseRegister(Register *Register) {
   switch (CurToken.Kind) {
   case TokenKind::bReg:

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -339,7 +339,7 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
     DescriptorTable(
       CBV(b0),
-      SRV(t42),
+      SRV(t42, numDescriptors = +4),
       Sampler(s987),
       UAV(u987234)
     ),
@@ -364,6 +364,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));
@@ -372,6 +373,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));
@@ -380,6 +382,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));
@@ -388,6 +391,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,16 @@ 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 rootsig
 } // namespace hlsl
 } // namespace llvm