Add support for lexing registers

Author: inbelic

clang/include/clang/Basic/DiagnosticLexKinds.td

@@ -1025,6 +1025,7 @@ let CategoryName = "Root Signature Lexical Issue" in {
   def err_hlsl_number_literal_overflow :
     Error<"provided %select{unsigned integer|signed integer}0 literal '%1' that overflows the maximum of 32 bits">;
   def err_hlsl_invalid_token: Error<"unable to lex a valid Root Signature token">;
+  def err_hlsl_invalid_register_literal: Error<"expected unsigned integer literal as part of register">;
 }
 
 }

clang/include/clang/Parse/HLSLRootSignatureTokenKinds.def

@@ -25,6 +25,12 @@
 TOK(invalid)
 TOK(int_literal)
 
+// Register Tokens:
+TOK(bReg)
+TOK(tReg)
+TOK(uReg)
+TOK(sReg)
+
 // Punctuators:
 PUNCTUATOR(l_paren, '(')
 PUNCTUATOR(r_paren, ')')

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -100,6 +100,53 @@ bool RootSignatureLexer::LexToken(RootSignatureToken &Result) {
   if (isdigit(C) || C == '-' || C == '+')
     return LexNumber(Result);
 
+  // All following tokens require at least one additional character
+  if (Buffer.size() <= 1) {
+    PP.getDiagnostics().Report(Result.TokLoc, diag::err_hlsl_invalid_token);
+    return true;
+  }
+
+  // Peek at the next character to deteremine token type
+  char NextC = Buffer[1];
+
+  // Registers: [tsub][0-9+]
+  if ((C == 't' || C == 's' || C == 'u' || C == 'b') && isdigit(NextC)) {
+    AdvanceBuffer();
+
+    if (LexNumber(Result))
+      return true; // Error parsing number which is already reported
+
+    // Lex number could also parse a float so ensure it was an unsigned int
+    if (Result.Kind != TokenKind::int_literal ||
+        Result.NumLiteral.getInt().isSigned()) {
+      // Return invalid number literal for register error
+      PP.getDiagnostics().Report(Result.TokLoc,
+                                 diag::err_hlsl_invalid_register_literal);
+      return true;
+    }
+
+    // Convert character to the register type.
+    // This is done after LexNumber to override the TokenKind
+    switch (C) {
+    case 'b':
+      Result.Kind = TokenKind::bReg;
+      break;
+    case 't':
+      Result.Kind = TokenKind::tReg;
+      break;
+    case 'u':
+      Result.Kind = TokenKind::uReg;
+      break;
+    case 's':
+      Result.Kind = TokenKind::sReg;
+      break;
+    default:
+      llvm_unreachable("Switch for an expected token was not provided");
+      return true;
+    }
+    return false;
+  }
+
   // Unable to match on any token type
   PP.getDiagnostics().Report(Result.TokLoc, diag::err_hlsl_invalid_token);
   return true;

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -144,6 +144,8 @@ TEST_F(ParseHLSLRootSignatureTest, ValidLexAllTokensTest) {
   const llvm::StringLiteral Source = R"cc(
     42
 
+    b0 t43 u987 s234
+
     (),|=
   )cc";
 
@@ -213,6 +215,27 @@ TEST_F(ParseHLSLRootSignatureTest, InvalidLexEmptyNumberTest) {
   ASSERT_TRUE(Consumer->IsSatisfied());
 }
 
+TEST_F(ParseHLSLRootSignatureTest, InvalidLexRegNumberTest) {
+  // This test will check that the lexing fails due to no integer being provided
+  const llvm::StringLiteral Source = R"cc(
+    b32.4
+  )cc";
+
+  TrivialModuleLoader ModLoader;
+  auto PP = CreatePP(Source, ModLoader);
+  auto TokLoc = SourceLocation();
+
+  // Test correct diagnostic produced
+  Consumer->SetExpected(diag::err_hlsl_invalid_register_literal);
+
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
+
+  SmallVector<hlsl::RootSignatureToken> Tokens;
+  ASSERT_TRUE(Lexer.Lex(Tokens));
+  // FIXME(#120472): This should be TRUE once we can lex a floating
+  ASSERT_FALSE(Consumer->IsSatisfied());
+}
+
 TEST_F(ParseHLSLRootSignatureTest, InvalidLexIdentifierTest) {
   // This test will check that the lexing fails due to no valid token
   const llvm::StringLiteral Source = R"cc(