[Lexer] Eliminate unnecessary calls to TriviaLexer::lexTrivia

If the lexer itself keeps track of where the first comment of a token
starts, we can avoid parsing trivia into pieces.

Author: ahoppen

include/swift/Parse/Lexer.h

@@ -129,6 +129,10 @@ class Lexer {
   /// The StringRef points into the source buffer that is currently being lexed.
   StringRef TrailingTrivia;
 
+  /// The location at which the comment of the next token starts. \c nullptr if
+  /// the next token doesn't have a comment.
+  const char *CommentStart;
+
   Lexer(const Lexer&) = delete;
   void operator=(const Lexer&) = delete;
 

lib/Parse/Lexer.cpp

@@ -283,14 +283,8 @@ void Lexer::formToken(tok Kind, const char *TokStart) {
   }
   unsigned CommentLength = 0;
   if (RetainComments == CommentRetentionMode::AttachToNextToken) {
-    auto LeadingTriviaPieces = TriviaLexer::lexTrivia(LeadingTrivia);
-    // 'CommentLength' here is the length from the *first* comment to the
-    // token text (or its backtick if exist).
-    auto Iter = llvm::find_if(LeadingTriviaPieces, [](const ParsedTriviaPiece &Piece) {
-      return isCommentTriviaKind(Piece.getKind());
-    });
-    for (auto End = LeadingTriviaPieces.end(); Iter != End; Iter++) {
-      CommentLength += Iter->getLength();
+    if (CommentStart) {
+      CommentLength = TokStart - CommentStart;
     }
   }
 
@@ -2533,6 +2527,7 @@ Token Lexer::getTokenAtLocation(const SourceManager &SM, SourceLoc Loc,
 
 StringRef Lexer::lexTrivia(bool IsForTrailingTrivia) {
   const char *AllTriviaStart = CurPtr;
+  CommentStart = nullptr;
 
 Restart:
   const char *TriviaStart = CurPtr;
@@ -2562,10 +2557,16 @@ StringRef Lexer::lexTrivia(bool IsForTrailingTrivia) {
       // Don't try to lex comments here if we are lexing comments as Tokens.
       break;
     } else if (*CurPtr == '/') {
+      if (CommentStart == nullptr) {
+        CommentStart = CurPtr - 1;
+      }
       // '// ...' comment.
       skipSlashSlashComment(/*EatNewline=*/false);
       goto Restart;
     } else if (*CurPtr == '*') {
+      if (CommentStart == nullptr) {
+        CommentStart = CurPtr - 1;
+      }
       // '/* ... */' comment.
       skipSlashStarComment();
       goto Restart;

lib/Parse/ParseExpr.cpp

@@ -2017,8 +2017,6 @@ ParserResult<Expr> Parser::parseExprStringLiteral() {
     SyntaxContext->addToken(CloseDelimiter, StringRef(), EntireTrailingTrivia);
   } else {
     // Without custom delimiter the quote owns trailing trivia.
-    auto EntireTrailingTriviaPieces =
-        TriviaLexer::lexTrivia(EntireTrailingTrivia);
     SyntaxContext->addToken(CloseQuote, StringRef(), EntireTrailingTrivia);
   }