[HLSL][RootSignature] Implement Lexing of DescriptorTables (#122981)

For the sake of scope, we will let the lexing of floating literals be
deferred until needed for Static Samplers. Other than that this pr
should allow us to simply define new enumerations/keywords in
`RootSignatureTokenKinds.def` for when they are used in the parser. We
could have defined all of these keywords here, but for the sake of
correctness in review we will let them be split up.

- Define `RootSignatureLexer` and provide a public `LexToken` method for
external use
- Define the file `RootSignatureTokenKinds` to define required tokens
and allow for future custom keywords/enums
- Implement the internal methods required to parse the different types
of tokens (integers, flag enums, puncuators...)
- Add test harness for unit testing and the respective unit tests for
lexing the tokens

Resolves #126563

---------

Co-authored-by: Chris B <beanz@abolishcrlf.org>

clang/include/clang/Lex/HLSLRootSignatureTokenKinds.def

@@ -0,0 +1,123 @@
+//===--- HLSLRootSignature.def - Tokens and Enum Database -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the TokenKinds used in the Root Signature DSL. This
+// includes keywords, enums and a small subset of punctuators. Users of this
+// file must optionally #define the TOK, KEYWORD, ENUM or specific ENUM macros
+// to make use of this file.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TOK
+#define TOK(X)
+#endif
+#ifndef PUNCTUATOR
+#define PUNCTUATOR(X,Y) TOK(pu_ ## X)
+#endif
+#ifndef KEYWORD
+#define KEYWORD(X) TOK(kw_ ## X)
+#endif
+#ifndef ENUM
+#define ENUM(NAME, LIT) TOK(en_ ## NAME)
+#endif
+
+// Defines the various types of enum
+#ifndef DESCRIPTOR_RANGE_OFFSET_ENUM
+#define DESCRIPTOR_RANGE_OFFSET_ENUM(NAME, LIT) ENUM(NAME, LIT)
+#endif
+#ifndef ROOT_DESCRIPTOR_FLAG_ENUM
+#define ROOT_DESCRIPTOR_FLAG_ENUM(NAME, LIT) ENUM(NAME, LIT)
+#endif
+// Note: ON denotes that the flag is unique from the above Root Descriptor
+//  Flags. This is required to avoid token kind enum conflicts.
+#ifndef DESCRIPTOR_RANGE_FLAG_ENUM_OFF
+#define DESCRIPTOR_RANGE_FLAG_ENUM_OFF(NAME, LIT)
+#endif
+#ifndef DESCRIPTOR_RANGE_FLAG_ENUM_ON
+#define DESCRIPTOR_RANGE_FLAG_ENUM_ON(NAME, LIT) ENUM(NAME, LIT)
+#endif
+#ifndef DESCRIPTOR_RANGE_FLAG_ENUM
+#define DESCRIPTOR_RANGE_FLAG_ENUM(NAME, LIT, ON) DESCRIPTOR_RANGE_FLAG_ENUM_##ON(NAME, LIT)
+#endif
+#ifndef SHADER_VISIBILITY_ENUM
+#define SHADER_VISIBILITY_ENUM(NAME, LIT) ENUM(NAME, LIT)
+#endif
+
+// General Tokens:
+TOK(invalid)
+TOK(end_of_stream)
+TOK(int_literal)
+
+// Register Tokens:
+TOK(bReg)
+TOK(tReg)
+TOK(uReg)
+TOK(sReg)
+
+// Punctuators:
+PUNCTUATOR(l_paren, '(')
+PUNCTUATOR(r_paren, ')')
+PUNCTUATOR(comma,   ',')
+PUNCTUATOR(or,      '|')
+PUNCTUATOR(equal,   '=')
+PUNCTUATOR(plus,    '+')
+PUNCTUATOR(minus,   '-')
+
+// RootElement Keywords:
+KEYWORD(DescriptorTable)
+
+// DescriptorTable Keywords:
+KEYWORD(CBV)
+KEYWORD(SRV)
+KEYWORD(UAV)
+KEYWORD(Sampler)
+
+// General Parameter Keywords:
+KEYWORD(space)
+KEYWORD(visibility)
+KEYWORD(flags)
+
+// View Parameter Keywords:
+KEYWORD(numDescriptors)
+KEYWORD(offset)
+
+// Descriptor Range Offset Enum:
+DESCRIPTOR_RANGE_OFFSET_ENUM(DescriptorRangeOffsetAppend, "DESCRIPTOR_RANGE_OFFSET_APPEND")
+
+// Root Descriptor Flag Enums:
+ROOT_DESCRIPTOR_FLAG_ENUM(DataVolatile, "DATA_VOLATILE")
+ROOT_DESCRIPTOR_FLAG_ENUM(DataStaticWhileSetAtExecute, "DATA_STATIC_WHILE_SET_AT_EXECUTE")
+ROOT_DESCRIPTOR_FLAG_ENUM(DataStatic, "DATA_STATIC")
+
+// Descriptor Range Flag Enums:
+DESCRIPTOR_RANGE_FLAG_ENUM(DescriptorsVolatile, "DESCRIPTORS_VOLATILE", ON)
+DESCRIPTOR_RANGE_FLAG_ENUM(DataVolatile, "DATA_VOLATILE", OFF)
+DESCRIPTOR_RANGE_FLAG_ENUM(DataStaticWhileSetAtExecute, "DATA_STATIC_WHILE_SET_AT_EXECUTE", OFF)
+DESCRIPTOR_RANGE_FLAG_ENUM(DataStatic, "DATA_STATIC", OFF)
+DESCRIPTOR_RANGE_FLAG_ENUM(DescriptorsStaticKeepingBufferBoundsChecks, "DESCRIPTORS_STATIC_KEEPING_BUFFER_BOUNDS_CHECKS", ON)
+
+// Shader Visibiliy Enums:
+SHADER_VISIBILITY_ENUM(All, "SHADER_VISIBILITY_ALL")
+SHADER_VISIBILITY_ENUM(Vertex, "SHADER_VISIBILITY_VERTEX")
+SHADER_VISIBILITY_ENUM(Hull, "SHADER_VISIBILITY_HULL")
+SHADER_VISIBILITY_ENUM(Domain, "SHADER_VISIBILITY_DOMAIN")
+SHADER_VISIBILITY_ENUM(Geometry, "SHADER_VISIBILITY_GEOMETRY")
+SHADER_VISIBILITY_ENUM(Pixel, "SHADER_VISIBILITY_PIXEL")
+SHADER_VISIBILITY_ENUM(Amplification, "SHADER_VISIBILITY_AMPLIFICATION")
+SHADER_VISIBILITY_ENUM(Mesh, "SHADER_VISIBILITY_MESH")
+
+#undef SHADER_VISIBILITY_ENUM
+#undef DESCRIPTOR_RANGE_FLAG_ENUM
+#undef DESCRIPTOR_RANGE_FLAG_ENUM_OFF
+#undef DESCRIPTOR_RANGE_FLAG_ENUM_ON
+#undef ROOT_DESCRIPTOR_FLAG_ENUM
+#undef DESCRIPTOR_RANGE_OFFSET_ENUM
+#undef ENUM
+#undef KEYWORD
+#undef PUNCTUATOR
+#undef TOK

clang/include/clang/Lex/LexHLSLRootSignature.h

@@ -0,0 +1,86 @@
+//===--- LexHLSLRootSignature.h ---------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines the LexHLSLRootSignature interface.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_LEX_LEXHLSLROOTSIGNATURE_H
+#define LLVM_CLANG_LEX_LEXHLSLROOTSIGNATURE_H
+
+#include "clang/Basic/SourceLocation.h"
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSwitch.h"
+
+namespace clang {
+namespace hlsl {
+
+struct RootSignatureToken {
+  enum Kind {
+#define TOK(X) X,
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  };
+
+  Kind Kind = Kind::invalid;
+
+  // Retain the SouceLocation of the token for diagnostics
+  clang::SourceLocation TokLoc;
+
+  // Retain spelling of an numeric constant to be parsed later
+  StringRef NumSpelling;
+
+  // Constructors
+  RootSignatureToken(clang::SourceLocation TokLoc) : TokLoc(TokLoc) {}
+  RootSignatureToken(enum Kind Kind, clang::SourceLocation TokLoc)
+      : Kind(Kind), TokLoc(TokLoc) {}
+};
+using TokenKind = enum RootSignatureToken::Kind;
+
+class RootSignatureLexer {
+public:
+  RootSignatureLexer(StringRef Signature, clang::SourceLocation SourceLoc)
+      : Buffer(Signature), SourceLoc(SourceLoc) {}
+
+  /// Consumes and returns the next token.
+  RootSignatureToken ConsumeToken();
+
+  /// Returns the token that proceeds CurToken
+  RootSignatureToken PeekNextToken();
+
+  bool EndOfBuffer() {
+    AdvanceBuffer(Buffer.take_while(isspace).size());
+    return Buffer.empty();
+  }
+
+private:
+  // Internal buffer to iterate over
+  StringRef Buffer;
+
+  // Current peek state
+  std::optional<RootSignatureToken> NextToken = std::nullopt;
+
+  // Passed down parameters from Sema
+  clang::SourceLocation SourceLoc;
+
+  /// Consumes the buffer and returns the lexed token.
+  RootSignatureToken LexToken();
+
+  /// Advance the buffer by the specified number of characters.
+  /// Updates the SourceLocation appropriately.
+  void AdvanceBuffer(unsigned NumCharacters = 1) {
+    Buffer = Buffer.drop_front(NumCharacters);
+    SourceLoc = SourceLoc.getLocWithOffset(NumCharacters);
+  }
+};
+
+} // namespace hlsl
+} // namespace clang
+
+#endif // LLVM_CLANG_LEX_PARSEHLSLROOTSIGNATURE_H

clang/lib/Lex/CMakeLists.txt

@@ -11,6 +11,7 @@ add_clang_library(clangLex
   HeaderSearch.cpp
   InitHeaderSearch.cpp
   Lexer.cpp
+  LexHLSLRootSignature.cpp
   LiteralSupport.cpp
   MacroArgs.cpp
   MacroInfo.cpp

clang/lib/Lex/LexHLSLRootSignature.cpp

@@ -0,0 +1,120 @@
+#include "clang/Lex/LexHLSLRootSignature.h"
+
+namespace clang {
+namespace hlsl {
+
+// Lexer Definitions
+
+static bool IsNumberChar(char C) {
+  // TODO(#126565): extend for float support exponents
+  return isdigit(C); // integer support
+}
+
+RootSignatureToken RootSignatureLexer::LexToken() {
+  // Discard any leading whitespace
+  AdvanceBuffer(Buffer.take_while(isspace).size());
+
+  if (EndOfBuffer())
+    return RootSignatureToken(TokenKind::end_of_stream, SourceLoc);
+
+  // Record where this token is in the text for usage in parser diagnostics
+  RootSignatureToken Result(SourceLoc);
+
+  char C = Buffer.front();
+
+  // Punctuators
+  switch (C) {
+#define PUNCTUATOR(X, Y)                                                       \
+  case Y: {                                                                    \
+    Result.Kind = TokenKind::pu_##X;                                           \
+    AdvanceBuffer();                                                           \
+    return Result;                                                             \
+  }
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  default:
+    break;
+  }
+
+  // Integer literal
+  if (isdigit(C)) {
+    Result.Kind = TokenKind::int_literal;
+    Result.NumSpelling = Buffer.take_while(IsNumberChar);
+    AdvanceBuffer(Result.NumSpelling.size());
+    return Result;
+  }
+
+  // All following tokens require at least one additional character
+  if (Buffer.size() <= 1) {
+    Result = RootSignatureToken(TokenKind::invalid, SourceLoc);
+    return Result;
+  }
+
+  // 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)) {
+    // Convert character to the register type.
+    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");
+    }
+
+    AdvanceBuffer();
+
+    // Lex the integer literal
+    Result.NumSpelling = Buffer.take_while(IsNumberChar);
+    AdvanceBuffer(Result.NumSpelling.size());
+
+    return Result;
+  }
+
+  // Keywords and Enums:
+  StringRef TokSpelling =
+      Buffer.take_while([](char C) { return isalnum(C) || C == '_'; });
+
+  // Define a large string switch statement for all the keywords and enums
+  auto Switch = llvm::StringSwitch<TokenKind>(TokSpelling);
+#define KEYWORD(NAME) Switch.Case(#NAME, TokenKind::kw_##NAME);
+#define ENUM(NAME, LIT) Switch.CaseLower(LIT, TokenKind::en_##NAME);
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+
+  // Then attempt to retreive a string from it
+  Result.Kind = Switch.Default(TokenKind::invalid);
+  AdvanceBuffer(TokSpelling.size());
+  return Result;
+}
+
+RootSignatureToken RootSignatureLexer::ConsumeToken() {
+  // If we previously peeked then just return the previous value over
+  if (NextToken && NextToken->Kind != TokenKind::end_of_stream) {
+    RootSignatureToken Result = *NextToken;
+    NextToken = std::nullopt;
+    return Result;
+  }
+  return LexToken();
+}
+
+RootSignatureToken RootSignatureLexer::PeekNextToken() {
+  // Already peeked from the current token
+  if (NextToken)
+    return *NextToken;
+
+  NextToken = LexToken();
+  return *NextToken;
+}
+
+} // namespace hlsl
+} // namespace clang

clang/unittests/Lex/CMakeLists.txt

@@ -7,6 +7,7 @@ add_clang_unittest(LexTests
   HeaderMapTest.cpp
   HeaderSearchTest.cpp
   LexerTest.cpp
+  LexHLSLRootSignatureTest.cpp
   ModuleDeclStateTest.cpp
   PPCallbacksTest.cpp
   PPConditionalDirectiveRecordTest.cpp

clang/unittests/Lex/LexHLSLRootSignatureTest.cpp

@@ -0,0 +1,155 @@
+//=== LexHLSLRootSignatureTest.cpp - Lex Root Signature tests -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Lex/LexHLSLRootSignature.h"
+#include "gtest/gtest.h"
+
+using namespace clang;
+
+namespace {
+
+// The test fixture.
+class LexHLSLRootSignatureTest : public ::testing::Test {
+protected:
+  LexHLSLRootSignatureTest() {}
+
+  void CheckTokens(hlsl::RootSignatureLexer &Lexer,
+                   SmallVector<hlsl::RootSignatureToken> &Computed,
+                   SmallVector<hlsl::TokenKind> &Expected) {
+    for (unsigned I = 0, E = Expected.size(); I != E; ++I) {
+      // Skip these to help with the macro generated test
+      if (Expected[I] == hlsl::TokenKind::invalid ||
+          Expected[I] == hlsl::TokenKind::end_of_stream)
+        continue;
+      hlsl::RootSignatureToken Result = Lexer.ConsumeToken();
+      ASSERT_EQ(Result.Kind, Expected[I]);
+      Computed.push_back(Result);
+    }
+    hlsl::RootSignatureToken EndOfStream = Lexer.ConsumeToken();
+    ASSERT_EQ(EndOfStream.Kind, hlsl::TokenKind::end_of_stream);
+    ASSERT_TRUE(Lexer.EndOfBuffer());
+  }
+};
+
+// Lexing Tests
+
+TEST_F(LexHLSLRootSignatureTest, ValidLexNumbersTest) {
+  // This test will check that we can lex different number tokens
+  const llvm::StringLiteral Source = R"cc(
+    -42 42 +42 +2147483648
+  )cc";
+
+  auto TokLoc = SourceLocation();
+
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc);
+
+  SmallVector<hlsl::RootSignatureToken> Tokens;
+  SmallVector<hlsl::TokenKind> Expected = {
+      hlsl::TokenKind::pu_minus,    hlsl::TokenKind::int_literal,
+      hlsl::TokenKind::int_literal, hlsl::TokenKind::pu_plus,
+      hlsl::TokenKind::int_literal, hlsl::TokenKind::pu_plus,
+      hlsl::TokenKind::int_literal,
+  };
+  CheckTokens(Lexer, Tokens, Expected);
+
+  // Sample negative: int component
+  hlsl::RootSignatureToken IntToken = Tokens[1];
+  ASSERT_EQ(IntToken.NumSpelling, "42");
+
+  // Sample unsigned int
+  IntToken = Tokens[2];
+  ASSERT_EQ(IntToken.NumSpelling, "42");
+
+  // Sample positive: int component
+  IntToken = Tokens[4];
+  ASSERT_EQ(IntToken.NumSpelling, "42");
+
+  // Sample positive int that would overflow the signed representation but
+  // is treated as an unsigned integer instead
+  IntToken = Tokens[6];
+  ASSERT_EQ(IntToken.NumSpelling, "2147483648");
+}
+
+TEST_F(LexHLSLRootSignatureTest, ValidLexAllTokensTest) {
+  // This test will check that we can lex all defined tokens as defined in
+  // HLSLRootSignatureTokenKinds.def, plus some additional integer variations
+  const llvm::StringLiteral Source = R"cc(
+    42
+
+    b0 t43 u987 s234
+
+    (),|=+-
+
+    DescriptorTable
+
+    CBV SRV UAV Sampler
+    space visibility flags
+    numDescriptors offset
+
+    DESCRIPTOR_RANGE_OFFSET_APPEND
+
+    DATA_VOLATILE
+    DATA_STATIC_WHILE_SET_AT_EXECUTE
+    DATA_STATIC
+    DESCRIPTORS_VOLATILE
+    DESCRIPTORS_STATIC_KEEPING_BUFFER_BOUNDS_CHECKS
+
+    shader_visibility_all
+    shader_visibility_vertex
+    shader_visibility_hull
+    shader_visibility_domain
+    shader_visibility_geometry
+    shader_visibility_pixel
+    shader_visibility_amplification
+    shader_visibility_mesh
+  )cc";
+  auto TokLoc = SourceLocation();
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc);
+
+  SmallVector<hlsl::RootSignatureToken> Tokens;
+  SmallVector<hlsl::TokenKind> Expected = {
+#define TOK(NAME) hlsl::TokenKind::NAME,
+#include "clang/Lex/HLSLRootSignatureTokenKinds.def"
+  };
+
+  CheckTokens(Lexer, Tokens, Expected);
+}
+
+TEST_F(LexHLSLRootSignatureTest, ValidLexPeekTest) {
+  // This test will check that we the peek api is correctly used
+  const llvm::StringLiteral Source = R"cc(
+    )1
+  )cc";
+  auto TokLoc = SourceLocation();
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc);
+
+  // Test basic peek
+  hlsl::RootSignatureToken Res = Lexer.PeekNextToken();
+  ASSERT_EQ(Res.Kind, hlsl::TokenKind::pu_r_paren);
+
+  // Ensure it doesn't peek past one element
+  Res = Lexer.PeekNextToken();
+  ASSERT_EQ(Res.Kind, hlsl::TokenKind::pu_r_paren);
+
+  Res = Lexer.ConsumeToken();
+  ASSERT_EQ(Res.Kind, hlsl::TokenKind::pu_r_paren);
+
+  // Invoke after reseting the NextToken
+  Res = Lexer.PeekNextToken();
+  ASSERT_EQ(Res.Kind, hlsl::TokenKind::int_literal);
+
+  // Ensure we can still consume the second token
+  Res = Lexer.ConsumeToken();
+  ASSERT_EQ(Res.Kind, hlsl::TokenKind::int_literal);
+
+  // Ensure end of stream token
+  Res = Lexer.PeekNextToken();
+  ASSERT_EQ(Res.Kind, hlsl::TokenKind::end_of_stream);
+}
+
+} // anonymous namespace