review: remove api for pre-allocating tokens

- we want to prevent the pre-allocation of all the tokens and rather let
the parser allocate the tokens as it goes along
- this allows for a quicker cycle to the user if there is a parsing
error as we don't need to lex all of the tokens before trying to parse

- as such, this change provides a new api for interacting with the lexer
through ConsumeToken and PeekNextToken as opposed to generated the
entire list of lexed tokens

Author: inbelic

clang/include/clang/Basic/DiagnosticLexKinds.td

@@ -1025,6 +1025,7 @@ let CategoryName = "Root Signature Lexical Issue" in {
     Error<"integer literal is too large to be represented in any %select{signed |}0 integer type">;
   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">;
+  def err_hlsl_rootsig_unexpected_eos : Error<"unexpected end to token stream">;
 }
 
 }

clang/include/clang/Parse/HLSLRootSignatureTokenKinds.def

@@ -51,6 +51,7 @@
 
 // General Tokens:
 TOK(invalid)
+TOK(end_of_stream)
 TOK(int_literal)
 
 // Register Tokens:

clang/include/clang/Parse/ParseHLSLRootSignature.h

@@ -39,7 +39,9 @@ struct RootSignatureToken {
   APValue NumLiteral = APValue();
 
   // Constructors
+  RootSignatureToken() : TokLoc(SourceLocation()) {}
   RootSignatureToken(clang::SourceLocation TokLoc) : TokLoc(TokLoc) {}
+  RootSignatureToken(enum Kind Kind, clang::SourceLocation TokLoc) : Kind(Kind), TokLoc(TokLoc) {}
 };
 using TokenKind = enum RootSignatureToken::Kind;
 
@@ -49,28 +51,38 @@ class RootSignatureLexer {
                      clang::Preprocessor &PP)
       : Buffer(Signature), SourceLoc(SourceLoc), PP(PP) {}
 
-  // Consumes the internal buffer as a list of tokens and will emplace them
-  // onto the given tokens.
-  //
-  // It will consume until it successfully reaches the end of the buffer,
-  // or, until the first error is encountered. The return value denotes if
-  // there was a failure.
-  bool Lex(SmallVector<RootSignatureToken> &Tokens);
+  /// Updates CurToken to the next token. Either it will take the previously
+  /// lexed NextToken, or it will lex a token.
+  ///
+  /// The return value denotes if there was a failure.
+  bool ConsumeToken();
+
+  /// Returns the token that comes after CurToken or std::nullopt if an
+  /// error is encountered during lexing of the next token.
+  std::optional<RootSignatureToken> PeekNextToken();
+
+  RootSignatureToken GetCurToken() { return CurToken; }
+
+  /// Check if we have reached the end of input
+  bool EndOfBuffer() {
+    AdvanceBuffer(Buffer.take_while(isspace).size());
+    return Buffer.empty();
+  }
 
 private:
   // Internal buffer to iterate over
   StringRef Buffer;
 
+  // Current Token state
+  RootSignatureToken CurToken;
+  std::optional<RootSignatureToken> NextToken = std::nullopt;
+
   // Passed down parameters from Sema
   clang::SourceLocation SourceLoc;
   clang::Preprocessor &PP;
 
   bool LexNumber(RootSignatureToken &Result);
-
-  // Consumes the internal buffer for a single token.
-  //
-  // The return value denotes if there was a failure.
-  bool LexToken(RootSignatureToken &Token);
+  bool LexToken(RootSignatureToken &Result);
 
   // Advance the buffer by the specified number of characters. Updates the
   // SourceLocation appropriately.

clang/lib/Parse/ParseHLSLRootSignature.cpp

@@ -60,27 +60,13 @@ bool RootSignatureLexer::LexNumber(RootSignatureToken &Result) {
   return false;
 }
 
-bool RootSignatureLexer::Lex(SmallVector<RootSignatureToken> &Tokens) {
+bool RootSignatureLexer::LexToken(RootSignatureToken &Result) {
   // Discard any leading whitespace
   AdvanceBuffer(Buffer.take_while(isspace).size());
 
-  while (!Buffer.empty()) {
-    // Record where this token is in the text for usage in parser diagnostics
-    RootSignatureToken Result(SourceLoc);
-    if (LexToken(Result))
-      return true;
-
-    // Successfully Lexed the token so we can store it
-    Tokens.push_back(Result);
+  // Record where this token is in the text for usage in parser diagnostics
+  Result = RootSignatureToken(SourceLoc);
 
-    // Discard any trailing whitespace
-    AdvanceBuffer(Buffer.take_while(isspace).size());
-  }
-
-  return false;
-}
-
-bool RootSignatureLexer::LexToken(RootSignatureToken &Result) {
   char C = Buffer.front();
 
   // Punctuators
@@ -168,5 +154,40 @@ bool RootSignatureLexer::LexToken(RootSignatureToken &Result) {
   return false;
 }
 
+bool RootSignatureLexer::ConsumeToken() {
+  // If we previously peeked then just copy the value over
+  if (NextToken && NextToken->Kind != TokenKind::end_of_stream) {
+    CurToken = *NextToken;
+    NextToken = std::nullopt;
+    return false;
+  }
+
+  // This will be implicity be true if NextToken->Kind == end_of_stream
+  if (EndOfBuffer()) {
+    // Report unexpected end of tokens error
+    PP.getDiagnostics().Report(SourceLoc, diag::err_hlsl_rootsig_unexpected_eos);
+    return true;
+  }
+
+  return LexToken(CurToken);
+}
+
+std::optional<RootSignatureToken> RootSignatureLexer::PeekNextToken() {
+  // Already peeked from the current token
+  if (NextToken.has_value())
+    return NextToken;
+
+  if (EndOfBuffer()) {
+    // We have reached the end of the stream, but only error on consume
+    return RootSignatureToken(TokenKind::end_of_stream, SourceLoc);
+  }
+
+  RootSignatureToken Result;
+  if (LexToken(Result))
+    return std::nullopt;
+  NextToken = Result;
+  return Result;
+}
+
 } // namespace hlsl
 } // namespace clang

clang/unittests/Parse/ParseHLSLRootSignatureTest.cpp

@@ -91,12 +91,19 @@ class ParseHLSLRootSignatureTest : public ::testing::Test {
     return PP;
   }
 
-  void CheckTokens(SmallVector<hlsl::RootSignatureToken> &Computed,
+  void CheckTokens(hlsl::RootSignatureLexer &Lexer,
+                   SmallVector<hlsl::RootSignatureToken> &Computed,
                    SmallVector<hlsl::TokenKind> &Expected) {
-    ASSERT_EQ(Computed.size(), Expected.size());
     for (unsigned I = 0, E = Expected.size(); I != E; ++I) {
-      ASSERT_EQ(Computed[I].Kind, Expected[I]);
+      if (Expected[I] == hlsl::TokenKind::invalid ||
+          Expected[I] == hlsl::TokenKind::end_of_stream)
+        continue;
+      ASSERT_FALSE(Lexer.ConsumeToken());
+      hlsl::RootSignatureToken Result = Lexer.GetCurToken();
+      ASSERT_EQ(Result.Kind, Expected[I]);
+      Computed.push_back(Result);
     }
+    ASSERT_TRUE(Lexer.EndOfBuffer());
   }
 
   FileSystemOptions FileMgrOpts;
@@ -128,16 +135,14 @@ TEST_F(ParseHLSLRootSignatureTest, ValidLexNumbersTest) {
   hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
 
   SmallVector<hlsl::RootSignatureToken> Tokens;
-  ASSERT_FALSE(Lexer.Lex(Tokens));
-  ASSERT_TRUE(Consumer->IsSatisfied());
-
   SmallVector<hlsl::TokenKind> Expected = {
       hlsl::TokenKind::int_literal,
       hlsl::TokenKind::int_literal,
       hlsl::TokenKind::int_literal,
       hlsl::TokenKind::int_literal,
   };
-  CheckTokens(Tokens, Expected);
+  CheckTokens(Lexer, Tokens, Expected);
+  ASSERT_TRUE(Consumer->IsSatisfied());
 
   // Sample negative int
   hlsl::RootSignatureToken IntToken = Tokens[0];
@@ -204,23 +209,77 @@ TEST_F(ParseHLSLRootSignatureTest, ValidLexAllTokensTest) {
 
   hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
 
-  SmallVector<hlsl::RootSignatureToken> Tokens = {
-      hlsl::RootSignatureToken(
-          SourceLocation()) // invalid token for completeness
-  };
-  ASSERT_FALSE(Lexer.Lex(Tokens));
-  ASSERT_TRUE(Consumer->IsSatisfied());
-
+  SmallVector<hlsl::RootSignatureToken> Tokens;
   SmallVector<hlsl::TokenKind> Expected = {
 #define TOK(NAME) hlsl::TokenKind::NAME,
 #include "clang/Parse/HLSLRootSignatureTokenKinds.def"
   };
 
-  CheckTokens(Tokens, Expected);
+  CheckTokens(Lexer, Tokens, Expected);
+  ASSERT_TRUE(Consumer->IsSatisfied());
+}
+
+TEST_F(ParseHLSLRootSignatureTest, ValidLexPeekTest) {
+  // 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(
+    )1
+  )cc";
+
+  TrivialModuleLoader ModLoader;
+  auto PP = CreatePP(Source, ModLoader);
+  auto TokLoc = SourceLocation();
+
+  // Test no diagnostics produced
+  Consumer->SetNoDiag();
+
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
+  // Test basic peek
+  auto Res = Lexer.PeekNextToken();
+  ASSERT_TRUE(Res.has_value());
+  ASSERT_EQ(Res->Kind, hlsl::TokenKind::pu_r_paren);
+
+  // Ensure it doesn't peek past one element
+  Res = Lexer.PeekNextToken();
+  ASSERT_TRUE(Res.has_value());
+  ASSERT_EQ(Res->Kind, hlsl::TokenKind::pu_r_paren);
+
+  ASSERT_FALSE(Lexer.ConsumeToken());
+
+  // Invoke after reseting the NextToken
+  Res = Lexer.PeekNextToken();
+  ASSERT_TRUE(Res.has_value());
+  ASSERT_EQ(Res->Kind, hlsl::TokenKind::int_literal);
+
+  // Ensure we can still consume the second token
+  ASSERT_FALSE(Lexer.ConsumeToken());
+
+  // Ensure no error raised when peeking past end of stream
+  Res = Lexer.PeekNextToken();
+  ASSERT_TRUE(Res.has_value());
+  ASSERT_EQ(Res->Kind, hlsl::TokenKind::end_of_stream);
+
+  ASSERT_TRUE(Consumer->IsSatisfied());
 }
 
 // Invalid Lexing Tests
 
+TEST_F(ParseHLSLRootSignatureTest, InvalidParseUnexpectedEOSTest) {
+  const llvm::StringLiteral Source = R"cc()cc";
+
+  TrivialModuleLoader ModLoader;
+  auto PP = CreatePP(Source, ModLoader);
+  auto TokLoc = SourceLocation();
+
+  hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
+
+  // Test correct diagnostic produced
+  Consumer->SetExpected(diag::err_hlsl_rootsig_unexpected_eos);
+  ASSERT_TRUE(Lexer.ConsumeToken());
+
+  ASSERT_TRUE(Consumer->IsSatisfied());
+}
+
 TEST_F(ParseHLSLRootSignatureTest, InvalidLexOverflowedNumberTest) {
   // This test will check that the lexing fails due to an integer overflow
   const llvm::StringLiteral Source = R"cc(
@@ -236,8 +295,7 @@ TEST_F(ParseHLSLRootSignatureTest, InvalidLexOverflowedNumberTest) {
 
   hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
 
-  SmallVector<hlsl::RootSignatureToken> Tokens;
-  ASSERT_TRUE(Lexer.Lex(Tokens));
+  ASSERT_TRUE(Lexer.ConsumeToken());
   ASSERT_TRUE(Consumer->IsSatisfied());
 }
 
@@ -256,8 +314,7 @@ TEST_F(ParseHLSLRootSignatureTest, InvalidLexEmptyNumberTest) {
 
   hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
 
-  SmallVector<hlsl::RootSignatureToken> Tokens;
-  ASSERT_TRUE(Lexer.Lex(Tokens));
+  ASSERT_TRUE(Lexer.ConsumeToken());
   ASSERT_TRUE(Consumer->IsSatisfied());
 }
 
@@ -276,9 +333,8 @@ TEST_F(ParseHLSLRootSignatureTest, InvalidLexRegNumberTest) {
 
   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_TRUE(Lexer.ConsumeToken());
+  // FIXME(#126565): This should be TRUE once we can lex a float
   ASSERT_FALSE(Consumer->IsSatisfied());
 }
 
@@ -297,8 +353,7 @@ TEST_F(ParseHLSLRootSignatureTest, InvalidLexIdentifierTest) {
 
   hlsl::RootSignatureLexer Lexer(Source, TokLoc, *PP);
 
-  SmallVector<hlsl::RootSignatureToken> Tokens;
-  ASSERT_TRUE(Lexer.Lex(Tokens));
+  ASSERT_TRUE(Lexer.ConsumeToken());
   ASSERT_TRUE(Consumer->IsSatisfied());
 }