[Syntax] Fix roundtrip test

include/swift/Parse/CodeCompletionCallbacks.h

@@ -34,7 +34,7 @@ class CodeCompletionCallbacks {
 protected:
   Parser &P;
   ASTContext &Context;
-  Parser::ParserPosition ExprBeginPosition;
+  ParserPosition ExprBeginPosition;
 
   /// The declaration parsed during delayed parsing that was caused by code
   /// completion. This declaration contained the code completion token.
@@ -65,7 +65,7 @@ class CodeCompletionCallbacks {
     return CompleteExprSelectorContext != ObjCSelectorContext::None;
   }
 
-  void setExprBeginning(Parser::ParserPosition PP) {
+  void setExprBeginning(ParserPosition PP) {
     ExprBeginPosition = PP;
   }
 

include/swift/Parse/Lexer.h

@@ -20,6 +20,7 @@
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/Basic/SourceLoc.h"
 #include "swift/Basic/SourceManager.h"
+#include "swift/Parse/LexerState.h"
 #include "swift/Parse/Token.h"
 #include "swift/Syntax/References.h"
 #include "swift/Syntax/Trivia.h"
@@ -59,13 +60,15 @@ enum class ConflictMarkerKind {
   /// separated by 4 "="s, and terminated by 4 "<"s.
   Perforce
 };
-  
+
 class Lexer {
   const LangOptions &LangOpts;
   const SourceManager &SourceMgr;
   DiagnosticEngine *Diags;
   const unsigned BufferID;
 
+  using State = LexerState;
+
   /// Pointer to the first character of the buffer, even in a lexer that
   /// scans a subrange of the buffer.
   const char *BufferStart;
@@ -128,28 +131,6 @@ class Lexer {
   /// `TriviaRetentionMode::WithTrivia`.
   syntax::TriviaList TrailingTrivia;
 
-public:
-  /// \brief Lexer state can be saved/restored to/from objects of this class.
-  class State {
-  public:
-    State() {}
-
-    bool isValid() const {
-      return Loc.isValid();
-    }
-
-    State advance(unsigned Offset) const {
-      assert(isValid());
-      return State(Loc.getAdvancedLoc(Offset));
-    }
-
-  private:
-    explicit State(SourceLoc Loc) : Loc(Loc) {}
-    SourceLoc Loc;
-    friend class Lexer;
-  };
-
-private:
   Lexer(const Lexer&) = delete;
   void operator=(const Lexer&) = delete;
 
@@ -262,24 +243,27 @@ 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 {
     return State(getLocForEndOfToken(SourceMgr, Loc));
   }
 
+  bool isStateForCurrentBuffer(LexerState State) const {
+    return SourceMgr.findBufferContainingLoc(State.Loc) == getBufferID();
+  }
+
   /// \brief Restore the lexer state to a given one, that can be located either
   /// before or after the current position.
   void restoreState(State S, bool enableDiagnostics = false) {
@@ -288,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

@@ -0,0 +1,48 @@
+//===--- LexerState.h - Lexer State -----------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines the LexerState object.
+//
+//===----------------------------------------------------------------------===//
+
+#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;
+
+/// \brief Lexer state can be saved/restored to/from objects of this class.
+class LexerState {
+public:
+  LexerState() {}
+
+  bool isValid() const { return Loc.isValid(); }
+
+  LexerState advance(unsigned Offset) const {
+    assert(isValid());
+    return LexerState(Loc.getAdvancedLoc(Offset));
+  }
+
+private:
+  explicit LexerState(SourceLoc Loc) : Loc(Loc) {}
+  SourceLoc Loc;
+  llvm::Optional<syntax::TriviaList> LeadingTrivia;
+  friend class Lexer;
+};
+
+} // end namespace swift
+
+#endif

include/swift/Parse/Parser.h

@@ -30,6 +30,7 @@
 #include "swift/Parse/LocalContext.h"
 #include "swift/Parse/PersistentParserState.h"
 #include "swift/Parse/Token.h"
+#include "swift/Parse/ParserPosition.h"
 #include "swift/Parse/ParserResult.h"
 #include "swift/Parse/SyntaxParserResult.h"
 #include "swift/Syntax/SyntaxParsingContext.h"
@@ -358,24 +359,6 @@ class Parser {
   //===--------------------------------------------------------------------===//
   // Routines to save and restore parser state.
 
-  class ParserPosition {
-  public:
-    ParserPosition() = default;
-    ParserPosition &operator=(const ParserPosition &) = default;
-
-    bool isValid() const {
-      return LS.isValid();
-    }
-
-  private:
-    ParserPosition(Lexer::State LS, SourceLoc PreviousLoc):
-        LS(LS), PreviousLoc(PreviousLoc)
-    {}
-    Lexer::State LS;
-    SourceLoc PreviousLoc;
-    friend class Parser;
-  };
-
   ParserPosition getParserPosition() {
     return ParserPosition(L->getStateForBeginningOfToken(Tok, LeadingTrivia),
                           PreviousLoc);

include/swift/Parse/ParserPosition.h

@@ -0,0 +1,41 @@
+//===--- ParserPosition.h - Parser Position ---------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// Parser position where Parser can jump to.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_PARSE_PARSERPOSITION_H
+#define SWIFT_PARSE_PARSERPOSITION_H
+
+#include "swift/Basic/SourceLoc.h"
+#include "swift/Parse/LexerState.h"
+
+namespace swift {
+
+class ParserPosition {
+  LexerState LS;
+  SourceLoc PreviousLoc;
+  friend class Parser;
+
+  ParserPosition(LexerState LS, SourceLoc PreviousLoc)
+      : LS(LS), PreviousLoc(PreviousLoc) {}
+public:
+  ParserPosition() = default;
+  ParserPosition &operator=(const ParserPosition &) = default;
+
+  bool isValid() const { return LS.isValid(); }
+};
+
+} // end namespace swift
+
+#endif

include/swift/Parse/PersistentParserState.h

@@ -19,6 +19,7 @@
 
 #include "swift/Basic/SourceLoc.h"
 #include "swift/Parse/LocalContext.h"
+#include "swift/Parse/ParserPosition.h"
 #include "swift/Parse/Scope.h"
 #include "llvm/ADT/DenseMap.h"
 
@@ -114,7 +115,7 @@ class PersistentParserState {
   DelayedAccessorBodiesTy DelayedAccessorBodies;
 
   /// \brief Parser sets this if it stopped parsing before the buffer ended.
-  ParserPos MarkedPos;
+  ParserPosition MarkedPos;
 
   std::unique_ptr<DelayedDeclState> CodeCompletionDelayedDeclState;
 
@@ -166,16 +167,16 @@ class PersistentParserState {
     return TopLevelCode;
   }
 
-  void markParserPosition(SourceLoc Loc, SourceLoc PrevLoc,
+  void markParserPosition(ParserPosition Pos,
                           bool InPoundLineEnvironment) {
-    MarkedPos = {Loc, PrevLoc};
+    MarkedPos = Pos;
     this->InPoundLineEnvironment = InPoundLineEnvironment;
   }
 
   /// \brief Returns the marked parser position and resets it.
-  ParserPos takeParserPosition() {
-    ParserPos Pos = MarkedPos;
-    MarkedPos = ParserPos();
+  ParserPosition takeParserPosition() {
+    ParserPosition Pos = MarkedPos;
+    MarkedPos = ParserPosition();
     return Pos;
   }
 };

lib/Parse/ParseDecl.cpp

@@ -270,7 +270,7 @@ bool Parser::parseTopLevel() {
 
   // Next time start relexing from the beginning of the comment so that we can
   // attach it to the token.
-  State->markParserPosition(Tok.getCommentRange().getStart(), PreviousLoc,
+  State->markParserPosition(getParserPosition(),
                             InPoundLineEnvironment);
 
   // If we are done parsing the whole file, finalize the token receiver.

lib/Parse/ParseExpr.cpp

@@ -1926,13 +1926,13 @@ ParserResult<Expr> Parser::parseExprStringLiteral() {
       TmpContext.setDiscard();
 
       // Create a temporary lexer that lexes from the body of the string.
-      Lexer::State BeginState =
+      LexerState BeginState =
           L->getStateForBeginningOfTokenLoc(Segment.Loc);
       // We need to set the EOF at r_paren, to prevent the Lexer from eagerly
       // trying to lex the token beyond it. Parser::parseList() does a special
       // check for a tok::EOF that is spelled with a ')'.
       // FIXME: This seems like a hack, there must be a better way..
-      Lexer::State EndState = BeginState.advance(Segment.Length-1);
+      LexerState EndState = BeginState.advance(Segment.Length-1);
       Lexer LocalLex(*L, BeginState, EndState);
 
       // Temporarily swap out the parser's current lexer with our new one.
@@ -2235,8 +2235,8 @@ Expr *Parser::parseExprEditorPlaceholder(Token PlaceholderTok,
       SourceLoc TypeStartLoc = PlaceholderTok.getLoc().getAdvancedLoc(Offset);
       SourceLoc TypeEndLoc = TypeStartLoc.getAdvancedLoc(TyStr.size());
 
-      Lexer::State StartState = L->getStateForBeginningOfTokenLoc(TypeStartLoc);
-      Lexer::State EndState = L->getStateForBeginningOfTokenLoc(TypeEndLoc);
+      LexerState StartState = L->getStateForBeginningOfTokenLoc(TypeStartLoc);
+      LexerState EndState = L->getStateForBeginningOfTokenLoc(TypeEndLoc);
 
       // Create a lexer for the type sub-string.
       Lexer LocalLex(*L, StartState, EndState);

lib/Parse/Parser.cpp

@@ -460,9 +460,8 @@ Parser::Parser(std::unique_ptr<Lexer> Lex, SourceFile &SF,
 
   auto ParserPos = State->takeParserPosition();
   if (ParserPos.isValid() &&
-      SourceMgr.findBufferContainingLoc(ParserPos.Loc) == L->getBufferID()) {
-    auto BeginParserPosition = getParserPosition(ParserPos);
-    restoreParserPosition(BeginParserPosition);
+      L->isStateForCurrentBuffer(ParserPos.LS)) {
+    restoreParserPosition(ParserPos);
     InPoundLineEnvironment = State->InPoundLineEnvironment;
   }
 }

unittests/Parse/LexerTests.cpp

@@ -132,7 +132,7 @@ TEST_F(LexerTest, RestoreBasic) {
   ASSERT_EQ("bbb", Tok.getText());
   ASSERT_FALSE(Tok.isAtStartOfLine());
 
-  Lexer::State S = L.getStateForBeginningOfToken(Tok);
+  LexerState S = L.getStateForBeginningOfToken(Tok);
 
   L.lex(Tok);
   ASSERT_EQ(tok::identifier, Tok.getKind());
@@ -179,7 +179,7 @@ TEST_F(LexerTest, RestoreNewlineFlag) {
   ASSERT_EQ("bbb", Tok.getText());
   ASSERT_TRUE(Tok.isAtStartOfLine());
 
-  Lexer::State S = L.getStateForBeginningOfToken(Tok);
+  LexerState S = L.getStateForBeginningOfToken(Tok);
 
   L.lex(Tok);
   ASSERT_EQ(tok::identifier, Tok.getKind());
@@ -231,7 +231,7 @@ TEST_F(LexerTest, RestoreStopAtCodeCompletion) {
   ASSERT_EQ("bbb", Tok.getText());
   ASSERT_FALSE(Tok.isAtStartOfLine());
 
-  Lexer::State S = L.getStateForBeginningOfToken(Tok);
+  LexerState S = L.getStateForBeginningOfToken(Tok);
 
   L.lex(Tok);
   ASSERT_EQ(tok::identifier, Tok.getKind());
@@ -290,7 +290,7 @@ TEST_F(LexerTest, RestoreWithTrivia) {
             (Trivia{{TriviaPiece::newlines(1), TriviaPiece::spaces(1)}}));
   ASSERT_EQ(TrailingTrivia, (Trivia{{TriviaPiece::spaces(1)}}));
 
-  Lexer::State S = L.getStateForBeginningOfToken(Tok, LeadingTrivia);
+  LexerState S = L.getStateForBeginningOfToken(Tok, LeadingTrivia);
 
   L.lex(Tok, LeadingTrivia, TrailingTrivia);
   ASSERT_EQ(tok::identifier, Tok.getKind());

utils/round-trip-syntax-test

@@ -172,7 +172,7 @@ This driver invokes swift-syntax-test using -round-trip-lex and
                                  args.skip_bad_syntax)
                    for filename in all_input_files]
 
-    failed = reduce(lambda a, b: a and b,
+    failed = reduce(lambda a, b: a or b,
                     map(run_task, lex_tasks + parse_tasks))
     sys.exit(1 if failed else 0)