From fbcf83aa7656a97b2a9cf98d470afaf0870eb59e Mon Sep 17 00:00:00 2001
From: Daniel Chao <dan.chao@apple.com>
Date: Tue, 28 Oct 2025 10:46:19 -0700
Subject: [PATCH] Adjust formatting of argument lists (#1260)

This changes code so that multiple lambda arguments makes the whole
argument list wrap.

Improves the readability of code like:

```
foo
  .toMap(
    (it) -> makeSomeKey(it),
    (it) -> makeSomeValue(it)
  )
```
---
 .../main/kotlin/org/pkl/formatter/Builder.kt  | 24 +++++++++++++++++--
 .../input/method-call-trailing-lambdas.pkl    | 18 ++++++++++++++
 .../output/method-call-trailing-lambdas.pkl   | 19 +++++++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 pkl-formatter/src/test/files/FormatterSnippetTests/input/method-call-trailing-lambdas.pkl
 create mode 100644 pkl-formatter/src/test/files/FormatterSnippetTests/output/method-call-trailing-lambdas.pkl

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 6cd48e1c..43127fcd 100644
--- a/pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt
+++ b/pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt
@@ -625,14 +625,34 @@ internal class Builder(sourceText: String, private val grammarVersion: GrammarVe
     }
   }
 
+  /**
+   * Only considered trailing lamdba if there is only one lambda/new expr/amends expr in the list.
+   *
+   * E.g. avoid formatting `toMap()` weirdly:
+   * ```
+   * foo.toMap(
+   *   (it) -> makeSomeKey(it),
+   *   (it) -> makeSomeValue(it),
+   * )
+   * ```
+   */
   private fun hasTrailingLambda(argList: Node): Boolean {
     val children = argList.firstProperChild()?.children ?: return false
+    var seenArg = false
+    var ret = false
     for (i in children.lastIndex downTo 0) {
       val child = children[i]
       if (!child.isProper()) continue
-      return child.type in SAME_LINE_EXPRS
+      if (child.type in SAME_LINE_EXPRS) {
+        if (seenArg) {
+          return false
+        } else {
+          seenArg = true
+          ret = true
+        }
+      }
     }
-    return false
+    return ret
   }
 
   private fun pairArguments(nodes: List<Node>): List<Node> {
diff --git a/pkl-formatter/src/test/files/FormatterSnippetTests/input/method-call-trailing-lambdas.pkl b/pkl-formatter/src/test/files/FormatterSnippetTests/input/method-call-trailing-lambdas.pkl
new file mode 100644
index 00000000..c1d8759d
--- /dev/null
+++ b/pkl-formatter/src/test/files/FormatterSnippetTests/input/method-call-trailing-lambdas.pkl
@@ -0,0 +1,18 @@
+// multiple lambda arguments means argument list wraps
+res1 = foo.toMap(
+  (elem) -> elem.toString(),
+  (elem) -> elem.toString().split("").map((str) -> str + "hello").toSet()
+)
+
+// new object and trailing lambda means argument list wraps
+res2 = foo.doFoo(
+  new Listing { 1; 2; 3 },
+  (elem) -> elem.toString().split("").map((str) -> str + "hello").toSet()
+)
+
+// single trailing lambda is broken on its own
+res4 = foo.foldLeft(
+  "heeeeeeelloooooooooooooooooooooooooooooooooooooooooooooo",
+  (elem) -> elem.foooooooooooooooooooooooooooooooooooooooooooooooooooo
+)
+
diff --git a/pkl-formatter/src/test/files/FormatterSnippetTests/output/method-call-trailing-lambdas.pkl b/pkl-formatter/src/test/files/FormatterSnippetTests/output/method-call-trailing-lambdas.pkl
new file mode 100644
index 00000000..e5c55402
--- /dev/null
+++ b/pkl-formatter/src/test/files/FormatterSnippetTests/output/method-call-trailing-lambdas.pkl
@@ -0,0 +1,19 @@
+// multiple lambda arguments means argument list wraps
+res1 =
+  foo.toMap(
+    (elem) -> elem.toString(),
+    (elem) -> elem.toString().split("").map((str) -> str + "hello").toSet(),
+  )
+
+// new object and trailing lambda means argument list wraps
+res2 =
+  foo.doFoo(
+    new Listing { 1; 2; 3 },
+    (elem) -> elem.toString().split("").map((str) -> str + "hello").toSet(),
+  )
+
+// single trailing lambda is broken on its own
+res4 =
+  foo.foldLeft("heeeeeeelloooooooooooooooooooooooooooooooooooooooooooooo", (elem) ->
+    elem.foooooooooooooooooooooooooooooooooooooooooooooooooooo
+  )