[Syntax] Restore Lexer state by restoring LeadingTrivia

Instead of re-lexing Trivias. Trivia might contain skipped garbage which
might be re-lexed as tokens.

Author: rintaro

include/swift/Parse/Lexer.h

@@ -243,18 +243,17 @@ class Lexer {
   /// there.
   State getStateForBeginningOfToken(const Token &Tok,
                                     const syntax::Trivia &LeadingTrivia = {}) const {
-    // If trivia parsing mode, start position of trivia is the position we want
-    // to restore.
-    if (TriviaRetention == TriviaRetentionMode::WithTrivia)
-      return State(Tok.getLoc().getAdvancedLoc(-LeadingTrivia.getTextLength()));
 
     // If the token has a comment attached to it, rewind to before the comment,
     // not just the start of the token.  This ensures that we will re-lex and
     // reattach the comment to the token if rewound to this state.
     SourceLoc TokStart = Tok.getCommentStart();
     if (TokStart.isInvalid())
       TokStart = Tok.getLoc();
-    return getStateForBeginningOfTokenLoc(TokStart);
+    auto S = getStateForBeginningOfTokenLoc(TokStart);
+    if (TriviaRetention == TriviaRetentionMode::WithTrivia)
+      S.LeadingTrivia = LeadingTrivia.Pieces;
+    return S;
   }
 
   State getStateForEndOfTokenLoc(SourceLoc Loc) const {
@@ -273,7 +272,13 @@ class Lexer {
     // Don't reemit diagnostics while readvancing the lexer.
     llvm::SaveAndRestore<DiagnosticEngine*>
       D(Diags, enableDiagnostics ? Diags : nullptr);
+
     lexImpl();
+
+    // Restore Trivia.
+    if (TriviaRetention == TriviaRetentionMode::WithTrivia)
+      if (auto &LTrivia = S.LeadingTrivia)
+        LeadingTrivia = std::move(*LTrivia);
   }
 
   /// \brief Restore the lexer state to a given state that is located before

include/swift/Parse/LexerState.h

@@ -17,10 +17,12 @@
 #ifndef SWIFT_LEXERSTATE_H
 #define SWIFT_LEXERSTATE_H
 
+#include "llvm/ADT/Optional.h"
 #include "swift/Basic/SourceLoc.h"
+#include "swift/Syntax/Trivia.h"
 
 namespace swift {
-  class Lexer;
+class Lexer;
 
 /// \brief Lexer state can be saved/restored to/from objects of this class.
 class LexerState {
@@ -37,6 +39,7 @@ class LexerState {
 private:
   explicit LexerState(SourceLoc Loc) : Loc(Loc) {}
   SourceLoc Loc;
+  llvm::Optional<syntax::TriviaList> LeadingTrivia;
   friend class Lexer;
 };