From 781eb2f025b6ef82ef35967f84b7308cba080d6a Mon Sep 17 00:00:00 2001
From: Omer Strulovich <ostrulovich@fb.com>
Date: Wed, 13 Apr 2022 14:02:02 -0700
Subject: [PATCH] Do not try to make block like syntax for expression with more
 than 1 lambda

Summary: Avoiding the block like syntax once more than one lambda shows up looks much better for many complicated chains.

Reviewed By: hick209

Differential Revision: D35618238

fbshipit-source-id: acb835a68439d052680fa35053545bdf2f810c2b
---
 .../ktfmt/format/KotlinInputAstVisitor.kt     | 11 +++---
 .../facebook/ktfmt/format/FormatterTest.kt    | 36 +++++++++++++++----
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
index c4633f2a..cdb3344b 100644
--- a/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
+++ b/core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt
@@ -487,8 +487,9 @@ class KotlinInputAstVisitor(
    */
   private fun emitQualifiedExpression(expression: KtExpression) {
     val parts = breakIntoParts(expression)
-    val hasTrailingLambda = parts.last().isLambda()
-    val groupingInfos = computeGroupingInfo(parts, hasTrailingLambda)
+    // whether we want to make a lambda look like a block, this make Kotlin DSLs look as expected
+    val useBlockLikeLambdaStyle = parts.last().isLambda() && parts.count { it.isLambda() } == 1
+    val groupingInfos = computeGroupingInfo(parts, useBlockLikeLambdaStyle)
     builder.block(expressionBreakIndent) {
       val nameTag = genSym() // allows adjusting arguments indentation if a break will be made
       for ((index, ktExpression) in parts.withIndex()) {
@@ -515,7 +516,7 @@ class KotlinInputAstVisitor(
                 builder.close()
               }
               // close group due to last lambda to allow block-like style in `as.forEach { ... }`
-              val isTrailingLambda = hasTrailingLambda && index == parts.size - 1
+              val isTrailingLambda = useBlockLikeLambdaStyle && index == parts.size - 1
               if (isTrailingLambda) {
                 builder.close()
               }
@@ -609,7 +610,7 @@ class KotlinInputAstVisitor(
    */
   private fun computeGroupingInfo(
       parts: List<KtExpression>,
-      hasTrailingLambda: Boolean
+      useBlockLikeLambdaStyle: Boolean
   ): List<GroupingInfo> {
     val groupingInfos = List(parts.size) { GroupingInfo() }
     var lastIndexToOpen = 0
@@ -639,7 +640,7 @@ class KotlinInputAstVisitor(
         }
       }
     }
-    if (hasTrailingLambda) {
+    if (useBlockLikeLambdaStyle) {
       // a trailing lambda adds a group that we stop before emitting the lambda
       groupingInfos[0].groupOpenCount++
     }
diff --git a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
index ac893e3a..875a47ea 100644
--- a/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
+++ b/core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt
@@ -652,13 +652,16 @@ class FormatterTest {
       |        foo
       |      }
       |
-      |  // All method calls with lambdas are on one line (because they fit),
-      |  // even though the apply() at the end requires a line break.
+      |  // All method calls with lambdas could fit, but we avoid a block like syntax
+      |  // and break to one call per line
       |  val items =
-      |      items.map { it + 1 }.filter { it > 0 }.apply {
-      |        //
-      |        foo
-      |      }
+      |      items
+      |          .map { it + 1 }
+      |          .filter { it > 0 }
+      |          .apply {
+      |            //
+      |            foo
+      |          }
       |
       |  // the lambda is indented properly with the break before it
       |  val items =
@@ -692,6 +695,27 @@ class FormatterTest {
       |""".trimMargin(),
           deduceMaxWidth = true)
 
+  @Test
+  fun `when two lambdas are in a chain, avoid block syntax`() =
+      assertFormatted(
+          """
+      |class Foo : Bar() {
+      |  fun doIt() {
+      |    fruit.onlyBananas().forEach { banana ->
+      |      val seeds = banana.find { it.type == SEED }
+      |      println(seeds)
+      |    }
+      |
+      |    fruit
+      |        .filter { isBanana(it, Bananas.types) }
+      |        .forEach { banana ->
+      |          val seeds = banana.find { it.type == SEED }
+      |          println(seeds)
+      |        }
+      |  }
+      |}
+      |""".trimMargin())
+
   @Test
   fun `indent parameters after a break when there's a lambda afterwards`() =
       assertFormatted(