Adjust formatting of qualified access chains (#1252)

This adjusts formatting of qualified access chains so that leading
dot calls are kept in the same line if possible.

pkl-formatter/src/main/kotlin/org/pkl/formatter/Builder.kt

@@ -131,6 +131,7 @@ internal class Builder(sourceText: String) {
       NodeType.AMENDS_EXPR -> formatNewExpr(node)
       NodeType.NEW_HEADER -> formatNewHeader(node)
       NodeType.UNQUALIFIED_ACCESS_EXPR -> formatUnqualifiedAccessExpression(node)
+      NodeType.QUALIFIED_ACCESS_EXPR -> formatQualifiedAccessExpression(node)
       NodeType.BINARY_OP_EXPR -> formatBinaryOpExpr(node)
       NodeType.FUNCTION_LITERAL_EXPR -> formatFunctionLiteralExpr(node)
       NodeType.FUNCTION_LITERAL_BODY -> formatFunctionLiteralBody(node)
@@ -193,7 +194,7 @@ internal class Builder(sourceText: String) {
     return if (prefixes.isEmpty()) {
       res
     } else {
-      val sep = getSeparator(prefixes.last(), nodes.first())
+      val sep = getSeparator(prefixes.last(), nodes.first(), spaceOrLine())
       Nodes(formatGeneric(prefixes, spaceOrLine()) + listOf(sep, res))
     }
   }
@@ -242,6 +243,123 @@ internal class Builder(sourceText: String) {
     }
   }
 
+  /**
+   * Special cases when formatting qualified access:
+   *
+   * Case 1: Dot calls followed by closing method call: wrap after the opening paren.
+   *
+   * ```
+   * foo.bar.baz(new {
+   *   qux = 1
+   * })
+   * ```
+   *
+   * Case 2: Dot calls, then method calls: group the leading access together.
+   *
+   * ```
+   * foo.bar
+   *   .baz(new {
+   *     qux = 1
+   *   })
+   *   .baz()
+   * ```
+   *
+   * Case 3: If there are multiple lambdas present, always force a newline.
+   *
+   * ```
+   * foo
+   *   .map((it) -> it + 1)
+   *   .filter((it) -> it.isEven)
+   * ```
+   */
+  private fun formatQualifiedAccessExpression(node: Node): FormatNode {
+    var lambdaCount = 0
+    var methodCallCount = 0
+    var indexBeforeFirstMethodCall = 0
+    val flat = mutableListOf<Node>()
+
+    fun gatherFacts(current: Node) {
+      for (child in current.children) {
+        if (child.type == NodeType.QUALIFIED_ACCESS_EXPR) {
+          gatherFacts(child)
+        } else {
+          flat.add(child)
+          when {
+            isMethodCall(child) -> {
+              methodCallCount++
+              if (hasFunctionLiteral(child, 2)) {
+                lambdaCount++
+              }
+            }
+            methodCallCount == 0 -> {
+              indexBeforeFirstMethodCall = flat.lastIndex
+            }
+          }
+        }
+      }
+    }
+
+    gatherFacts(node)
+
+    val separator: (Node, Node) -> FormatNode? = { prev, next ->
+      when {
+        prev.type == NodeType.OPERATOR -> null
+        next.type == NodeType.OPERATOR -> if (lambdaCount > 1) forceLine() else line()
+        else -> spaceOrLine()
+      }
+    }
+
+    val nodes =
+      when {
+        methodCallCount == 1 && isMethodCall(flat.lastProperNode()!!) -> {
+          // lift argument list into its own node
+          val (callChain, argsList) = splitFunctionCallNode(flat)
+          val leadingNodes = indentAfterFirstNewline(formatGeneric(callChain, separator), true)
+          val trailingNodes = formatGeneric(argsList, separator)
+          val sep = getBaseSeparator(callChain.last(), argsList.first())
+          if (sep != null) {
+            leadingNodes + sep + trailingNodes
+          } else {
+            leadingNodes + trailingNodes
+          }
+        }
+        methodCallCount > 0 && indexBeforeFirstMethodCall > 0 -> {
+          val leading = flat.subList(0, indexBeforeFirstMethodCall)
+          val trailing = flat.subList(indexBeforeFirstMethodCall, flat.size)
+          val leadingNodes = indentAfterFirstNewline(formatGeneric(leading, separator), true)
+          val trailingNodes = formatGeneric(trailing, separator)
+          leadingNodes + line() + trailingNodes
+        }
+        else -> formatGeneric(flat, separator)
+      }
+
+    val shouldGroup = node.children.size == flat.size
+    return Group(newId(), indentAfterFirstNewline(nodes, shouldGroup))
+  }
+
+  /**
+   * Split a function call node to extract its identifier into the leading group. For example,
+   * `foo.bar(5)` becomes: leading gets `foo.bar`, rest gets `(5)`.
+   */
+  private fun splitFunctionCallNode(nodes: List<Node>): Pair<List<Node>, List<Node>> {
+    assert(nodes.isNotEmpty())
+    val lastNode = nodes.last()
+    val argListIdx = lastNode.children.indexOfFirst { it.type == NodeType.ARGUMENT_LIST }
+    val leading = nodes.subList(0, nodes.lastIndex) + lastNode.children.subList(0, argListIdx)
+    val trailing = lastNode.children.subList(argListIdx, lastNode.children.size)
+    return leading to trailing
+  }
+
+  private fun isMethodCall(node: Node): Boolean {
+    if (node.type != NodeType.UNQUALIFIED_ACCESS_EXPR) return false
+    for (child in node.children) {
+      if (child.type == NodeType.ARGUMENT_LIST) {
+        return true
+      }
+    }
+    return false
+  }
+
   private fun formatAmendsExtendsClause(node: Node): FormatNode {
     val prefix = formatGeneric(node.children.dropLast(1), spaceOrLine())
     // string constant
@@ -688,7 +806,7 @@ internal class Builder(sourceText: String) {
           val prevNoNewlines = noNewlines
           noNewlines = true
           val elems = cursor.takeUntilBefore { it.isTerminal(")") }
-          getSeparator(prev!!, elems.first(), { _, _ -> null })?.let { add(it) }
+          getBaseSeparator(prev!!, elems.first())?.let { add(it) }
           val formatted =
             try {
               formatGeneric(elems, null)
@@ -697,7 +815,7 @@ internal class Builder(sourceText: String) {
               formatGeneric(elems, null)
             }
           addAll(formatted)
-          getSeparator(elems.last(), cursor.peek(), { _, _ -> null })?.let { add(it) }
+          getBaseSeparator(elems.last(), cursor.peek())?.let { add(it) }
           noNewlines = prevNoNewlines
           isInStringInterpolation = false
           continue
@@ -848,26 +966,21 @@ internal class Builder(sourceText: String) {
 
   private fun formatBinaryOpExpr(node: Node): FormatNode {
     val flat = flattenBinaryOperatorExprs(node)
-    val callChainSize = flat.count { it.isOperator(".", "?.") }
-    val hasMultipleLambdas = callChainSize > 1 && flat.hasMoreThan(1) { hasFunctionLiteral(it, 2) }
     val nodes =
       formatGeneric(flat) { prev, next ->
         if (prev.type == NodeType.OPERATOR) {
           when (prev.text()) {
-            ".",
-            "?." -> null
             "-" -> spaceOrLine()
             else -> Space
           }
         } else if (next.type == NodeType.OPERATOR) {
           when (next.text()) {
-            ".",
-            "?." -> if (hasMultipleLambdas) forceLine() else line()
             "-" -> Space
             else -> spaceOrLine()
           }
         } else spaceOrLine()
       }
+
     val shouldGroup = node.children.size == flat.size
     return Group(newId(), indentAfterFirstNewline(nodes, shouldGroup))
   }
@@ -1095,7 +1208,7 @@ internal class Builder(sourceText: String) {
     val nodes = children.subList(index, children.size)
     val res = mutableListOf<FormatNode>()
     res += formatGeneric(prefixes, spaceOrLine())
-    res += getSeparator(prefixes.last(), nodes.first())
+    res += getSeparator(prefixes.last(), nodes.first(), spaceOrLine())
     res += groupFn(nodes)
     return res
   }
@@ -1103,12 +1216,8 @@ internal class Builder(sourceText: String) {
   private fun getImportUrl(node: Node): String =
     node.findChildByType(NodeType.STRING_CHARS)!!.text().drop(1).dropLast(1)
 
-  private fun getSeparator(
-    prev: Node,
-    next: Node,
-    separator: FormatNode = spaceOrLine(),
-  ): FormatNode {
-    return getSeparator(prev, next) { _, _ -> separator }!!
+  private fun getSeparator(prev: Node, next: Node, separator: FormatNode): FormatNode {
+    return getBaseSeparator(prev, next) ?: separator
   }
 
   private fun getSeparator(
@@ -1116,6 +1225,10 @@ internal class Builder(sourceText: String) {
     next: Node,
     separatorFn: (Node, Node) -> FormatNode?,
   ): FormatNode? {
+    return getBaseSeparator(prev, next) ?: separatorFn(prev, next)
+  }
+
+  private fun getBaseSeparator(prev: Node, next: Node): FormatNode? {
     return when {
       prevNode?.type == NodeType.LINE_COMMENT -> {
         if (prev.linesBetween(next) > 1) {
@@ -1124,6 +1237,7 @@ internal class Builder(sourceText: String) {
           mustForceLine()
         }
       }
+
       hasTrailingAffix(prev, next) -> Space
       prev.type == NodeType.DOC_COMMENT -> mustForceLine()
       prev.type == NodeType.ANNOTATION -> forceLine()
@@ -1134,17 +1248,21 @@ internal class Builder(sourceText: String) {
           mustForceLine()
         }
       }
+
       prev.type == NodeType.BLOCK_COMMENT ->
         if (prev.linesBetween(next) > 0) forceSpaceyLine() else Space
+
       next.type in EMPTY_SUFFIXES ||
         prev.isTerminal("[", "!", "@", "[[") ||
-        next.isTerminal("]", "?", ",") -> null
+        next.isTerminal("]", "?", ",") -> Empty
+
       prev.isTerminal("class", "function", "new") ||
         next.isTerminal("=", "{", "->", "class", "function") ||
         next.type == NodeType.OBJECT_BODY ||
         prev.type == NodeType.MODIFIER_LIST -> Space
+
       next.type == NodeType.DOC_COMMENT -> TWO_NEWLINES
-      else -> separatorFn(prev, next)
+      else -> null
     }
   }
 
@@ -1260,9 +1378,6 @@ internal class Builder(sourceText: String) {
   private fun Node.isTerminal(vararg texts: String): Boolean =
     type == NodeType.TERMINAL && text(source) in texts
 
-  private fun Node.isOperator(vararg texts: String): Boolean =
-    type == NodeType.OPERATOR && text(source) in texts
-
   private fun newId(): Int {
     return id++
   }
@@ -1293,6 +1408,13 @@ internal class Builder(sourceText: String) {
     return null
   }
 
+  private fun List<Node>.lastProperNode(): Node? {
+    for (i in lastIndex downTo 0) {
+      if (this[i].isProper()) return this[i]
+    }
+    return null
+  }
+
   // returns true if this node is not an affix or terminal
   private fun Node.isProper(): Boolean = !type.isAffix && type != NodeType.TERMINAL
 
@@ -1305,18 +1427,6 @@ 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
-  }
-
   class PeekableIterator<T>(private val iterator: Iterator<T>) : Iterator<T> {
     private var peek: T? = null
 

pkl-formatter/src/test/files/FormatterSnippetTests/input/expr-chain-grouping.pkl

@@ -0,0 +1,31 @@
+res1 = foo.bar(new {
+  foo = 1
+})
+
+res2 = foo.bar.baz(new {
+  foo =1
+})
+
+res3 = foo.bar(3).baz(new {
+  qux = 3
+})
+
+res4 = foooooooooooooooooooooooooooooooooooooo.foooooooooooooooooooooooooooooooooooooo.foooooooooooooooooooooooooooooooooooooo.bar(5)
+
+res5 = foooooooooooooo.foooooooooooooo.foooooooooooooo.foooooooooooooo.foooooooooooooo(4).foooooooooooooo.bar(5)
+
+res6 = foooooooooooooooooooooooooooooooooooooo.foooooooooooooooooooooooooooooooooooooo.foooooooooooooooooooooooooooooooooooooo.foooooooooooooo(4).foooooooooooooo.bar(5)
+
+res7 = foooooooooooooo.foooooooooooooo.foooooooooooooo.foooooooooooooo /* hello */ .foooooooooooooo(4).foooooooooooooo.bar(5)
+
+res8 = foooooooooooooo.foooooooooooooo.foooooooooooooo().foooooooooooooo.foooooooooooooo(4).foooooooooooooo.bar(5)
+
+res9 = foo.bar(new { 
+  foo = 1
+  bar = 2
+}).qux
+
+res10 = foo.bar /* some comment */ (5)
+
+res11 = foo.bar /* some comment */.baz(new { foooooooooooooooooooooooooo = 1; baaaaaaaaaaaaaaaaaaaaar = 2 }).buz
+

pkl-formatter/src/test/files/FormatterSnippetTests/output/expr-chain-grouping.pkl

@@ -0,0 +1,65 @@
+res1 =
+  foo.bar(new {
+    foo = 1
+  })
+
+res2 =
+  foo.bar.baz(new {
+    foo = 1
+  })
+
+res3 =
+  foo
+    .bar(3)
+    .baz(new {
+      qux = 3
+    })
+
+res4 =
+  foooooooooooooooooooooooooooooooooooooo
+    .foooooooooooooooooooooooooooooooooooooo
+    .foooooooooooooooooooooooooooooooooooooo
+    .bar(5)
+
+res5 =
+  foooooooooooooo.foooooooooooooo.foooooooooooooo.foooooooooooooo
+    .foooooooooooooo(4)
+    .foooooooooooooo
+    .bar(5)
+
+res6 =
+  foooooooooooooooooooooooooooooooooooooo
+    .foooooooooooooooooooooooooooooooooooooo
+    .foooooooooooooooooooooooooooooooooooooo
+    .foooooooooooooo(4)
+    .foooooooooooooo
+    .bar(5)
+
+res7 =
+  foooooooooooooo.foooooooooooooo.foooooooooooooo.foooooooooooooo /* hello */
+    .foooooooooooooo(4)
+    .foooooooooooooo
+    .bar(5)
+
+res8 =
+  foooooooooooooo.foooooooooooooo
+    .foooooooooooooo()
+    .foooooooooooooo
+    .foooooooooooooo(4)
+    .foooooooooooooo
+    .bar(5)
+
+res9 =
+  foo
+    .bar(new {
+      foo = 1
+      bar = 2
+    })
+    .qux
+
+res10 = foo.bar /* some comment */ (5)
+
+res11 =
+  foo.bar /* some comment */
+    .baz(new { foooooooooooooooooooooooooo = 1; baaaaaaaaaaaaaaaaaaaaar = 2 })
+    .buz

pkl-parser/src/main/java/org/pkl/parser/GenericParser.java

@@ -679,6 +679,10 @@ public class GenericParser {
           ff(children);
           expect(Token.RBRACK, children, "unexpectedToken", "]");
         }
+        case DOT, QDOT -> {
+          nodeType = NodeType.QUALIFIED_ACCESS_EXPR;
+          children.add(parseUnqualifiedAccessExpr());
+        }
         case NON_NULL -> nodeType = NodeType.NON_NULL_EXPR;
         default -> children.add(parseExpr(expectation, nextMinPrec));
       }
@@ -719,6 +723,16 @@ public class GenericParser {
     };
   }
 
+  private Node parseUnqualifiedAccessExpr() {
+    var children = new ArrayList<Node>();
+    children.add(parseIdentifier());
+    if (lookahead() == Token.LPAREN && noSemicolonInbetween() && _lookahead.newLinesBetween == 0) {
+      ff(children);
+      children.add(parseArgumentList());
+    }
+    return new Node(NodeType.UNQUALIFIED_ACCESS_EXPR, children);
+  }
+
   private Node parseExprAtom(@Nullable String expectation) {
     var expr =
         switch (lookahead) {
@@ -884,16 +898,7 @@ public class GenericParser {
           case FLOAT -> new Node(NodeType.FLOAT_LITERAL_EXPR, next().span);
           case STRING_START -> parseSingleLineStringLiteralExpr();
           case STRING_MULTI_START -> parseMultiLineStringLiteralExpr();
-          case IDENTIFIER -> {
-            var children = new ArrayList<Node>();
-            children.add(parseIdentifier());
-            if (lookahead == Token.LPAREN
-                && noSemicolonInbetween()
-                && _lookahead.newLinesBetween == 0) {
-              children.add(parseArgumentList());
-            }
-            yield new Node(NodeType.UNQUALIFIED_ACCESS_EXPR, children);
-          }
+          case IDENTIFIER -> parseUnqualifiedAccessExpr();
           case EOF ->
               throw parserError(
                   ErrorMessages.create("unexpectedEndOfFile"), prev().span.stopSpan());

pkl-parser/src/main/java/org/pkl/parser/syntax/generic/NodeType.java

@@ -110,6 +110,7 @@ public enum NodeType {
   SUPER_SUBSCRIPT_EXPR(NodeKind.EXPR),
   SUPER_ACCESS_EXPR(NodeKind.EXPR),
   SUBSCRIPT_EXPR(NodeKind.EXPR),
+  QUALIFIED_ACCESS_EXPR(NodeKind.EXPR),
   IF_EXPR(NodeKind.EXPR),
   IF_HEADER,
   IF_CONDITION,

pkl-parser/src/test/kotlin/org/pkl/parser/GenericSexpRenderer.kt

@@ -34,10 +34,6 @@ class GenericSexpRenderer(code: String) {
       renderUnionType(node)
       return
     }
-    if (node.type == NodeType.BINARY_OP_EXPR && binopName(node).endsWith("ualifiedAccessExpr")) {
-      renderQualifiedAccess(node)
-      return
-    }
     doRender(name(node), collectChildren(node))
   }
 
@@ -143,6 +139,10 @@ class GenericSexpRenderer(code: String) {
       NodeType.AMENDS_CLAUSE -> "extendsOrAmendsClause"
       NodeType.TYPEALIAS -> "typeAlias"
       NodeType.STRING_ESCAPE -> "stringChars"
+      NodeType.QUALIFIED_ACCESS_EXPR -> {
+        val op = node.findChildByType(NodeType.OPERATOR)!!
+        if (op.text(source) == ".") "qualifiedAccessExpr" else "nullableQualifiedAccessExpr"
+      }
       NodeType.READ_EXPR -> {
         val terminal = node.children.find { it.type == NodeType.TERMINAL }!!.text(source)
         when (terminal) {

pkl-parser/src/test/kotlin/org/pkl/parser/SexpRenderer.kt

@@ -28,14 +28,14 @@ class SexpRenderer {
   private var tab = ""
   private var buf = StringBuilder()
 
-  fun render(mod: org.pkl.parser.syntax.Module): String {
+  fun render(mod: Module): String {
     renderModule(mod)
     val res = buf.toString()
     reset()
     return res
   }
 
-  fun renderModule(mod: org.pkl.parser.syntax.Module) {
+  fun renderModule(mod: Module) {
     buf.append(tab)
     buf.append("(module")
     val oldTab = increaseTab()
@@ -515,13 +515,11 @@ class SexpRenderer {
     tab = oldTab
   }
 
-  fun renderQualifiedAccessExpr(expr: QualifiedAccessExpr) {
+  fun renderUnqualifiedAccessExprOfQualified(expr: QualifiedAccessExpr) {
     buf.append(tab)
-    buf.append(if (expr.isNullable) "(nullableQualifiedAccessExpr" else "(qualifiedAccessExpr")
+    buf.append("(unqualifiedAccessExpr")
     val oldTab = increaseTab()
     buf.append('\n')
-    renderExpr(expr.expr)
-    buf.append('\n')
     buf.append(tab)
     buf.append("(identifier)")
     if (expr.argumentList !== null) {
@@ -532,6 +530,18 @@ class SexpRenderer {
     tab = oldTab
   }
 
+  fun renderQualifiedAccessExpr(expr: QualifiedAccessExpr) {
+    buf.append(tab)
+    buf.append(if (expr.isNullable) "(nullableQualifiedAccessExpr" else "(qualifiedAccessExpr")
+    val oldTab = increaseTab()
+    buf.append('\n')
+    renderExpr(expr.expr)
+    buf.append('\n')
+    renderUnqualifiedAccessExprOfQualified(expr)
+    buf.append(')')
+    tab = oldTab
+  }
+
   fun renderSuperAccessExpr(expr: SuperAccessExpr) {
     buf.append(tab)
     buf.append("(superAccessExpr")