add support for parsing registers

Author: inbelic

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -115,6 +115,9 @@ class RootSignatureParser {
   bool ParseDescriptorTable();
   bool ParseDescriptorTableClause();
 
+  // Helper dispatch method
+  bool ParseRegister(llvm::hlsl::rootsig::Register *Reg);
+
   /// 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

@@ -313,13 +313,43 @@ bool RootSignatureParser::ParseDescriptorTableClause() {
   if (ConsumeExpectedToken(TokenKind::pu_l_paren))
     return true;
 
+  // Consume mandatory Register paramater
+  if (ConsumeExpectedToken(
+          {TokenKind::bReg, TokenKind::tReg, TokenKind::uReg, TokenKind::sReg}))
+    return true;
+  if (ParseRegister(&Clause.Register))
+    return true;
+
   if (ConsumeExpectedToken(TokenKind::pu_r_paren))
     return true;
 
   Elements.push_back(Clause);
   return false;
 }
 
+bool RootSignatureParser::ParseRegister(Register *Register) {
+  switch (CurToken.Kind) {
+  case TokenKind::bReg:
+    Register->ViewType = RegisterType::BReg;
+    break;
+  case TokenKind::tReg:
+    Register->ViewType = RegisterType::TReg;
+    break;
+  case TokenKind::uReg:
+    Register->ViewType = RegisterType::UReg;
+    break;
+  case TokenKind::sReg:
+    Register->ViewType = RegisterType::SReg;
+    break;
+  default:
+    llvm_unreachable("Switch for an expected token was not provided");
+  }
+
+  Register->Number = CurToken.NumLiteral.getInt().getExtValue();
+
+  return false;
+}
+
 // 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,10 +338,10 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseEmptyTest) {
 TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   const llvm::StringLiteral Source = R"cc(
     DescriptorTable(
-      CBV(),
-      SRV(),
-      Sampler(),
-      UAV()
+      CBV(b0),
+      SRV(t42),
+      Sampler(s987),
+      UAV(u987234)
     ),
     DescriptorTable()
   )cc";
@@ -361,18 +361,33 @@ TEST_F(ParseHLSLRootSignatureTest, ValidParseDTClausesTest) {
   RootElement Elem = Elements[0];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::CBuffer);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.ViewType,
+            RegisterType::BReg);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number, (uint32_t)0);
 
   Elem = Elements[1];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::SRV);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.ViewType,
+            RegisterType::TReg);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number,
+            (uint32_t)42);
 
   Elem = Elements[2];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::Sampler);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.ViewType,
+            RegisterType::SReg);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number,
+            (uint32_t)987);
 
   Elem = Elements[3];
   ASSERT_TRUE(std::holds_alternative<DescriptorTableClause>(Elem));
   ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Type, ClauseType::UAV);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.ViewType,
+            RegisterType::UReg);
+  ASSERT_EQ(std::get<DescriptorTableClause>(Elem).Register.Number,
+            (uint32_t)987234);
 
   Elem = Elements[4];
   ASSERT_TRUE(std::holds_alternative<DescriptorTable>(Elem));

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

@@ -23,6 +23,13 @@ namespace rootsig {
 
 // Definitions of the in-memory data layout structures
 
+// Models the different registers: bReg | tReg | uReg | sReg
+enum class RegisterType { BReg, TReg, UReg, SReg };
+struct Register {
+  RegisterType ViewType;
+  uint32_t Number;
+};
+
 // Models the end of a descriptor table and stores its visibility
 struct DescriptorTable {
   uint32_t NumClauses = 0; // The number of clauses in the table
@@ -32,6 +39,7 @@ struct DescriptorTable {
 using ClauseType = llvm::dxil::ResourceClass;
 struct DescriptorTableClause {
   ClauseType Type;
+  Register Register;
 };
 
 // Models RootElement : DescriptorTable | DescriptorTableClause