[clang-format] Fix a bug that wraps before function arguments

Fixes a long-standing bug that erroneously placed function arguments on a
new line despite all arguments being able to fit on the same line.

The original diff that introduced the bug implemented behaviour that pushed
the first argument to a function onto a new line under certain circumstances
relating passing lambdas as arguments.

This behaviour was implemented in TokenAnnotator::mustBreakBefore() which
meant the code lacked the necessary context to figure out whether subsequent
arguments might be able to all fit on one line. As such, I've moved the
implementation to ContinuationIndenter and, instead of forcing a line break
at the first argument in all cases, we now allow the OptimizingLineFormatter
to consider placing the first argument on the same line as the function call
but don't allow further line breaks in this case.

The end result is that either the first argument must go on a new line (as
before) or all arguments must be put on the current line.

Closes llvm#44486.

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

Author: jp4a50

Diff for: clang/lib/Format/ContinuationIndenter.cpp

@@ -260,6 +260,7 @@ LineState ContinuationIndenter::getInitialState(unsigned FirstIndent,
                                    /*NoLineBreak=*/false));
   State.NoContinuation = false;
   State.StartOfStringLiteral = 0;
+  State.NoLineBreak = false;
   State.StartOfLineLevel = 0;
   State.LowestLevelOnLine = 0;
   State.IgnoreStackForComparison = false;
@@ -342,7 +343,7 @@ bool ContinuationIndenter::canBreak(const LineState &State) {
     return true;
   }
 
-  return !CurrentState.NoLineBreak;
+  return !State.NoLineBreak && !CurrentState.NoLineBreak;
 }
 
 bool ContinuationIndenter::mustBreak(const LineState &State) {
@@ -653,6 +654,38 @@ void ContinuationIndenter::addTokenOnCurrentLine(LineState &State, bool DryRun,
   const FormatToken &Previous = *State.NextToken->Previous;
   auto &CurrentState = State.Stack.back();
 
+  bool DisallowLineBreaksOnThisLine = Style.isCpp() && [&Current] {
+    // Deal with lambda arguments in C++. The aim here is to ensure that we
+    // don't over-indent lambda function bodies when lambdas are passed as
+    // arguments to function calls. We do this by ensuring that either all
+    // arguments (including any lambdas) go on the same line as the function
+    // call, or we break before the first argument.
+    auto PrevNonComment = Current.getPreviousNonComment();
+    if (!PrevNonComment || PrevNonComment->isNot(tok::l_paren))
+      return false;
+    if (Current.isOneOf(tok::comment, tok::l_paren, TT_LambdaLSquare))
+      return false;
+    auto BlockParameterCount = PrevNonComment->BlockParameterCount;
+    if (BlockParameterCount == 0)
+      return false;
+
+    // Multiple lambdas in the same function call.
+    if (BlockParameterCount > 1)
+      return true;
+
+    // A lambda followed by another arg.
+    if (!PrevNonComment->Role)
+      return false;
+    auto Comma = PrevNonComment->Role->lastComma();
+    if (!Comma)
+      return false;
+    auto Next = Comma->getNextNonComment();
+    return Next && !Next->isOneOf(TT_LambdaLSquare, tok::l_brace, tok::caret);
+  }();
+
+  if (DisallowLineBreaksOnThisLine)
+    State.NoLineBreak = true;
+
   if (Current.is(tok::equal) &&
       (State.Line->First->is(tok::kw_for) || Current.NestingLevel == 0) &&
       CurrentState.VariablePos == 0) {

Diff for: clang/lib/Format/ContinuationIndenter.h

@@ -433,6 +433,9 @@ struct LineState {
   /// literal sequence, 0 otherwise.
   unsigned StartOfStringLiteral;
 
+  /// Disallow line breaks for this line.
+  bool NoLineBreak;
+
   /// A stack keeping track of properties applying to parenthesis
   /// levels.
   SmallVector<ParenState> Stack;

Diff for: clang/lib/Format/TokenAnnotator.cpp

@@ -5193,30 +5193,6 @@ bool TokenAnnotator::mustBreakBefore(const AnnotatedLine &Line,
       return true;
   }
 
-  // Deal with lambda arguments in C++ - we want consistent line breaks whether
-  // they happen to be at arg0, arg1 or argN. The selection is a bit nuanced
-  // as aggressive line breaks are placed when the lambda is not the last arg.
-  if ((Style.Language == FormatStyle::LK_Cpp ||
-       Style.Language == FormatStyle::LK_ObjC) &&
-      Left.is(tok::l_paren) && Left.BlockParameterCount > 0 &&
-      !Right.isOneOf(tok::l_paren, TT_LambdaLSquare)) {
-    // Multiple lambdas in the same function call force line breaks.
-    if (Left.BlockParameterCount > 1)
-      return true;
-
-    // A lambda followed by another arg forces a line break.
-    if (!Left.Role)
-      return false;
-    auto Comma = Left.Role->lastComma();
-    if (!Comma)
-      return false;
-    auto Next = Comma->getNextNonComment();
-    if (!Next)
-      return false;
-    if (!Next->isOneOf(TT_LambdaLSquare, tok::l_brace, tok::caret))
-      return true;
-  }
-
   return false;
 }
 

Diff for: clang/unittests/Format/FormatTest.cpp

@@ -22224,8 +22224,25 @@ TEST_F(FormatTest, FormatsLambdas) {
                "  }\n"
                "};");
 
-  // Multiple lambdas in the same parentheses change indentation rules. These
-  // lambdas are forced to start on new lines.
+  // Lambdas that fit on a single line within an argument list are not forced
+  // onto new lines.
+  verifyFormat("SomeFunction([] {});");
+  verifyFormat("SomeFunction(0, [] {});");
+  verifyFormat("SomeFunction([] {}, 0);");
+  verifyFormat("SomeFunction(0, [] {}, 0);");
+  verifyFormat("SomeFunction([] { return 0; }, 0);");
+  verifyFormat("SomeFunction(a, [] { return 0; }, b);");
+  verifyFormat("SomeFunction([] { return 0; }, [] { return 0; });");
+  verifyFormat("SomeFunction([] { return 0; }, [] { return 0; }, b);");
+  verifyFormat("auto loooooooooooooooooooooooooooong =\n"
+               "    SomeFunction([] { return 0; }, [] { return 0; }, b);");
+  // Exceeded column limit. We need to break.
+  verifyFormat("auto loooooooooooooooooooooooooooongName = SomeFunction(\n"
+               "    [] { return anotherLooooooooooonoooooooongName; }, [] { "
+               "return 0; }, b);");
+
+  // Multiple multi-line lambdas in the same parentheses change indentation
+  // rules. These lambdas are always forced to start on new lines.
   verifyFormat("SomeFunction(\n"
                "    []() {\n"
                "      //\n"
@@ -22234,7 +22251,7 @@ TEST_F(FormatTest, FormatsLambdas) {
                "      //\n"
                "    });");
 
-  // A lambda passed as arg0 is always pushed to the next line.
+  // A multi-line lambda passed as arg0 is always pushed to the next line.
   verifyFormat("SomeFunction(\n"
                "    [this] {\n"
                "      //\n"