[clang] [clang-format] Correctly compute SplitPenalty of TrailingReturnArrow (PR #105613)

[Owen Pan via cfe-commits]

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

Fixes #105480.

>From 4bb6856deb83b89a5530efca2155da40663753e1 Mon Sep 17 00:00:00 2001
From: Owen Pan <owenpiano at gmail.com>
Date: Wed, 21 Aug 2024 21:45:12 -0700
Subject: [PATCH] [clang-format] Correctly compute SplitPenalty of
 TrailingReturnArrow

Fixes #105480.
---
 clang/lib/Format/TokenAnnotator.cpp           |  4 ++--
 clang/unittests/Format/TokenAnnotatorTest.cpp | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/clang/lib/Format/TokenAnnotator.cpp b/clang/lib/Format/TokenAnnotator.cpp
index 9d4204655b8ed6..fc1bacf8c5ac66 100644
--- a/clang/lib/Format/TokenAnnotator.cpp
+++ b/clang/lib/Format/TokenAnnotator.cpp
@@ -4050,7 +4050,7 @@ void TokenAnnotator::calculateFormattingInformation(AnnotatedLine &Line) const {
                              ChildSize + Current->SpacesRequiredBefore;
     }
 
-    if (Current->is(TT_CtorInitializerColon))
+    if (Current->isOneOf(TT_CtorInitializerColon, TT_LambdaLSquare))
       InFunctionDecl = false;
 
     // FIXME: Only calculate this if CanBreakBefore is true once static
@@ -4211,7 +4211,7 @@ unsigned TokenAnnotator::splitPenalty(const AnnotatedLine &Line,
   }
   if (Right.is(TT_PointerOrReference))
     return 190;
-  if (Right.is(TT_TrailingReturnArrow))
+  if (!InFunctionDecl && Right.is(TT_TrailingReturnArrow))
     return 110;
   if (Left.is(tok::equal) && Right.is(tok::l_brace))
     return 160;
diff --git a/clang/unittests/Format/TokenAnnotatorTest.cpp b/clang/unittests/Format/TokenAnnotatorTest.cpp
index 99798de43e70ff..b057ac49659acf 100644
--- a/clang/unittests/Format/TokenAnnotatorTest.cpp
+++ b/clang/unittests/Format/TokenAnnotatorTest.cpp
@@ -42,6 +42,8 @@ class TokenAnnotatorTest : public testing::Test {
   EXPECT_EQ((FormatTok)->getPrecedence(), Prec) << *(FormatTok)
 #define EXPECT_BRACE_KIND(FormatTok, Kind)                                     \
   EXPECT_EQ(FormatTok->getBlockKind(), Kind) << *(FormatTok)
+#define EXPECT_SPLIT_PENALTY(FormatTok, Penalty)                               \
+  EXPECT_EQ(FormatTok->SplitPenalty, (unsigned)(Penalty)) << *(FormatTok)
 #define EXPECT_TOKEN(FormatTok, Kind, Type)                                    \
   do {                                                                         \
     EXPECT_TOKEN_KIND(FormatTok, Kind);                                        \
@@ -3369,6 +3371,20 @@ TEST_F(TokenAnnotatorTest, GNULanguageStandard) {
   EXPECT_TOKEN(Tokens[2], tok::spaceship, TT_BinaryOperator);
 }
 
+TEST_F(TokenAnnotatorTest, SplitPenalty) {
+  auto Style = getLLVMStyle();
+  Style.ColumnLimit = 20;
+
+  auto Tokens = annotate("class foo {\n"
+                         "  auto bar()\n"
+                         "      -> bool;\n"
+                         "};",
+                         Style);
+  ASSERT_EQ(Tokens.size(), 13u);
+  EXPECT_TOKEN(Tokens[7], tok::arrow, TT_TrailingReturnArrow);
+  EXPECT_SPLIT_PENALTY(Tokens[7], 23);
+}
+
 } // namespace
 } // namespace format
 } // namespace clang