[clang-format] Fix bugs in parsing C++20 module import statements

[owenpan]

Summary

Also fixes https://github.com/llvm/llvm-project/issues/60145.

[owenpan]

Also fixed a potential bug in export import ....

[owenpan]

Simplified parseModuleImport() a little bit.

[owenpan]

Added test cases with import followed by comments.

[owenpan]

Put back the inadvertently removed line nextToken();, which was just before the outer while loop in parseModuleImport().

clang/lib/Format/UnwrappedLineParser.h

private:
   void parseTryCatch();
   void parseLoopBody(bool KeepBraces, bool WrapRightBrace);
   void parseForOrWhileLoop();
   void parseDoWhile();
   void parseLabel(bool LeftAlignLabel = false);
   void parseCaseLabel();
   void parseSwitch();
   void parseNamespace();
-  void parseModuleImport();
+  bool parseModuleImport();
   void parseNew();
   void parseAccessSpecifier();
   bool parseEnum();
   bool parseStructLike();
   bool parseRequires();
   void parseRequiresClause(FormatToken *RequiresToken);
   void parseRequiresExpression(FormatToken *RequiresToken);
   void parseConstraintExpression();

clang/lib/Format/UnwrappedLineParser.cpp

public:
   virtual FormatToken *getNextToken() = 0;
 
   // Returns the token preceding 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;
+  virtual FormatToken *peekNextToken(bool SkipComment = false) = 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;
if (eof())
       return &FakeEOF;
     return Token;
   }
 
   FormatToken *getPreviousToken() override {
     return PreviousTokenSource->getPreviousToken();
   }
 
-  FormatToken *peekNextToken() override {
+  FormatToken *peekNextToken(bool SkipComment) override {
     if (eof())
       return &FakeEOF;
-    return PreviousTokenSource->peekNextToken();
+    return PreviousTokenSource->peekNextToken(SkipComment);
   }
 
   bool isEOF() override { return PreviousTokenSource->isEOF(); }
 
   unsigned getPosition() override { return PreviousTokenSource->getPosition(); }
 
   FormatToken *setPosition(unsigned Position) override {
     PreviousToken = nullptr;
FormatToken *getNextToken() override {
     });
     return Tokens[Position];
   }
 
   FormatToken *getPreviousToken() override {
     return Position > 0 ? Tokens[Position - 1] : nullptr;
   }
 
-  FormatToken *peekNextToken() override {
+  FormatToken *peekNextToken(bool SkipComment) override {
     int Next = Position + 1;
+    if (SkipComment)
+      while (Tokens[Next]->is(tok::comment))
+        ++Next;
     LLVM_DEBUG({
       llvm::dbgs() << "Peeking ";
       dbgToken(Next);
     });
     return Tokens[Next];
   }
 
   bool isEOF() override { return Tokens[Position]->is(tok::eof); }
static bool isC78ParameterDecl(const FormatToken *Tok, const FormatToken *Next,
 
   Tok = Tok->Previous;
   if (!Tok || Tok->isNot(tok::identifier))
     return false;
 
   return Tok->Previous && Tok->Previous->isOneOf(tok::l_paren, tok::comma);
 }
 
-void UnwrappedLineParser::parseModuleImport() {
+bool UnwrappedLineParser::parseModuleImport() {
+  assert(FormatTok->is(Keywords.kw_import) && "'import' expected");
+
+  if (auto Token = Tokens->peekNextToken(/*SkipComment=*/true);
+      !Token->Tok.getIdentifierInfo() &&
+      !Token->isOneOf(tok::colon, tok::less, tok::string_literal)) {
+    return false;
+  }
+
   nextToken();
   while (!eof()) {
     if (FormatTok->is(tok::colon)) {
       FormatTok->setFinalizedType(TT_ModulePartitionColon);
     }
     // Handle import <foo/bar.h> as we would an include statement.
     else if (FormatTok->is(tok::less)) {
       nextToken();
while (!eof()) {
     if (FormatTok->is(tok::semi)) {
       nextToken();
       break;
     }
     nextToken();
   }
 
   addUnwrappedLine();
+  return true;
 }
 
 // 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.
 //
 // This method is conservative - it cannot cover all edge cases of JavaScript,
 // but only aims to correctly handle certain well known cases. It *must not*
case tok::kw_extern:
     break;
   case tok::kw_export:
     if (Style.isJavaScript()) {
       parseJavaScriptEs6ImportExport();
       return;
     }
     if (Style.isCpp()) {
       nextToken();
-      if (FormatTok->is(Keywords.kw_import)) {
-        parseModuleImport();
-        return;
-      }
       if (FormatTok->is(tok::kw_namespace)) {
         parseNamespace();
         return;
       }
+      if (FormatTok->is(Keywords.kw_import) && parseModuleImport())
+        return;
     }
     break;
   case tok::kw_inline:
     nextToken();
     if (FormatTok->is(tok::kw_namespace)) {
       parseNamespace();
       return;
     }
if (FormatTok->is(Keywords.kw_import)) {
         if (!FormatTok->is(tok::string_literal))
           return;
         nextToken();
         if (FormatTok->is(tok::semi))
           nextToken();
         addUnwrappedLine();
         return;
       }
-      if (Style.isCpp()) {
-        parseModuleImport();
+      if (Style.isCpp() && parseModuleImport())
         return;
-      }
-    }
     if (Style.isCpp() &&
         FormatTok->isOneOf(Keywords.kw_signals, Keywords.kw_qsignals,
                            Keywords.kw_slots, Keywords.kw_qslots)) {
       nextToken();
       if (FormatTok->is(tok::colon)) {
         nextToken();
         addUnwrappedLine();
         return;

clang/unittests/Format/FormatTest.cpp

TEST_F(FormatTest, HandlesIncludeDirectives) {
   FormatStyle Style = getLLVMStyle();
   Style.AlwaysBreakBeforeMultilineStrings = true;
   Style.ColumnLimit = 0;
   verifyFormat("#import \"abc.h\"", Style);
 
   // But 'import' might also be a regular C++ namespace.
   verifyFormat("import::SomeFunction(aaaaaaaaaaaaaaaaaaaaaaaaaaa,\n"
                "                     aaaaaaaaaaaaaaaaaaaaaaaaaaaaa);");
+  verifyFormat("import::Bar foo(val ? 2 : 1);");
 }
 
 //===----------------------------------------------------------------------===//
 // Error recovery tests.
 //===----------------------------------------------------------------------===//
 
 TEST_F(FormatTest, IncompleteParameterLists) {
   FormatStyle NoBinPacking = getLLVMStyle();
verifyFormat("export {\n"
                "};",
                Style);
   verifyFormat("export export char const *hello() { return \"hello\"; }");
 
   verifyFormat("import bar;", Style);
   verifyFormat("import foo.bar;", Style);
   verifyFormat("import foo:bar;", Style);
   verifyFormat("import :bar;", Style);
+  verifyFormat("import /* module partition */ :bar;", Style);
   verifyFormat("import <ctime>;", Style);
   verifyFormat("import \"header\";", Style);
 
   verifyFormat("module foo;", Style);
   verifyFormat("module foo:bar;", Style);
   verifyFormat("module foo.bar;", Style);
   verifyFormat("module;", Style);
 
TEST_F(FormatTest, Cpp20ModulesSupport) {
   verifyFormat("import a", Style);
   verifyFormat("module a", Style);
   verifyFormat("export import a", Style);
   verifyFormat("export module a", Style);
 
   verifyFormat("import", Style);
   verifyFormat("module", Style);
   verifyFormat("export", Style);
+
+  verifyFormat("import /* not keyword */ = val ? 2 : 1;");
 }
 
 TEST_F(FormatTest, CoroutineForCoawait) {
   FormatStyle Style = getLLVMStyle();
   verifyFormat("for co_await (auto x : range())\n  ;");
   verifyFormat("for (auto i : arr) {\n"
                "}",
                Style);