J2K: preserving line breaks between enum members from original code

j2k/src/org/jetbrains/kotlin/j2k/CodeBuilder.kt

@@ -21,9 +21,13 @@ import com.intellij.psi.*
 import com.intellij.psi.javadoc.PsiDocComment
 import org.jetbrains.kotlin.j2k.ast.CommentsAndSpacesInheritance
 import org.jetbrains.kotlin.j2k.ast.Element
+import org.jetbrains.kotlin.j2k.ast.SpacesInheritance
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
-import java.util.*
+import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.ArrayList
+import java.util.HashSet
+import java.util.LinkedHashSet
 import kotlin.platform.platformName
 
 fun<T> CodeBuilder.append(generators: Collection<() -> T>, separator: String, prefix: String = "", suffix: String = ""): CodeBuilder {
@@ -117,31 +121,25 @@ class CodeBuilder(private val topElement: PsiElement?) {
         }
 
         val notInsideElements = HashSet<PsiElement>()
-        val prefixElements = ArrayList<PsiElement>(1)
+        var prefix = Prefix.Empty
         val postfixElements = ArrayList<PsiElement>(1)
         for ((prototype, inheritance) in element.prototypes!!) {
             assert(prototype !is PsiComment)
             assert(prototype !is PsiWhiteSpace)
             if (!topElement.isAncestor(prototype)) continue
-            prefixElements.collectPrefixElements(prototype, inheritance, notInsideElements)
+            prefix += collectPrefixElements(prototype, inheritance, notInsideElements)
             postfixElements.collectPostfixElements(prototype, inheritance, notInsideElements)
         }
 
-        commentsAndSpacesUsed.addAll(prefixElements)
+        if (prefix.lineBreaksBefore > 0) {
-        commentsAndSpacesUsed.addAll(postfixElements)
+            val lineBreaksToAdd = prefix.lineBreaksBefore - builder.trailingLineBreakCount()
-
+            for (_ in 1..lineBreaksToAdd) {
-        for ((i, e) in prefixElements.withIndex()) {
+                append("\n", false)
-            if (i == 0 && e is PsiWhiteSpace) {
-                val blankLines = e.newLinesCount() - 1
-                for (_ in 1..blankLines) {
-                    append("\n", false)
-                }
-            }
-            else {
-                appendCommentOrWhiteSpace(e)
             }
         }
 
+        prefix.elements.forEach { appendCommentOrWhiteSpace(it) }
+
         element.generateCode(this)
 
         // scan for all comments inside which are not yet used in the text and put them here to not loose any comment from code
@@ -164,29 +162,52 @@ class CodeBuilder(private val topElement: PsiElement?) {
         return this
     }
 
-    private fun MutableList<PsiElement>.collectPrefixElements(element: PsiElement,
+    private data class Prefix(val elements: Collection<PsiElement>, val lineBreaksBefore: Int) {
-                                                              inheritance: CommentsAndSpacesInheritance,
+        fun plus(other: Prefix) = Prefix(elements + other.elements, Math.max(lineBreaksBefore, other.lineBreaksBefore))
-                                                              notInsideElements: MutableSet<PsiElement>) {
+
+        companion object {
+            val Empty = Prefix(emptyList(), 0)
+        }
+    }
+
+    private fun collectPrefixElements(
+            element: PsiElement,
+            inheritance: CommentsAndSpacesInheritance,
+            notInsideElements: MutableSet<PsiElement>
+    ): Prefix {
         val before = ArrayList<PsiElement>(1).collectCommentsAndSpacesBefore(element)
         val atStart = ArrayList<PsiElement>(1).collectCommentsAndSpacesAtStart(element)
         notInsideElements.addAll(atStart)
 
-        if (!inheritance.blankLinesBefore && !inheritance.commentsBefore) return
+        if (inheritance.spacesBefore == SpacesInheritance.NONE && !inheritance.commentsBefore) return Prefix.Empty
 
         val firstSpace = before.lastOrNull() as? PsiWhiteSpace
-        if (!inheritance.commentsBefore) { // take only first whitespace
+        var lineBreaks = 0
-            if (firstSpace != null) {
+        if (firstSpace != null) {
-                add(firstSpace)
+            lineBreaks = firstSpace.newLinesCount()
+            when (inheritance.spacesBefore) {
+                SpacesInheritance.NONE -> lineBreaks = 0
+
+                SpacesInheritance.LINE_BREAKS -> commentsAndSpacesUsed.add(firstSpace)
+
+                SpacesInheritance.BLANK_LINES_ONLY -> {
+                    commentsAndSpacesUsed.add(firstSpace)
+                    if (lineBreaks == 1) lineBreaks = 0
+                }
             }
-            return
         }
 
-        if (!inheritance.blankLinesBefore && firstSpace != null) {
+        if (!inheritance.commentsBefore) { // take only whitespace
+            return Prefix(emptyList(), lineBreaks)
+        }
+
+        if (firstSpace != null) {
             before.remove(before.lastIndex)
         }
 
-        addAll(before.reverse())
+        val elements = before.reverse() + atStart
-        addAll(atStart)
+        commentsAndSpacesUsed.addAll(elements)
+        return Prefix(elements, lineBreaks)
     }
 
     private fun MutableList<PsiElement>.collectPostfixElements(element: PsiElement, inheritance: CommentsAndSpacesInheritance, notInsideElements: MutableSet<PsiElement>) {
@@ -205,6 +226,9 @@ class CodeBuilder(private val topElement: PsiElement?) {
 
         addAll(atEnd.reverse())
         addAll(after)
+
+        commentsAndSpacesUsed.addAll(atEnd)
+        commentsAndSpacesUsed.addAll(after)
     }
 
     private fun MutableList<PsiElement>.collectCommentsAndSpacesBefore(element: PsiElement): MutableList<PsiElement> {
@@ -253,14 +277,14 @@ class CodeBuilder(private val topElement: PsiElement?) {
     private fun MutableList<PsiElement>.collectCommentsAndSpacesAtStart(element: PsiElement): MutableList<PsiElement> {
         var child = element.getFirstChild()
         while(child != null) {
-            if (child!!.isCommentOrSpace()) {
+            if (child.isCommentOrSpace()) {
-                if (child !in commentsAndSpacesUsed) add(child!!) else break
+                if (child !in commentsAndSpacesUsed) add(child) else break
             }
-            else if (!child!!.isEmptyElement()) {
+            else if (!child.isEmptyElement()) {
-                collectCommentsAndSpacesAtStart(child!!)
+                collectCommentsAndSpacesAtStart(child)
                 break
             }
-            child = child!!.getNextSibling()
+            child = child.getNextSibling()
         }
         return this
     }
@@ -268,14 +292,14 @@ class CodeBuilder(private val topElement: PsiElement?) {
     private fun MutableList<PsiElement>.collectCommentsAndSpacesAtEnd(element: PsiElement): MutableList<PsiElement> {
         var child = element.getLastChild()
         while(child != null) {
-            if (child!!.isCommentOrSpace()) {
+            if (child.isCommentOrSpace()) {
-                if (child !in commentsAndSpacesUsed) add(child!!) else break
+                if (child !in commentsAndSpacesUsed) add(child) else break
             }
-            else if (!child!!.isEmptyElement()) {
+            else if (!child.isEmptyElement()) {
-                collectCommentsAndSpacesAtEnd(child!!)
+                collectCommentsAndSpacesAtEnd(child)
                 break
             }
-            child = child!!.getPrevSibling()
+            child = child.getPrevSibling()
         }
         return this
     }
@@ -289,5 +313,11 @@ class CodeBuilder(private val topElement: PsiElement?) {
     private fun PsiWhiteSpace.newLinesCount() = StringUtil.getLineBreakCount(getText()!!)
 
     private fun PsiWhiteSpace.hasNewLines() = StringUtil.containsLineBreak(getText()!!)
+
+    private fun CharSequence.trailingLineBreakCount(): Int {
+        val index = ((length()-1 downTo 0).firstOrNull { val c = this[it]; c != '\n' && c != '\r' } ?: -1) + 1
+        if (index == length()) return 0
+        return StringUtil.getLineBreakCount(subSequence(index, length()))
+    }
 }
 

j2k/src/org/jetbrains/kotlin/j2k/ConstructorConverter.kt

@@ -214,7 +214,7 @@ class ConstructorConverter(
                                   if (isVal(converter.referenceSearcher, field)) Parameter.VarValModifier.Val else Parameter.VarValModifier.Var,
                                   converter.convertAnnotations(parameter) + converter.convertAnnotations(field),
                                   accessModifiers,
-                                  default).assignPrototypes(listOf(parameter, field), CommentsAndSpacesInheritance(blankLinesBefore = false))
+                                  default).assignPrototypes(listOf(parameter, field), CommentsAndSpacesInheritance(spacesBefore = SpacesInheritance.NONE))
                     }
                 },
                 correctCodeConverter = { correct() })

j2k/src/org/jetbrains/kotlin/j2k/Converter.kt

@@ -244,7 +244,7 @@ class Converter private constructor(
 
     private fun convertAnnotationType(psiClass: PsiClass): Class {
         val paramModifiers = Modifiers(listOf(Modifier.PUBLIC)).assignNoPrototype()
-        val noBlankLinesInheritance = CommentsAndSpacesInheritance(blankLinesBefore = false)
+        val noBlankLinesInheritance = CommentsAndSpacesInheritance(spacesBefore = SpacesInheritance.NONE)
         val annotationMethods = psiClass.getMethods().filterIsInstance<PsiAnnotationMethod>()
         val (methodsNamedValue, otherMethods) = annotationMethods.partition { it.getName() == "value" }
 
@@ -313,13 +313,14 @@ class Converter private constructor(
         }
 
         val name = correction?.identifier ?: field.declarationIdentifier()
-        val converted = if (field is PsiEnumConstant) {
+        if (field is PsiEnumConstant) {
             val argumentList = field.getArgumentList()
             val params = deferredElement { codeConverter ->
                 ExpressionList(codeConverter.convertExpressions(argumentList?.getExpressions() ?: arrayOf<PsiExpression>())).assignPrototype(argumentList)
             }
             val body = field.getInitializingClass()?.let { convertAnonymousClassBody(it) }
-            EnumConstant(name, annotations, modifiers, params, body)
+            return EnumConstant(name, annotations, modifiers, params, body)
+                    .assignPrototype(field, CommentsAndSpacesInheritance(spacesBefore = SpacesInheritance.LINE_BREAKS))
         }
         else {
             val isVal = isVal(referenceSearcher, field)
@@ -330,7 +331,7 @@ class Converter private constructor(
 
             addUsageProcessing(FieldToPropertyProcessing(field, correction?.name ?: field.getName(), propertyType.isNullable))
 
-            Property(name,
+            return Property(name,
                   annotations,
                   modifiers,
                   propertyType,
@@ -338,9 +339,8 @@ class Converter private constructor(
                   isVal,
                   typeToDeclare != null,
                   shouldGenerateDefaultInitializer(referenceSearcher, field),
-                  if (correction != null) correction.setterAccess else modifiers.accessModifier())
+                  if (correction != null) correction.setterAccess else modifiers.accessModifier()).assignPrototype(field)
         }
-        return converted.assignPrototype(field)
     }
 
     public fun variableTypeToDeclare(variable: PsiVariable, specifyAlways: Boolean, canChangeType: Boolean): Type? {
@@ -556,7 +556,7 @@ class Converter private constructor(
 
     public fun convertModifiers(owner: PsiModifierListOwner): Modifiers {
         return Modifiers(MODIFIERS_MAP.filter { owner.hasModifierProperty(it.first) }.map { it.second })
-                .assignPrototype(owner.getModifierList(), CommentsAndSpacesInheritance(blankLinesBefore = false))
+                .assignPrototype(owner.getModifierList(), CommentsAndSpacesInheritance(spacesBefore = SpacesInheritance.NONE))
     }
 
     public fun convertAnonymousClassBody(anonymousClass: PsiAnonymousClass): AnonymousClassBody {

j2k/src/org/jetbrains/kotlin/j2k/ast/ClassBody.kt

@@ -45,7 +45,7 @@ class ClassBody (
         else {
             val (constants, otherMembers) = membersFiltered.partition { it is EnumConstant }
 
-            builder.append(constants, ",\n")
+            builder.append(constants, ", ")
 
             if (otherMembers.isNotEmpty() || companionObjectMembers.isNotEmpty()) {
                 builder.append(";\n")

j2k/src/org/jetbrains/kotlin/j2k/ast/Constructors.kt

@@ -43,7 +43,7 @@ class PrimaryConstructor(
 
         // assign prototypes later because we don't know yet whether the body is empty or not
         converter.addPostUnfoldDeferredElementsAction {
-            val inheritance = CommentsAndSpacesInheritance(blankLinesBefore = false, commentsAfter = body!!.isEmpty, commentsInside = body.isEmpty)
+            val inheritance = CommentsAndSpacesInheritance(spacesBefore = SpacesInheritance.NONE, commentsAfter = body!!.isEmpty, commentsInside = body.isEmpty)
             signature.assignPrototypesFrom(this, inheritance)
         }
 

j2k/src/org/jetbrains/kotlin/j2k/ast/Element.kt

@@ -45,10 +45,16 @@ fun <TElement: Element> TElement.assignPrototypesFrom(element: Element, inherita
 
 data class PrototypeInfo(val element: PsiElement, val commentsAndSpacesInheritance: CommentsAndSpacesInheritance)
 
-data class CommentsAndSpacesInheritance(val blankLinesBefore: Boolean = true,
+enum class SpacesInheritance {
-                                        val commentsBefore: Boolean = true,
+    NONE, BLANK_LINES_ONLY, LINE_BREAKS
-                                        val commentsAfter: Boolean = true,
+}
-                                        val commentsInside: Boolean = true)
+
+data class CommentsAndSpacesInheritance(
+        val spacesBefore: SpacesInheritance = SpacesInheritance.BLANK_LINES_ONLY,
+        val commentsBefore: Boolean = true,
+        val commentsAfter: Boolean = true,
+        val commentsInside: Boolean = true
+)
 
 fun Element.canonicalCode(): String {
     val builder = CodeBuilder(null)

j2k/testData/fileOrElement/annotations/annotationInterface3.kt

@@ -2,8 +2,7 @@
 annotation class Anon(public val value: String) {
 
     public enum class E {
-        A,
+        A, B
-        B
     }
 
     companion object {

j2k/testData/fileOrElement/enum/fieldsWithPrimaryPrivateConstructor.kt

@@ -1,7 +1,3 @@
 enum class Color private constructor(public val code: Int) {
-    WHITE(21),
+    WHITE(21), BLACK(22), RED(23), YELLOW(24), BLUE(25)
-    BLACK(22),
-    RED(23),
-    YELLOW(24),
-    BLUE(25)
 }

j2k/testData/fileOrElement/enum/overrideToString.kt

@@ -1,9 +1,5 @@
 enum class Color {
-    WHITE,
+    WHITE, BLACK, RED, YELLOW, BLUE;
-    BLACK,
-    RED,
-    YELLOW,
-    BLUE;
     override fun toString(): String {
         return "COLOR"
     }

j2k/testData/fileOrElement/enum/runnableImplementation.kt

@@ -1,9 +1,5 @@
 enum class Color : Runnable {
-    WHITE,
+    WHITE, BLACK, RED, YELLOW, BLUE;
-    BLACK,
-    RED,
-    YELLOW,
-    BLUE;
 
     override fun run() {
         println("name()=" + name() + ", toString()=" + toString())

j2k/testData/fileOrElement/enum/typeSafeEnum.kt

@@ -1,6 +1,3 @@
 enum class Coin {
-    PENNY,
+    PENNY, NICKEL, DIME, QUARTER
-    NICKEL,
-    DIME,
-    QUARTER
 }