From 3a29ea8998a2ef3d55793e77b223adf8d5a45bcb Mon Sep 17 00:00:00 2001
From: Daniel Chao <dan.chao@apple.com>
Date: Thu, 9 Oct 2025 10:16:13 -0700
Subject: [PATCH] Format lambda chains (#1235)

This changes the formatter to only force line on call chains
with multiple lambdas

With this change, this is kept as a single line:

```
foo.bar.map((it) -> it + 1)
```
---
 .../src/main/kotlin/org/pkl/formatter/Builder.kt | 16 ++++++++++++++--
 .../FormatterSnippetTests/input/expr-chain.pkl   |  2 ++
 .../FormatterSnippetTests/output/expr-chain.pkl  |  2 ++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt b/pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt
index 0e07007e..02e6f0dd 100644
--- a/pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt
+++ b/pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt
@@ -803,7 +803,7 @@ internal class Builder(sourceText: String) {
   private fun formatBinaryOpExpr(node: Node): FormatNode {
     val flat = flattenBinaryOperatorExprs(node)
     val callChainSize = flat.count { it.isOperator(".", "?.") }
-    val hasLambda = callChainSize > 1 && flat.any { hasFunctionLiteral(it, 2) }
+    val hasMultipleLambdas = callChainSize > 1 && flat.hasMoreThan(1) { hasFunctionLiteral(it, 2) }
     val nodes =
       formatGeneric(flat) { prev, next ->
         if (prev.type == NodeType.OPERATOR) {
@@ -816,7 +816,7 @@ internal class Builder(sourceText: String) {
         } else if (next.type == NodeType.OPERATOR) {
           when (next.text()) {
             ".",
-            "?." -> if (hasLambda) ForceLine else Line
+            "?." -> if (hasMultipleLambdas) ForceLine else Line
             "-" -> Space
             else -> SpaceOrLine
           }
@@ -1231,6 +1231,18 @@ internal class Builder(sourceText: String) {
     }
   }
 
+  private inline fun <T> Iterable<T>.hasMoreThan(num: Int, predicate: (T) -> Boolean): Boolean {
+    if (this is Collection && isEmpty()) return false
+    var count = 0
+    for (element in this) {
+      if (predicate(element)) count++
+      if (count > num) {
+        return true
+      }
+    }
+    return false
+  }
+
   companion object {
     private val ABSOLUTE_URL_REGEX = Regex("""\w+:.*""")
 
diff --git a/pkl-formatter/src/test/files/FormatterSnippetTests/input/expr-chain.pkl b/pkl-formatter/src/test/files/FormatterSnippetTests/input/expr-chain.pkl
index 1b9ff4b1..6161f242 100644
--- a/pkl-formatter/src/test/files/FormatterSnippetTests/input/expr-chain.pkl
+++ b/pkl-formatter/src/test/files/FormatterSnippetTests/input/expr-chain.pkl
@@ -3,3 +3,5 @@ res = myList.map((it) -> it.partition).filter((it) -> someList.contains(it))
 res2 = myList.map(lambda1).filter(lambda2)
 
 res3 = myList.map((it) -> it.partition)
+
+res4 = myList.foo.map((it) -> it.partition)
diff --git a/pkl-formatter/src/test/files/FormatterSnippetTests/output/expr-chain.pkl b/pkl-formatter/src/test/files/FormatterSnippetTests/output/expr-chain.pkl
index ef7be664..fa462c36 100644
--- a/pkl-formatter/src/test/files/FormatterSnippetTests/output/expr-chain.pkl
+++ b/pkl-formatter/src/test/files/FormatterSnippetTests/output/expr-chain.pkl
@@ -6,3 +6,5 @@ res =
 res2 = myList.map(lambda1).filter(lambda2)
 
 res3 = myList.map((it) -> it.partition)
+
+res4 = myList.foo.map((it) -> it.partition)