add support for empty descriptor table clause

inbelic · inbelic · commit 6f57e87caecf · 2025-01-28T21:08:44.000Z
diff --git a/clang/include/clang/Parse/ParseHLSLRootSignature.h b/clang/include/clang/Parse/ParseHLSLRootSignature.h
@@ -103,15 +103,28 @@ class RootSignatureParser {
   // Root Element helpers
   bool ParseRootElement(bool First);
   bool ParseDescriptorTable();
+  bool ParseDescriptorTableClause();
 
+  // Helper dispatch method
   // Increment the token iterator if we have not reached the end.
   // Return value denotes if we were already at the last token.
   bool ConsumeNextToken();
 
+  // Attempt to retrieve the next token, if TokenKind is invalid then there was
+  // no next token.
+  RootSignatureToken PeekNextToken();
+
   // Is the current token one of the expected kinds
   bool EnsureExpectedToken(TokenKind AnyExpected);
   bool EnsureExpectedToken(ArrayRef<TokenKind> AnyExpected);
 
+  // Peek if the next token is of the expected kind.
+  //
+  // Return value denotes if it failed to match the expected kind, either it is
+  // the end of the stream or it didn't match any of the expected kinds.
+  bool PeekExpectedToken(TokenKind Expected);
+  bool PeekExpectedToken(ArrayRef<TokenKind> AnyExpected);
+
   // Consume the next token and report an error if it is not of the expected
   // kind.
   //
@@ -120,6 +133,15 @@ class RootSignatureParser {
   bool ConsumeExpectedToken(TokenKind Expected);
   bool ConsumeExpectedToken(ArrayRef<TokenKind> AnyExpected);
 
+  // Peek if the next token is of the expected kind and if it is then consume
+  // it.
+  //
+  // Return value denotes if it failed to match the expected kind, either it is
+  // the end of the stream or it didn't match any of the expected kinds. It will
+  // not report an error if there isn't a match.
+  bool TryConsumeExpectedToken(TokenKind Expected);
+  bool TryConsumeExpectedToken(ArrayRef<TokenKind> Expected);
+
 private:
   SmallVector<rs::RootElement> &Elements;
   SmallVector<RootSignatureToken>::const_iterator CurTok;
diff --git a/clang/lib/Parse/ParseHLSLRootSignature.cpp b/clang/lib/Parse/ParseHLSLRootSignature.cpp
@@ -273,12 +273,64 @@ bool RootSignatureParser::ParseDescriptorTable() {
     return true;
 
   // Empty case:
-  if (!ConsumeExpectedToken(TokenKind::pu_r_paren)) {
+  if (!TryConsumeExpectedToken(TokenKind::pu_r_paren)) {
     Elements.push_back(Table);
     return false;
   }
 
-  return true;
+  // Iterate through all the defined clauses
+  do {
+    if (ParseDescriptorTableClause())
+      return true;
+    Table.NumClauses++;
+  } while (!TryConsumeExpectedToken(TokenKind::pu_comma));
+
+  if (ConsumeExpectedToken(TokenKind::pu_r_paren))
+    return true;
+
+  Elements.push_back(Table);
+  return false;
+}
+
+bool RootSignatureParser::ParseDescriptorTableClause() {
+  if (ConsumeExpectedToken({TokenKind::kw_CBV, TokenKind::kw_SRV,
+                            TokenKind::kw_UAV, TokenKind::kw_Sampler}))
+    return true;
+
+  DescriptorTableClause Clause;
+  switch (CurTok->Kind) {
+  case TokenKind::kw_CBV:
+    Clause.Type = ClauseType::CBuffer;
+    break;
+  case TokenKind::kw_SRV:
+    Clause.Type = ClauseType::SRV;
+    break;
+  case TokenKind::kw_UAV:
+    Clause.Type = ClauseType::UAV;
+    break;
+  case TokenKind::kw_Sampler:
+    Clause.Type = ClauseType::Sampler;
+    break;
+  default:
+    llvm_unreachable("Switch for an expected token was not provided");
+    return true;
+  }
+  if (ConsumeExpectedToken(TokenKind::pu_l_paren))
+    return true;
+
+  if (ConsumeExpectedToken(TokenKind::pu_r_paren))
+    return true;
+
+  Elements.push_back(Clause);
+  return false;
+}
+
+RootSignatureToken RootSignatureParser::PeekNextToken() {
+  // Create an invalid token
+  RootSignatureToken Token = RootSignatureToken(SourceLocation());
+  if (CurTok != LastTok)
+    Token = *(CurTok + 1);
+  return Token;
 }
 
 bool RootSignatureParser::ConsumeNextToken() {
@@ -313,6 +365,19 @@ bool RootSignatureParser::EnsureExpectedToken(ArrayRef<TokenKind> AnyExpected) {
   return true;
 }
 
+bool RootSignatureParser::PeekExpectedToken(TokenKind Expected) {
+  return PeekExpectedToken(ArrayRef{Expected});
+}
+
+bool RootSignatureParser::PeekExpectedToken(ArrayRef<TokenKind> AnyExpected) {
+  RootSignatureToken Token = PeekNextToken();
+  if (Token.Kind == TokenKind::invalid)
+    return true;
+  if (IsExpectedToken(Token.Kind, AnyExpected))
+    return false;
+  return true;
+}
+
 bool RootSignatureParser::ConsumeExpectedToken(TokenKind Expected) {
   return ConsumeExpectedToken(ArrayRef{Expected});
 }
@@ -325,5 +390,16 @@ bool RootSignatureParser::ConsumeExpectedToken(
   return EnsureExpectedToken(AnyExpected);
 }
 
+bool RootSignatureParser::TryConsumeExpectedToken(TokenKind Expected) {
+  return TryConsumeExpectedToken(ArrayRef{Expected});
+}
+
+bool RootSignatureParser::TryConsumeExpectedToken(
+    ArrayRef<TokenKind> AnyExpected) {
+  if (PeekExpectedToken(AnyExpected))
+    return true;
+  return ConsumeNextToken();
+}
+
 } // namespace hlsl
 } // namespace clang
diff --git a/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp b/clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp
@@ -307,6 +307,12 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseEmptyTest) {
 
 TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
+    DescriptorTable(
+      CBV(),
+      SRV(),
+      Sampler(),
+      UAV()
+    ),
     DescriptorTable()
   )cc";
 
@@ -327,7 +333,26 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
 
   ASSERT_FALSE(Parser.Parse());
   RootElement Elem = Elements[0];
+  ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::CBuffer);
 
+  Elem = Elements[1];
+  ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::SRV);
+
+  Elem = Elements[2];
+  ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::Sampler);
+
+  Elem = Elements[3];
+  ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::UAV);
+
+  Elem = Elements[4];
+  ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));
+  ASSERT_EQ(std::get<DescriptorTable>(Elem).NumClauses, (uint32_t)4);
+
+  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);
diff --git a/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h b/llvm/include/llvm/Frontend/HLSL/HLSLRootSignature.h
@@ -14,6 +14,7 @@
 #ifndef LLVM_FRONTEND_HLSL_HLSLROOTSIGNATURE_H
 #define LLVM_FRONTEND_HLSL_HLSLROOTSIGNATURE_H
 
+#include "llvm/Support/DXILABI.h"
 #include <variant>
 
 namespace llvm {
@@ -27,8 +28,14 @@ struct DescriptorTable {
   uint32_t NumClauses = 0; // The number of clauses in the table
 };
 
-// Models RootElement : DescriptorTable
-using RootElement = std::variant<DescriptorTable>;
+// Models DTClause : CBV | SRV | UAV | Sampler, by collecting like parameters
+using ClauseType = llvm::dxil::ResourceClass;
+struct DescriptorTableClause {
+  ClauseType Type;
+};
+
+// Models RootElement : DescriptorTable | DescriptorTableClause
+using RootElement = std::variant<DescriptorTable, DescriptorTableClause>;
 
 } // namespace root_signature
 } // namespace hlsl
