Smarter wrapping logic for long parameter lists

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

@@ -18,6 +18,9 @@ package com.squareup.kotlinpoet
 /** Sentinel value that indicates that no user-provided package has been set.  */
 private val NO_PACKAGE = String()
 
+private const val DEFAULT_INDENT_LEVEL = 1
+private const val CONTINUATION_INDENT_LEVEL = 2
+
 private fun extractMemberName(part: String): String {
   require(Character.isJavaIdentifierStart(part[0])) { "not an identifier: $part" }
   for (i in 1..part.length) {
@@ -40,6 +43,7 @@ internal class CodeWriter constructor(
 ) {
   private val out = LineWrapper(out, indent, 100)
   private var indentLevel = 0
+  private var indentLevelIncrement = CONTINUATION_INDENT_LEVEL
 
   private var kdoc = false
   private var comment = false
@@ -249,7 +253,7 @@ internal class CodeWriter constructor(
           statementLine = -1
         }
 
-        "%W" -> out.wrappingSpace(indentLevel + 2)
+        "%W" -> out.wrappingSpace(indentLevel + indentLevelIncrement)
 
         else -> {
           // Handle deferred type.
@@ -275,8 +279,18 @@ internal class CodeWriter constructor(
     }
   }
 
+  fun openWrappingGroup() {
+    out.openWrappingGroup()
+    indentLevelIncrement = DEFAULT_INDENT_LEVEL
+  }
+
+  fun closeWrappingGroup() {
+    trailingNewline = out.closeWrappingGroup()
+    indentLevelIncrement = CONTINUATION_INDENT_LEVEL
+  }
+
   fun emitWrappingSpace() = apply {
-    out.wrappingSpace(indentLevel + 2)
+    out.wrappingSpace(indentLevel + indentLevelIncrement)
   }
 
   private fun emitStaticImportMember(canonical: String, part: String): Boolean {

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

@@ -65,7 +65,7 @@ class LambdaTypeName internal constructor(
       out.emitCode("%T.", it)
     }
 
-    parameters.emit(out)
+    parameters.emit(out, wrap = false)
     out.emitCode(" -> %T", returnType)
 
     if (nullable) {

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

@@ -26,71 +26,185 @@ internal class LineWrapper(
 ) {
   private var closed = false
 
-  /** Characters written since the last wrapping space that haven't yet been flushed.  */
-  private val buffer = StringBuilder()
-
   /** The number of characters since the most recent newline. Includes both out and the buffer.  */
   private var column = 0
 
-  /** -1 if we have no buffering; otherwise the number of spaces to write after wrapping.  */
-  private var indentLevel = -1
+  private var helper: BufferedLineWrapperHelper = DefaultLineWrapperHelper()
 
   /** Emit `s`. This may be buffered to permit line wraps to be inserted.  */
   fun append(s: String) {
     check(!closed) { "closed" }
 
-    if (indentLevel != -1) {
+    if (helper.isBuffering) {
       val nextNewline = s.indexOf('\n')
 
       // If s doesn't cause the current line to cross the limit, buffer it and return. We'll decide
       // whether or not we have to wrap it later.
       if (nextNewline == -1 && column + s.length <= columnLimit) {
-        buffer.append(s)
+        helper.buffer(s)
         column += s.length
         return
       }
 
       // Wrap if appending s would overflow the current line.
       val wrap = nextNewline == -1 || column + nextNewline > columnLimit
-      flush(wrap)
+      helper.flush(wrap)
     }
 
-    out.append(s)
+    helper.append(s)
     val lastNewline = s.lastIndexOf('\n')
     column = if (lastNewline != -1)
       s.length - lastNewline - 1 else
       column + s.length
   }
 
+  fun openWrappingGroup() {
+    check(!closed) { "closed" }
+
+    helper = GroupLineWrapperHelper()
+  }
+
   /** Emit either a space or a newline character.  */
   fun wrappingSpace(indentLevel: Int) {
     check(!closed) { "closed" }
 
-    if (this.indentLevel != -1) flush(false)
+    helper.wrappingSpace(indentLevel)
     this.column++
-    this.indentLevel = indentLevel
+  }
+
+  fun closeWrappingGroup(): Boolean {
+    check(!closed) { "closed" }
+
+    val wrapped = helper.close()
+    helper = DefaultLineWrapperHelper()
+    return wrapped
   }
 
   /** Flush any outstanding text and forbid future writes to this line wrapper.  */
   fun close() {
-    if (indentLevel != -1) flush(false)
+    helper.close()
     closed = true
   }
 
   /** Write the space followed by any buffered text that follows it.  */
-  private fun flush(wrap: Boolean) {
+  private fun flush(buffered: String, wrap: Boolean) {
     if (wrap) {
       out.append('\n')
-      for (i in 0 until indentLevel) {
+      for (i in 0 until helper.indentLevel) {
         out.append(indent)
       }
-      column = indentLevel * indent.length
-      column += buffer.length
+      column = helper.indentLevel * indent.length
+      column += buffered.length
     } else {
       out.append(' ')
     }
-    out.append(buffer)
-    buffer.delete(0, buffer.length)
-    indentLevel = -1
+    out.append(buffered)
+  }
+
+  /**
+   * Contract for helpers that handle buffering, post-processing and flushing of the input.
+   */
+  internal interface BufferedLineWrapperHelper {
+
+    val indentLevel: Int
+
+    val isBuffering get() = indentLevel != -1
+
+    /** Append to out, bypassing the buffer */
+    fun append(s: String): Appendable
+
+    /** Append to buffer */
+    fun buffer(s: String): Appendable
+
+    /**
+     * Indicates that a new wrapping space occurred in input.
+     *
+     * @param indentLevel Indentation level for the new line
+     */
+    fun wrappingSpace(indentLevel: Int)
+
+    /**
+     * Flush any buffered text.
+     *
+     * @param wrap `true` if buffer contents should be flushed a on new line
+     * */
+    fun flush(wrap: Boolean)
+
+    /**
+     * Flush and clear the buffer.
+     *
+     * @return `true` if input wrapped to new line
+     */
+    fun close(): Boolean
+  }
+
+  /** Flushes the buffer each time the wrapping space is encountered */
+  internal inner class DefaultLineWrapperHelper : BufferedLineWrapperHelper {
+
+    private val buffer = StringBuilder()
+
+    private var _indentLevel = -1
+
+    override val indentLevel get() = _indentLevel
+
+    override fun append(s: String): Appendable = out.append(s)
+
+    override fun buffer(s: String): Appendable = buffer.append(s)
+
+    override fun wrappingSpace(indentLevel: Int) {
+      if (isBuffering) flush(false)
+      _indentLevel = indentLevel
+    }
+
+    override fun flush(wrap: Boolean) {
+      flush(buffer.toString(), wrap)
+      buffer.delete(0, buffer.length)
+      _indentLevel = -1
+    }
+
+    override fun close(): Boolean {
+      if (isBuffering) flush(false)
+      return false
+    }
+  }
+
+  /**
+   * Holds multiple buffers and only flushes when the group is closed. If wrapping happened within
+   * a group - each buffer will be flushed on a new line.
+   */
+  internal inner class GroupLineWrapperHelper : BufferedLineWrapperHelper {
+
+    private val buffer = mutableListOf(StringBuilder())
+    private var wrapped = false
+
+    private var _indentLevel = -1
+
+    override val indentLevel get() = _indentLevel
+
+    override fun append(s: String): Appendable = buffer.last().append(s)
+
+    override fun buffer(s: String): Appendable = buffer.last().append(s)
+
+    override fun wrappingSpace(indentLevel: Int) {
+      _indentLevel = indentLevel
+      buffer += StringBuilder()
+    }
+
+    override fun flush(wrap: Boolean) {
+      wrapped = wrap
+    }
+
+    override fun close(): Boolean {
+      if (wrapped) buffer.last().append('\n')
+      buffer.forEachIndexed { index, segment ->
+        if (index == 0 && !wrapped) {
+          out.append(segment)
+        } else {
+          flush(segment.toString(), wrapped)
+        }
+      }
+      _indentLevel = -1
+      return wrapped
+    }
   }
 }

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

@@ -145,28 +145,15 @@ class ParameterSpec private constructor(builder: ParameterSpec.Builder) {
 
 internal fun List<ParameterSpec>.emit(
   codeWriter: CodeWriter,
+  wrap: Boolean = true,
   emitBlock: (ParameterSpec) -> Unit = { it.emit(codeWriter) }
 ) = with(codeWriter) {
-  val params = this@emit
   emit("(")
-  when (size) {
-    0 -> emit("")
-    1 -> emitBlock(params[0])
-    2 -> {
-      emitBlock(params[0])
-      emit(", ")
-      emitBlock(params[1])
-    }
-    else -> {
-      emit("\n")
-      indent(2)
-      forEachIndexed { index, parameter ->
-        if (index > 0) emit(",\n")
-        emitBlock(parameter)
-      }
-      unindent(2)
-      emit("\n")
-    }
+  if (wrap) codeWriter.openWrappingGroup()
+  forEachIndexed { index, parameter ->
+    if (index > 0) if (wrap) emitCode(",%W") else emit(", ")
+    emitBlock(parameter)
   }
+  if (wrap) codeWriter.closeWrappingGroup()
   emit(")")
 }

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

@@ -17,6 +17,7 @@ package com.squareup.kotlinpoet
 
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
+import java.io.Serializable
 
 class KotlinPoetTest {
   private val tacosPackage = "com.squareup.tacos"
@@ -107,11 +108,7 @@ class KotlinPoetTest {
         |import kotlin.Boolean
         |import kotlin.String
         |
-        |class Taco(
-        |        val cheese: String,
-        |        var cilantro: String,
-        |        lettuce: String
-        |) {
+        |class Taco(val cheese: String, var cilantro: String, lettuce: String) {
         |    val lettuce: String = lettuce.trim()
         |
         |    val onion: Boolean = true
@@ -365,11 +362,7 @@ class KotlinPoetTest {
       |import kotlin.String
       |import kotlin.Unit
       |
-      |fun ((
-      |        name: String,
-      |        Int,
-      |        age: Long
-      |) -> Unit).whatever(): Unit = Unit
+      |fun ((name: String, Int, age: Long) -> Unit).whatever(): Unit = Unit
       |""".trimMargin())
   }
 
@@ -580,4 +573,50 @@ class KotlinPoetTest {
       |}
       |""".trimMargin())
   }
+
+  @Test fun longParameterListWrapping() {
+    val source = FunSpec.builder("sum")
+        .addParameter(ParameterSpec.builder("a", Int::class).build())
+        .addParameter(ParameterSpec.builder("b", Int::class).build())
+        .addParameter(ParameterSpec.builder("c", Int::class).build())
+        .addParameter(ParameterSpec.builder("d", Int::class).build())
+        .addParameter(ParameterSpec.builder("e", Int::class).build())
+        .addParameter(ParameterSpec.builder("f", Int::class).build())
+        .addParameter(ParameterSpec.builder("g", Int::class).build())
+        .addStatement("return a + b + c")
+        .build()
+    assertThat(source.toString()).isEqualTo("""
+      |fun sum(
+      |    a: kotlin.Int,
+      |    b: kotlin.Int,
+      |    c: kotlin.Int,
+      |    d: kotlin.Int,
+      |    e: kotlin.Int,
+      |    f: kotlin.Int,
+      |    g: kotlin.Int
+      |) = a + b + c
+      |""".trimMargin())
+  }
+
+  @Test fun longLambdaParameterListWrapping() {
+    val source = FunSpec.builder("veryLongFunctionName")
+        .addParameter(ParameterSpec.builder(
+            "veryLongParameterName",
+            LambdaTypeName.get(
+                parameters = listOf(
+                    ParameterSpec.unnamed(Serializable::class),
+                    ParameterSpec.unnamed(Appendable::class),
+                    ParameterSpec.unnamed(Cloneable::class)),
+                returnType = Unit::class.asTypeName()))
+            .build())
+        .addParameter("i", Int::class)
+        .addStatement("return %T", Unit::class)
+        .build()
+    assertThat(source.toString()).isEqualTo("""
+      |fun veryLongFunctionName(
+      |    veryLongParameterName: (java.io.Serializable, java.lang.Appendable, kotlin.Cloneable) -> kotlin.Unit,
+      |    i: kotlin.Int
+      |) = kotlin.Unit
+      |""".trimMargin())
+  }
 }