[HLSL][RootSignature] Implement Lexing of DescriptorTables

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