Fix a bug caused by too long return expressions

kotlinpoet/src/main/java/com/squareup/kotlinpoet/CodeBlock.kt

@@ -458,6 +458,8 @@ public class CodeBlock private constructor(
     private const val TYPE_NAME = 2
     private val NO_ARG_PLACEHOLDERS = setOf("⇥", "⇤", "«", "»")
     internal val EMPTY = CodeBlock(emptyList(), emptyList())
+    internal const val RETURN_WITH_SPACE = "return "
+    internal const val RETURN_WITH_NBSP = "return·"
 
     @JvmStatic public fun of(format: String, vararg args: Any?): CodeBlock =
       Builder().add(format, *args).build()

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,25 @@ public class FunSpec private constructor(
     return null
   }
 
+  private fun CodeBlock.returnsWithoutLinebreak(): CodeBlock {
+    val originCodeBlockBuilder = toBuilder()
+    val returnWithSpace = CodeBlock.RETURN_WITH_SPACE
+    originCodeBlockBuilder.formatParts.clear()
+    formatParts.mapTo(originCodeBlockBuilder.formatParts) { formatPart ->
+      if (formatPart.isEmpty()) return@mapTo formatPart
+      var startReturnIndex = 0
+      while (startReturnIndex < formatPart.length && formatPart[startReturnIndex] == ' ') startReturnIndex++
+      var endIndex = startReturnIndex + returnWithSpace.length
+      if (endIndex > formatPart.length) return@mapTo formatPart
+      if (formatPart.substring(startReturnIndex, endIndex) != returnWithSpace) {
+        return@mapTo formatPart
+      }
+      while (endIndex < formatPart.length && formatPart[endIndex] == ' ') endIndex++
+      CodeBlock.RETURN_WITH_NBSP + formatPart.substring(endIndex, formatPart.length)
+    }
+    return originCodeBlockBuilder.build()
+  }
+
   override fun equals(other: Any?): Boolean {
     if (this === other) return true
     if (other == null) return false
@@ -558,8 +577,8 @@ public class FunSpec private constructor(
     internal val String.isConstructor get() = this == CONSTRUCTOR
     internal val String.isAccessor get() = this.isOneOf(GETTER, SETTER)
 
-    private val RETURN_EXPRESSION_BODY_PREFIX_SPACE = CodeBlock.of("return ")
-    private val RETURN_EXPRESSION_BODY_PREFIX_NBSP = CodeBlock.of("return·")
+    private val RETURN_EXPRESSION_BODY_PREFIX_SPACE = CodeBlock.of(CodeBlock.RETURN_WITH_SPACE)
+    private val RETURN_EXPRESSION_BODY_PREFIX_NBSP = CodeBlock.of(CodeBlock.RETURN_WITH_NBSP)
     private val THROW_EXPRESSION_BODY_PREFIX_SPACE = CodeBlock.of("throw ")
     private val THROW_EXPRESSION_BODY_PREFIX_NBSP = CodeBlock.of("throw·")
 

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

@@ -206,6 +206,62 @@ class FunSpecTest {
     )
   }
 
+  @Test fun returnsLongExpression() {
+    val funSpec = FunSpec.builder("foo")
+      .returns(String::class)
+      .addStatement("val placeholder = 1")
+      .addStatement("return     \"Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong\"")
+      .build()
+    val sb = StringBuilder()
+    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 returnsLongExpressions() {
+    val funSpec = FunSpec.builder("foo")
+      .returns(String::class)
+      .addStatement("val placeholder = 1")
+      .beginControlFlow("if (placeholder == 0)")
+      .addStatement("  return     \"LooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooongExpressionInControlFlow\"")
+      .endControlFlow()
+      .addStatement("return \"Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong\"")
+      .build()
+    val sb = StringBuilder()
+    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
+      |  if (placeholder == 0) {
+      |    return "LooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooongExpressionInControlFlow"
+      |  }
+      |  return "Loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooong"
+      |}
+      |""".trimMargin()
+    )
+  }
+
   @Test fun functionParamWithKdoc() {
     val funSpec = FunSpec.builder("foo")
       .addParameter(