Don't leak an expression body's indent into later declarations

docs/changelog.md

@@ -14,6 +14,7 @@ Change Log
  * Fix: `get` and `set` operator function names are no longer escaped with backticks. (#1869)
  * Fix: Don't special case varargs in `KSAnnotation.toAnnotationSpec`. (#2360)
  * Fix: Emit context parameters after annotations in `FunSpec` and `PropertySpec`. (#2374)
+ * Fix: An expression body no longer leaks indentation into later declarations. (#1421)
 
 ## Version 2.3.0
 

kotlinpoet/src/jvmMain/kotlin/com/squareup/kotlinpoet/CodeBlock.kt

@@ -119,8 +119,34 @@ private constructor(internal val formatParts: List<String>, internal val args: L
   /**
    * Returns a copy of the code block without leading and trailing no-arg placeholders (`⇥`, `⇤`,
    * `«`, `»`).
+   *
+   * With [keepBalanced], placeholders are taken back from either end until nothing left in the
+   * block is unmatched, which would otherwise leak the indent or the statement into whatever is
+   * emitted next. A body ending inside `withIndent` loses the `⇤` that closed it, e.g.:
+   * ```
+   * ["return ", "%S", "⇥", ".trim()", "⇤"]
+   *   trim()                    -> ["return ", "%S", "⇥", ".trim()"]
+   *   trim(keepBalanced = true) -> ["return ", "%S", "⇥", ".trim()", "⇤"]
+   * ```
+   *
+   * Both ends move when both are unmatched, e.g.:
+   * ```
+   * ["⇥", "body1", "⇤", "⇥", "body2", "⇤"]
+   *   trim()                    -> ["body1", "⇤", "⇥", "body2"]
+   *   trim(keepBalanced = true) -> ["⇥", "body1", "⇤", "⇥", "body2", "⇤"]
+   * ```
+   *
+   * A block that was already unbalanced before trimming is left to the plain behavior, since taking
+   * one end back cannot balance it and moves the following declaration further out, e.g.:
+   * ```
+   * ["⇥", "return ", "%S", "⇤", "⇥", ".trim()"]
+   *   trim()                    -> ["return ", "%S", "⇤", "⇥", ".trim()"]
+   *   trim(keepBalanced = true) -> ["return ", "%S", "⇤", "⇥", ".trim()"]
+   * ```
+   *
+   * [trimTrailingNewLine] relies on the default, which always drops both runs.
    */
-  internal fun trim(): CodeBlock {
+  internal fun trim(keepBalanced: Boolean = false): CodeBlock {
     var start = 0
     var end = formatParts.size
     while (start < end && formatParts[start] in NO_ARG_PLACEHOLDERS) {
@@ -129,6 +155,79 @@ private constructor(internal val formatParts: List<String>, internal val args: L
     while (start < end && formatParts[end - 1] in NO_ARG_PLACEHOLDERS) {
       end--
     }
+    if (keepBalanced && (start > 0 || end < formatParts.size)) {
+      // Both stripped runs are placeholders by definition, so they are the first `start` and the
+      // last `formatParts.size - end` entries of this list, and the kept range is what sits
+      // between them. Everything below walks placeholders rather than format parts, which keeps
+      // the balance check to a single pass over the parts.
+      val placeholders = CharArray(formatParts.size)
+      var placeholderCount = 0
+      for (formatPart in formatParts) {
+        if (formatPart.length == 1 && formatPart[0].isSingleCharNoArgPlaceholder) {
+          placeholders[placeholderCount++] = formatPart[0]
+        }
+      }
+      val keptEnd = placeholderCount - (formatParts.size - end)
+      if (isBalanced(placeholders, placeholderCount)) {
+        // Track the lowest point reached, not just the running total. Two unrelated halves can
+        // cancel out to a total of zero and still emit an unindent ahead of its indent.
+        var indentLow = 0
+        var indentTotal = 0
+        var statementLow = 0
+        var statementTotal = 0
+        for (i in start..<keptEnd) {
+          when (placeholders[i]) {
+            '⇥' -> indentTotal++
+            '⇤' -> indentTotal--
+            '«' -> statementTotal++
+            '»' -> statementTotal--
+          }
+          if (indentTotal < indentLow) indentLow = indentTotal
+          if (statementTotal < statementLow) statementLow = statementTotal
+        }
+        var first = start
+        while (first > 0 && (indentLow < 0 || statementLow < 0)) {
+          when (placeholders[first - 1]) {
+            '⇥' -> {
+              indentLow++
+              indentTotal++
+            }
+            '⇤' -> {
+              indentLow--
+              indentTotal--
+            }
+            '«' -> {
+              statementLow++
+              statementTotal++
+            }
+            '»' -> {
+              statementLow--
+              statementTotal--
+            }
+          }
+          first--
+          if (indentLow >= 0 && statementLow >= 0) {
+            start = first
+            break
+          }
+        }
+        var last = keptEnd
+        while (last < placeholderCount && (indentTotal > 0 || statementTotal > 0)) {
+          when (placeholders[last]) {
+            '⇥' -> indentTotal++
+            '⇤' -> indentTotal--
+            '«' -> statementTotal++
+            '»' -> statementTotal--
+          }
+          if (indentTotal < 0 || statementTotal < 0) break
+          last++
+          if (indentTotal == 0 && statementTotal == 0) {
+            end += last - keptEnd
+            break
+          }
+        }
+      }
+    }
     return when {
       start > 0 || end < formatParts.size -> CodeBlock(formatParts.subList(start, end), args)
       else -> this
@@ -148,6 +247,21 @@ private constructor(internal val formatParts: List<String>, internal val args: L
 
   internal fun hasStatements() = formatParts.any { "«" in it }
 
+  private fun isBalanced(placeholders: CharArray, count: Int): Boolean {
+    var indent = 0
+    var statement = 0
+    for (i in 0..<count) {
+      when (placeholders[i]) {
+        '⇥' -> indent++
+        '⇤' -> indent--
+        '«' -> statement++
+        '»' -> statement--
+      }
+      if (indent < 0 || statement < 0) return false
+    }
+    return indent == 0 && statement == 0
+  }
+
   internal fun hasUnmatchedClosingStatement(): Boolean {
     var openCount = 0
     for (formatPart in formatParts) {

kotlinpoet/src/jvmMain/kotlin/com/squareup/kotlinpoet/FunSpec.kt

@@ -243,7 +243,7 @@ private constructor(
   }
 
   private fun CodeBlock.asExpressionBody(): CodeBlock? {
-    val codeBlock = this.trim()
+    val codeBlock = this.trim(keepBalanced = true)
 
     // If after trimming there are unmatched closing statement symbols, we can't have an expression
     // body.

kotlinpoet/src/jvmTest/kotlin/com/squareup/kotlinpoet/CodeBlockTest.kt

@@ -494,6 +494,13 @@ class CodeBlockTest {
     assertThat(CodeBlock.of("«»⇥⇤").trim()).isEqualTo(CodeBlock.of(""))
   }
 
+  @Test
+  fun trimKeepsAStatementAndItsIndentTogether() {
+    val codeBlock = CodeBlock.of("«⇥return %S⇤».trim()", "taco")
+
+    assertThat(codeBlock.trim(keepBalanced = true)).isEqualTo(codeBlock)
+  }
+
   @Test
   fun replaceSimple() {
     assertThat(CodeBlock.of("%%⇥%%").replaceAll("%%", "")).isEqualTo(CodeBlock.of("⇥"))
@@ -631,6 +638,23 @@ class CodeBlockTest {
       )
   }
 
+  @Test
+  fun ensureEndsWithNewLineKeepsArgsWhenBlockEndsInIndent() {
+    val codeBlock =
+      CodeBlock.builder().add("%S", "taco\n").indent().add("\nmore").unindent().build()
+
+    assertThat(codeBlock.ensureEndsWithNewLine().toString())
+      .isEqualTo(
+        """
+        |""${'"'}
+        ||taco
+        ||""${'"'}.trimMargin()
+        |  more
+        |"""
+          .trimMargin()
+      )
+  }
+
   @Test
   fun `N escapes keywords`() {
     val funSpec = FunSpec.builder("object").build()