From 35da0d54d321b8160b21bfea58290cb5e881cc46 Mon Sep 17 00:00:00 2001
From: zsub <49520432+zsqw123@users.noreply.github.com>
Date: Wed, 1 Jun 2022 22:28:48 +0800
Subject: [PATCH] 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>
---
 .../java/com/squareup/kotlinpoet/FunSpec.kt   | 16 ++++++++++-
 .../com/squareup/kotlinpoet/FunSpecTest.kt    | 27 +++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt b/kotlinpoet/src/main/java/com/squareup/kotlinpoet/FunSpec.kt
index 2266f6dd79..28c5ca23e2 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 df93cfcb6f..a3fe63863e 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(