Fix a bug caused by too long return expressions (#1256)

* Fix a bug caused by too long return expressions

* Fix a bug caused by too long return expressions

* Auto-identify extra spaces before and after return

* Auto-identify extra spaces before and after return

* Fix possible array out-of-bounds

* Making judgement easy

* add comments

* Update kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt

Co-authored-by: Egor Andreevich <github@egorand.dev>

* Update kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt

Co-authored-by: Egor Andreevich <github@egorand.dev>

* Update FunSpec.kt

* format if required

* spotless

Co-authored-by: yuejunyu <yuejunyu.0@bytedance.com>
Co-authored-by: Egor Andreevich <github@egorand.dev>

2 files changed, 42 insertions, 1 deletions

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
index 2266f6dd..28c5ca23 100644
--- a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
+++ b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
@@ -114,7 +114,7 @@ public class FunSpec private constructor(
     } else if (!isEmptySetter) {
       codeWriter.emitCode("·{\n")
       codeWriter.indent()
-      codeWriter.emitCode(body, ensureTrailingNewline = true)
+      codeWriter.emitCode(body.returnsWithoutLinebreak(), ensureTrailingNewline = true)
       codeWriter.unindent()
       codeWriter.emit("}\n")
     } else {
@@ -245,6 +245,20 @@ public class FunSpec private constructor(
     return null
   }
 
+  private fun CodeBlock.returnsWithoutLinebreak(): CodeBlock {
+    val returnWithSpace = RETURN_EXPRESSION_BODY_PREFIX_SPACE.formatParts[0]
+    val returnWithNbsp = RETURN_EXPRESSION_BODY_PREFIX_NBSP.formatParts[0]
+    var originCodeBlockBuilder: CodeBlock.Builder? = null
+    for ((i, formatPart) in formatParts.withIndex()) {
+      if (formatPart.startsWith(returnWithSpace)) {
+        val builder = originCodeBlockBuilder ?: toBuilder()
+        originCodeBlockBuilder = builder
+        builder.formatParts[i] = formatPart.replaceFirst(returnWithSpace, returnWithNbsp)
+      }
+    }
+    return originCodeBlockBuilder?.build() ?: this
+  }
+
   override fun equals(other: Any?): Boolean {
     if (this === other) return true
     if (other == null) return false
diff --git a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
index df93cfcb..a3fe6386 100644
--- a/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
+++ b/kotlinpoet/src/test/java/com/squareup/kotlinpoet/FunSpecTest.kt
@@ -206,6 +206,33 @@ class FunSpecTest {
     )
   }
 
+  @Test fun returnsLongExpression() {
+    val funSpec = FunSpec.builder("foo")
+      .returns(String::class)
+      .addStatement("val placeholder = 1")
+      .addStatement("return %S", "Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong")
+      .build()
+    val sb = StringBuilder()
+    // The FunSpec#toString columnLimit is Integer.MAX_VALUE,
+    // It will not cause problems with returns long expressions.
+    CodeWriter(sb).use {
+      funSpec.emit(
+        codeWriter = it,
+        enclosingName = null,
+        implicitModifiers = setOf(KModifier.PUBLIC),
+        includeKdocTags = false
+      )
+    }
+    assertThat(sb.toString()).isEqualTo(
+      """
+      |public fun foo(): kotlin.String {
+      |  val placeholder = 1
+      |  return "Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong"
+      |}
+      |""".trimMargin()
+    )
+  }
+
   @Test fun functionParamWithKdoc() {
     val funSpec = FunSpec.builder("foo")
       .addParameter(