[clang] [clang-format] Fix a bug in annotating angles containing FatArrow (PR #108671)

[Owen Pan via cfe-commits]

https://github.com/owenca created https://github.com/llvm/llvm-project/pull/108671

Fixes #108536.

>From 90a2015196595bb72fb64490d1579f01b7c18a2a Mon Sep 17 00:00:00 2001
From: Owen Pan <owenpiano at gmail.com>
Date: Fri, 13 Sep 2024 19:36:48 -0700
Subject: [PATCH] [clang-format] Fix a bug in annotating angles containing
 FatArrow

Fixes #108536.
---
 clang/lib/Format/TokenAnnotator.cpp           | 10 ++++++----
 clang/unittests/Format/TokenAnnotatorTest.cpp |  9 +++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/clang/lib/Format/TokenAnnotator.cpp b/clang/lib/Format/TokenAnnotator.cpp
index dfa703aed0d34d..4313cc20dd61a3 100644
--- a/clang/lib/Format/TokenAnnotator.cpp
+++ b/clang/lib/Format/TokenAnnotator.cpp
@@ -189,7 +189,8 @@ class AnnotatingParser {
       next();
     }
 
-    for (bool SeenTernaryOperator = false; CurrentToken;) {
+    for (bool SeenTernaryOperator = false, SeenFatArrow = false;
+         CurrentToken;) {
       const bool InExpr = Contexts[Contexts.size() - 2].IsExpression;
       if (CurrentToken->is(tok::greater)) {
         const auto *Next = CurrentToken->Next;
@@ -243,7 +244,7 @@ class AnnotatingParser {
       // operator that was misinterpreted because we are parsing template
       // parameters.
       // FIXME: This is getting out of hand, write a decent parser.
-      if (InExpr && !Line.startsWith(tok::kw_template) &&
+      if (InExpr && !SeenFatArrow && !Line.startsWith(tok::kw_template) &&
           Prev.is(TT_BinaryOperator)) {
         const auto Precedence = Prev.getPrecedence();
         if (Precedence > prec::Conditional && Precedence < prec::Relational)
@@ -251,6 +252,8 @@ class AnnotatingParser {
       }
       if (Prev.isOneOf(tok::question, tok::colon) && !Style.isProto())
         SeenTernaryOperator = true;
+      else if (Prev.is(TT_FatArrow))
+        SeenFatArrow = true;
       updateParameterCount(Left, CurrentToken);
       if (Style.Language == FormatStyle::LK_Proto) {
         if (FormatToken *Previous = CurrentToken->getPreviousNonComment()) {
@@ -260,8 +263,7 @@ class AnnotatingParser {
             Previous->setType(TT_SelectorName);
           }
         }
-      }
-      if (Style.isTableGen()) {
+      } else if (Style.isTableGen()) {
         if (CurrentToken->isOneOf(tok::comma, tok::equal)) {
           // They appear as separators. Unless they are not in class definition.
           next();
diff --git a/clang/unittests/Format/TokenAnnotatorTest.cpp b/clang/unittests/Format/TokenAnnotatorTest.cpp
index 5c28e3a4ea5a1f..c3ac5df12d085c 100644
--- a/clang/unittests/Format/TokenAnnotatorTest.cpp
+++ b/clang/unittests/Format/TokenAnnotatorTest.cpp
@@ -616,6 +616,15 @@ TEST_F(TokenAnnotatorTest, UnderstandsTernaryInTemplate) {
   EXPECT_TOKEN(Tokens[8], tok::greater, TT_TemplateCloser);
 }
 
+TEST_F(TokenAnnotatorTest, FatArrowInAngleBrackets) {
+  auto Tokens = annotate("foo = new Bar<(id: int) => X | Y>();",
+                         getLLVMStyle(FormatStyle::LK_JavaScript));
+  ASSERT_EQ(Tokens.size(), 19u) << Tokens;
+  EXPECT_TOKEN(Tokens[4], tok::less, TT_TemplateOpener);
+  EXPECT_TOKEN(Tokens[10], tok::equal, TT_FatArrow);
+  EXPECT_TOKEN(Tokens[14], tok::greater, TT_TemplateCloser);
+}
+
 TEST_F(TokenAnnotatorTest, UnderstandsNonTemplateAngleBrackets) {
   auto Tokens = annotate("return a < b && c > d;");
   ASSERT_EQ(Tokens.size(), 10u) << Tokens;