[clang-format] Add a space between an overloaded operator and '>'

The token annotator doesn't annotate the template opener and closer
as such if they enclose an overloaded operator. This causes the
space between the operator and the closer to be removed, resulting
in invalid C++ code.

Fixes llvm#58602.

Differential Revision: https://reviews.llvm.org/D143755

Author: owenca

clang/lib/Format/FormatToken.h

@@ -743,8 +743,8 @@ struct FormatToken {
   }
 
   /// Returns the next token ignoring comments.
-  [[nodiscard]] const FormatToken *getNextNonComment() const {
-    const FormatToken *Tok = Next;
+  [[nodiscard]] FormatToken *getNextNonComment() const {
+    FormatToken *Tok = Next;
     while (Tok && Tok->is(tok::comment))
       Tok = Tok->Next;
     return Tok;

clang/lib/Format/FormatTokenLexer.cpp

@@ -103,6 +103,8 @@ void FormatTokenLexer::tryMergePreviousTokens() {
     return;
   if (tryMergeLessLess())
     return;
+  if (tryMergeGreaterGreater())
+    return;
   if (tryMergeForEach())
     return;
   if (Style.isCpp() && tryTransformTryUsageForC())
@@ -460,12 +462,11 @@ bool FormatTokenLexer::tryMergeLessLess() {
     return false;
 
   auto X = Tokens.size() > 3 ? First[-1] : nullptr;
-  auto Y = First[2];
-  if ((X && X->is(tok::less)) || Y->is(tok::less))
+  if (X && X->is(tok::less))
     return false;
 
-  // Do not remove a whitespace between the two "<" e.g. "operator< <>".
-  if (X && X->is(tok::kw_operator) && Y->is(tok::greater))
+  auto Y = First[2];
+  if ((!X || X->isNot(tok::kw_operator)) && Y->is(tok::less))
     return false;
 
   First[0]->Tok.setKind(tok::lessless);
@@ -475,6 +476,30 @@ bool FormatTokenLexer::tryMergeLessLess() {
   return true;
 }
 
+bool FormatTokenLexer::tryMergeGreaterGreater() {
+  // Merge kw_operator,greater,greater into kw_operator,greatergreater.
+  if (Tokens.size() < 2)
+    return false;
+
+  auto First = Tokens.end() - 2;
+  if (First[0]->isNot(tok::greater) || First[1]->isNot(tok::greater))
+    return false;
+
+  // Only merge if there currently is no whitespace between the first two ">".
+  if (First[1]->hasWhitespaceBefore())
+    return false;
+
+  auto Tok = Tokens.size() > 2 ? First[-1] : nullptr;
+  if (Tok && Tok->isNot(tok::kw_operator))
+    return false;
+
+  First[0]->Tok.setKind(tok::greatergreater);
+  First[0]->TokenText = ">>";
+  First[0]->ColumnWidth += 1;
+  Tokens.erase(Tokens.end() - 1);
+  return true;
+}
+
 bool FormatTokenLexer::tryMergeTokens(ArrayRef<tok::TokenKind> Kinds,
                                       TokenType NewType) {
   if (Tokens.size() < Kinds.size())

clang/lib/Format/FormatTokenLexer.h

@@ -51,6 +51,7 @@ class FormatTokenLexer {
   void tryMergePreviousTokens();
 
   bool tryMergeLessLess();
+  bool tryMergeGreaterGreater();
   bool tryMergeNSStringLiteral();
   bool tryMergeJSPrivateIdentifier();
   bool tryMergeCSharpStringLiteral();

clang/lib/Format/TokenAnnotator.cpp

@@ -1219,19 +1219,25 @@ class AnnotatingParser {
              !CurrentToken->isOneOf(tok::l_paren, tok::semi, tok::r_paren)) {
         if (CurrentToken->isOneOf(tok::star, tok::amp))
           CurrentToken->setType(TT_PointerOrReference);
-        consumeToken();
-        if (!CurrentToken)
-          continue;
-        if (CurrentToken->is(tok::comma) &&
-            CurrentToken->Previous->isNot(tok::kw_operator)) {
+        auto Next = CurrentToken->getNextNonComment();
+        if (!Next)
           break;
-        }
-        if (CurrentToken->Previous->isOneOf(TT_BinaryOperator, TT_UnaryOperator,
-                                            tok::comma, tok::star, tok::arrow,
-                                            tok::amp, tok::ampamp) ||
+        if (Next->is(tok::less))
+          next();
+        else
+          consumeToken();
+        assert(CurrentToken);
+        auto Previous = CurrentToken->getPreviousNonComment();
+        assert(Previous);
+        if (CurrentToken->is(tok::comma) && Previous->isNot(tok::kw_operator))
+          break;
+        if (Previous->isOneOf(TT_BinaryOperator, TT_UnaryOperator, tok::comma,
+                              tok::star, tok::arrow, tok::amp, tok::ampamp) ||
             // User defined literal.
-            CurrentToken->Previous->TokenText.startswith("\"\"")) {
-          CurrentToken->Previous->setType(TT_OverloadedOperator);
+            Previous->TokenText.startswith("\"\"")) {
+          Previous->setType(TT_OverloadedOperator);
+          if (CurrentToken->isOneOf(tok::less, tok::greater))
+            break;
         }
       }
       if (CurrentToken && CurrentToken->is(tok::l_paren))
@@ -3893,6 +3899,10 @@ bool TokenAnnotator::spaceRequiredBefore(const AnnotatedLine &Line,
     return true;
 
   if (Style.isCpp()) {
+    if (Left.is(TT_OverloadedOperator) &&
+        Right.isOneOf(TT_TemplateOpener, TT_TemplateCloser)) {
+      return true;
+    }
     // Space between UDL and dot: auto b = 4s .count();
     if (Right.is(tok::period) && Left.is(tok::numeric_constant))
       return true;

clang/unittests/Format/FormatTest.cpp

@@ -10663,6 +10663,14 @@ TEST_F(FormatTest, UnderstandsOverloadedOperators) {
   verifyFormat("foo() { ::operator new(n * sizeof(foo)); }");
 }
 
+TEST_F(FormatTest, SpaceBeforeTemplateCloser) {
+  verifyFormat("C<&operator- > minus;");
+  verifyFormat("C<&operator> > gt;");
+  verifyFormat("C<&operator>= > ge;");
+  verifyFormat("C<&operator<= > le;");
+  verifyFormat("C<&operator< <X>> lt;");
+}
+
 TEST_F(FormatTest, UnderstandsFunctionRefQualification) {
   verifyFormat("void A::b() && {}");
   verifyFormat("void A::b() && noexcept {}");

clang/unittests/Format/TokenAnnotatorTest.cpp

@@ -574,6 +574,71 @@ TEST_F(TokenAnnotatorTest, UnderstandsOverloadedOperators) {
   EXPECT_TOKEN(Tokens[4], tok::l_paren, TT_OverloadedOperatorLParen);
 }
 
+TEST_F(TokenAnnotatorTest, OverloadedOperatorInTemplate) {
+  struct {
+    const char *Text;
+    tok::TokenKind Kind;
+  } Operators[] = {{"+", tok::plus},
+                   {"-", tok::minus},
+                   // FIXME:
+                   // {"*", tok::star},
+                   {"/", tok::slash},
+                   {"%", tok::percent},
+                   {"^", tok::caret},
+                   // FIXME:
+                   // {"&", tok::amp},
+                   {"|", tok::pipe},
+                   {"~", tok::tilde},
+                   {"!", tok::exclaim},
+                   {"=", tok::equal},
+                   // FIXME:
+                   // {"<", tok::less},
+                   {">", tok::greater},
+                   {"+=", tok::plusequal},
+                   {"-=", tok::minusequal},
+                   {"*=", tok::starequal},
+                   {"/=", tok::slashequal},
+                   {"%=", tok::percentequal},
+                   {"^=", tok::caretequal},
+                   {"&=", tok::ampequal},
+                   {"|=", tok::pipeequal},
+                   {"<<", tok::lessless},
+                   {">>", tok::greatergreater},
+                   {">>=", tok::greatergreaterequal},
+                   {"<<=", tok::lesslessequal},
+                   {"==", tok::equalequal},
+                   {"!=", tok::exclaimequal},
+                   {"<=", tok::lessequal},
+                   {">=", tok::greaterequal},
+                   {"<=>", tok::spaceship},
+                   {"&&", tok::ampamp},
+                   {"||", tok::pipepipe},
+                   {"++", tok::plusplus},
+                   {"--", tok::minusminus},
+                   {",", tok::comma},
+                   {"->*", tok::arrowstar},
+                   {"->", tok::arrow}};
+
+  for (const auto &Operator : Operators) {
+    std::string Input("C<&operator");
+    Input += Operator.Text;
+    Input += " > a;";
+    auto Tokens = annotate(std::string(Input));
+    ASSERT_EQ(Tokens.size(), 9u) << Tokens;
+    EXPECT_TOKEN(Tokens[1], tok::less, TT_TemplateOpener);
+    EXPECT_TOKEN(Tokens[4], Operator.Kind, TT_OverloadedOperator);
+    EXPECT_TOKEN(Tokens[5], tok::greater, TT_TemplateCloser);
+  }
+
+  auto Tokens = annotate("C<&operator< <X>> lt;");
+  ASSERT_EQ(Tokens.size(), 12u) << Tokens;
+  EXPECT_TOKEN(Tokens[1], tok::less, TT_TemplateOpener);
+  EXPECT_TOKEN(Tokens[4], tok::less, TT_OverloadedOperator);
+  EXPECT_TOKEN(Tokens[5], tok::less, TT_TemplateOpener);
+  EXPECT_TOKEN(Tokens[7], tok::greater, TT_TemplateCloser);
+  EXPECT_TOKEN(Tokens[8], tok::greater, TT_TemplateCloser);
+}
+
 TEST_F(TokenAnnotatorTest, UnderstandsRequiresClausesAndConcepts) {
   auto Tokens = annotate("template <typename T>\n"
                          "concept C = (Foo && Bar) && (Bar && Baz);");