1 files changed, 129 insertions, 30 deletions

diff --git a/clang/lib/Format/UnwrappedLineParser.cpp b/clang/lib/Format/UnwrappedLineParser.cpp
index 28d925858f77..5b9fe267aae6 100644
--- a/clang/lib/Format/UnwrappedLineParser.cpp
+++ b/clang/lib/Format/UnwrappedLineParser.cpp
@@ -28,9 +28,28 @@ namespace format {
 class FormatTokenSource {
 public:
   virtual ~FormatTokenSource() {}
+
+  // Returns the next token in the token stream.
   virtual FormatToken *getNextToken() = 0;
 
+  // Returns the token precedint the token returned by the last call to
+  // getNextToken() in the token stream, or nullptr if no such token exists.
+  virtual FormatToken *getPreviousToken() = 0;
+
+  // Returns the token that would be returned by the next call to
+  // getNextToken().
+  virtual FormatToken *peekNextToken() = 0;
+
+  // Returns whether we are at the end of the file.
+  // This can be different from whether getNextToken() returned an eof token
+  // when the FormatTokenSource is a view on a part of the token stream.
+  virtual bool isEOF() = 0;
+
+  // Gets the current position in the token stream, to be used by setPosition().
   virtual unsigned getPosition() = 0;
+
+  // Resets the token stream to the state it was in when getPosition() returned
+  // Position, and return the token at that position in the stream.
   virtual FormatToken *setPosition(unsigned Position) = 0;
 };
 
@@ -108,6 +127,18 @@ public:
     return Token;
   }
 
+  FormatToken *getPreviousToken() override {
+    return PreviousTokenSource->getPreviousToken();
+  }
+
+  FormatToken *peekNextToken() override {
+    if (eof())
+      return &FakeEOF;
+    return PreviousTokenSource->peekNextToken();
+  }
+
+  bool isEOF() override { return PreviousTokenSource->isEOF(); }
+
   unsigned getPosition() override { return PreviousTokenSource->getPosition(); }
 
   FormatToken *setPosition(unsigned Position) override {
@@ -199,16 +230,45 @@ public:
       : Tokens(Tokens), Position(-1) {}
 
   FormatToken *getNextToken() override {
+    if (Position >= 0 && Tokens[Position]->is(tok::eof)) {
+      LLVM_DEBUG({
+        llvm::dbgs() << "Next ";
+        dbgToken(Position);
+      });
+      return Tokens[Position];
+    }
     ++Position;
+    LLVM_DEBUG({
+      llvm::dbgs() << "Next ";
+      dbgToken(Position);
+    });
     return Tokens[Position];
   }
 
+  FormatToken *getPreviousToken() override {
+    assert(Position > 0);
+    return Tokens[Position - 1];
+  }
+
+  FormatToken *peekNextToken() override {
+    int Next = Position + 1;
+    LLVM_DEBUG({
+      llvm::dbgs() << "Peeking ";
+      dbgToken(Next);
+    });
+    return Tokens[Next];
+  }
+
+  bool isEOF() override { return Tokens[Position]->is(tok::eof); }
+
   unsigned getPosition() override {
+    LLVM_DEBUG(llvm::dbgs() << "Getting Position: " << Position << "\n");
     assert(Position >= 0);
     return Position;
   }
 
   FormatToken *setPosition(unsigned P) override {
+    LLVM_DEBUG(llvm::dbgs() << "Setting Position: " << P << "\n");
     Position = P;
     return Tokens[Position];
   }
@@ -216,6 +276,13 @@ public:
   void reset() { Position = -1; }
 
 private:
+  void dbgToken(int Position, llvm::StringRef Indent = "") {
+    FormatToken *Tok = Tokens[Position];
+    llvm::dbgs() << Indent << "[" << Position
+                 << "] Token: " << Tok->Tok.getName() << " / " << Tok->TokenText
+                 << ", Macro: " << !!Tok->MacroCtx << "\n";
+  }
+
   ArrayRef<FormatToken *> Tokens;
   int Position;
 };
@@ -399,7 +466,7 @@ void UnwrappedLineParser::parseLevel(bool HasOpeningBrace) {
       FormatToken *Next;
       do {
         Next = Tokens->getNextToken();
-      } while (Next && Next->is(tok::comment));
+      } while (Next->is(tok::comment));
       FormatTok = Tokens->setPosition(StoredPosition);
       if (Next && Next->isNot(tok::colon)) {
         // default not followed by ':' is not a case label; treat it like
@@ -875,10 +942,7 @@ void UnwrappedLineParser::parsePPEndIf() {
   parsePPUnknown();
   // If the #endif of a potential include guard is the last thing in the file,
   // then we found an include guard.
-  unsigned TokenPosition = Tokens->getPosition();
-  FormatToken *PeekNext = AllTokens[TokenPosition];
-  if (IncludeGuard == IG_Defined && PPBranchLevel == -1 &&
-      PeekNext->is(tok::eof) &&
+  if (IncludeGuard == IG_Defined && PPBranchLevel == -1 && Tokens->isEOF() &&
       Style.IndentPPDirectives != FormatStyle::PPDIS_None)
     IncludeGuard = IG_Found;
 }
@@ -1050,6 +1114,35 @@ static bool isC78ParameterDecl(const FormatToken *Tok, const FormatToken *Next,
   return Tok->Previous && Tok->Previous->isOneOf(tok::l_paren, tok::comma);
 }
 
+void UnwrappedLineParser::parseModuleImport() {
+  nextToken();
+  while (!eof()) {
+    if (FormatTok->is(tok::colon)) {
+      FormatTok->setType(TT_ModulePartitionColon);
+    }
+    // Handle import <foo/bar.h> as we would an include statement.
+    else if (FormatTok->is(tok::less)) {
+      nextToken();
+      while (!FormatTok->isOneOf(tok::semi, tok::greater, tok::eof)) {
+        // Mark tokens up to the trailing line comments as implicit string
+        // literals.
+        if (FormatTok->isNot(tok::comment) &&
+            !FormatTok->TokenText.startswith("//"))
+          FormatTok->setType(TT_ImplicitStringLiteral);
+        nextToken();
+      }
+    }
+    if (FormatTok->is(tok::semi)) {
+      nextToken();
+      break;
+    }
+    nextToken();
+  }
+
+  addUnwrappedLine();
+  return;
+}
+
 // readTokenWithJavaScriptASI reads the next token and terminates the current
 // line if JavaScript Automatic Semicolon Insertion must
 // happen between the current token and the next token.
@@ -1097,7 +1190,6 @@ void UnwrappedLineParser::readTokenWithJavaScriptASI() {
 }
 
 void UnwrappedLineParser::parseStructuralElement(bool IsTopLevel) {
-  assert(!FormatTok->is(tok::l_brace));
   if (Style.Language == FormatStyle::LK_TableGen &&
       FormatTok->is(tok::pp_include)) {
     nextToken();
@@ -1249,6 +1341,10 @@ void UnwrappedLineParser::parseStructuralElement(bool IsTopLevel) {
         addUnwrappedLine();
         return;
       }
+      if (Style.isCpp()) {
+        parseModuleImport();
+        return;
+      }
     }
     if (Style.isCpp() &&
         FormatTok->isOneOf(Keywords.kw_signals, Keywords.kw_qsignals,
@@ -1402,9 +1498,7 @@ void UnwrappedLineParser::parseStructuralElement(bool IsTopLevel) {
       // declaration.
       if (!IsTopLevel || !Style.isCpp() || !Previous || FormatTok->is(tok::eof))
         break;
-      const unsigned Position = Tokens->getPosition() + 1;
-      assert(Position < AllTokens.size());
-      if (isC78ParameterDecl(FormatTok, AllTokens[Position], Previous)) {
+      if (isC78ParameterDecl(FormatTok, Tokens->peekNextToken(), Previous)) {
         addUnwrappedLine();
         return;
       }
@@ -1488,7 +1582,7 @@ void UnwrappedLineParser::parseStructuralElement(bool IsTopLevel) {
           unsigned StoredPosition = Tokens->getPosition();
           FormatToken *Next = Tokens->getNextToken();
           FormatTok = Tokens->setPosition(StoredPosition);
-          if (Next && !mustBeJSIdent(Keywords, Next)) {
+          if (!mustBeJSIdent(Keywords, Next)) {
             nextToken();
             break;
           }
@@ -2099,8 +2193,8 @@ void UnwrappedLineParser::parseIfThenElse() {
       parseBlock();
       addUnwrappedLine();
     } else if (FormatTok->Tok.is(tok::kw_if)) {
-      FormatToken *Previous = AllTokens[Tokens->getPosition() - 1];
-      bool PrecededByComment = Previous->is(tok::comment);
+      FormatToken *Previous = Tokens->getPreviousToken();
+      bool PrecededByComment = Previous && Previous->is(tok::comment);
       if (PrecededByComment) {
         addUnwrappedLine();
         ++Line->Level;
@@ -2305,6 +2399,8 @@ void UnwrappedLineParser::parseForOrWhileLoop() {
   if (Style.Language == FormatStyle::LK_JavaScript &&
       FormatTok->is(Keywords.kw_await))
     nextToken();
+  if (Style.isCpp() && FormatTok->is(tok::kw_co_await))
+    nextToken();
   if (FormatTok->Tok.is(tok::l_paren))
     parseParens();
   if (FormatTok->Tok.is(tok::l_brace)) {
@@ -2653,23 +2749,25 @@ bool UnwrappedLineParser::tryToParseSimpleAttribute() {
   ScopedTokenPosition AutoPosition(Tokens);
   FormatToken *Tok = Tokens->getNextToken();
   // We already read the first [ check for the second.
-  if (Tok && !Tok->is(tok::l_square)) {
+  if (!Tok->is(tok::l_square)) {
     return false;
   }
   // Double check that the attribute is just something
   // fairly simple.
-  while (Tok) {
+  while (Tok->isNot(tok::eof)) {
     if (Tok->is(tok::r_square)) {
       break;
     }
     Tok = Tokens->getNextToken();
   }
+  if (Tok->is(tok::eof))
+    return false;
   Tok = Tokens->getNextToken();
-  if (Tok && !Tok->is(tok::r_square)) {
+  if (!Tok->is(tok::r_square)) {
     return false;
   }
   Tok = Tokens->getNextToken();
-  if (Tok && Tok->is(tok::semi)) {
+  if (Tok->is(tok::semi)) {
     return false;
   }
   return true;
@@ -2682,7 +2780,7 @@ void UnwrappedLineParser::parseJavaEnumBody() {
   unsigned StoredPosition = Tokens->getPosition();
   bool IsSimple = true;
   FormatToken *Tok = Tokens->getNextToken();
-  while (Tok) {
+  while (!Tok->is(tok::eof)) {
     if (Tok->is(tok::r_brace))
       break;
     if (Tok->isOneOf(tok::l_brace, tok::semi)) {
@@ -3292,6 +3390,20 @@ void UnwrappedLineParser::readToken(int LevelDifference) {
   do {
     FormatTok = Tokens->getNextToken();
     assert(FormatTok);
+    while (FormatTok->getType() == TT_ConflictStart ||
+           FormatTok->getType() == TT_ConflictEnd ||
+           FormatTok->getType() == TT_ConflictAlternative) {
+      if (FormatTok->getType() == TT_ConflictStart) {
+        conditionalCompilationStart(/*Unreachable=*/false);
+      } else if (FormatTok->getType() == TT_ConflictAlternative) {
+        conditionalCompilationAlternative();
+      } else if (FormatTok->getType() == TT_ConflictEnd) {
+        conditionalCompilationEnd();
+      }
+      FormatTok = Tokens->getNextToken();
+      FormatTok->MustBreakBefore = true;
+    }
+
     while (!Line->InPPDirective && FormatTok->Tok.is(tok::hash) &&
            (FormatTok->HasUnescapedNewline || FormatTok->IsFirst)) {
       distributeComments(Comments, FormatTok);
@@ -3313,19 +3425,6 @@ void UnwrappedLineParser::readToken(int LevelDifference) {
       flushComments(isOnNewLine(*FormatTok));
       parsePPDirective();
     }
-    while (FormatTok->getType() == TT_ConflictStart ||
-           FormatTok->getType() == TT_ConflictEnd ||
-           FormatTok->getType() == TT_ConflictAlternative) {
-      if (FormatTok->getType() == TT_ConflictStart) {
-        conditionalCompilationStart(/*Unreachable=*/false);
-      } else if (FormatTok->getType() == TT_ConflictAlternative) {
-        conditionalCompilationAlternative();
-      } else if (FormatTok->getType() == TT_ConflictEnd) {
-        conditionalCompilationEnd();
-      }
-      FormatTok = Tokens->getNextToken();
-      FormatTok->MustBreakBefore = true;
-    }
 
     if (!PPStack.empty() && (PPStack.back().Kind == PP_Unreachable) &&
         !Line->InPPDirective) {