New J2K: rework Java code formatting collecting

#KT-33687 fixed

nj2k/src/org/jetbrains/kotlin/nj2k/FormattingCollector.kt

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.nj2k
+
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiKeyword
+import com.intellij.psi.PsiWhiteSpace
+import com.intellij.psi.javadoc.PsiDocComment
+import org.jetbrains.kotlin.idea.j2k.IdeaDocCommentConverter
+import org.jetbrains.kotlin.nj2k.tree.JKComment
+import org.jetbrains.kotlin.nj2k.tree.JKFormattingOwner
+
+class FormattingCollector {
+    private val commentCache = mutableMapOf<PsiElement, JKComment>()
+
+    fun takeFormattingFrom(
+        element: JKFormattingOwner,
+        psi: PsiElement?,
+        saveLineBreaks: Boolean,
+        takeTrailingComments: Boolean,
+        takeLeadingComments: Boolean
+    ) {
+        if (psi == null) return
+        val (leftTokens, rightTokens) = psi.collectComments(takeTrailingComments, takeLeadingComments)
+        element.trailingComments += leftTokens
+        element.leadingComments += rightTokens
+        if (saveLineBreaks) {
+            element.hasLeadingLineBreak = psi.hasLineBreakAfter()
+            element.hasTrailingLineBreak = psi.hasLineBreakBefore()
+        }
+    }
+
+    fun takeLineBreaksFrom(element: JKFormattingOwner, psi: PsiElement?) {
+        if (psi == null) return
+        element.hasLeadingLineBreak = psi.hasLineBreakAfter()
+        element.hasTrailingLineBreak = psi.hasLineBreakBefore()
+    }
+
+    private fun PsiElement.asComment(): JKComment? {
+        if (this in commentCache) return commentCache.getValue(this)
+        val token = when (this) {
+            is PsiDocComment -> JKComment(
+                IdeaDocCommentConverter.convertDocComment(
+                    this
+                )
+            )
+            is PsiComment -> JKComment(text)
+            else -> null
+        } ?: return null
+        commentCache[this] = token
+        return token
+    }
+
+    private fun Sequence<PsiElement>.toComments(): List<JKComment> =
+        takeWhile { it is PsiComment || it is PsiWhiteSpace || it.text == ";" }
+            .mapNotNull { it.asComment() }
+            .toList()
+
+    fun PsiElement.leadingCommentsWithParent(): Sequence<JKComment> {
+        val innerElements = leadingComments()
+        return (if (innerElements.lastOrNull()?.nextSibling == null && this is PsiKeyword)
+            innerElements + parent?.leadingComments().orEmpty()
+        else innerElements).mapNotNull { it.asComment() }
+    }
+
+    private fun PsiElement.trailingCommentsWithParent(): Sequence<JKComment> {
+        val innerElements = trailingComments()
+        return (if (innerElements.firstOrNull()?.prevSibling == null && this is PsiKeyword)
+            innerElements + parent?.trailingComments().orEmpty()
+        else innerElements).mapNotNull { it.asComment() }
+    }
+
+    private fun PsiElement.isNonCodeElement() =
+        this is PsiComment || this is PsiWhiteSpace || textMatches(";") || textLength == 0
+
+    private fun PsiElement.leadingComments() =
+        generateSequence(nextSibling) { it.nextSibling }
+            .takeWhile { it.isNonCodeElement() }
+
+    private fun PsiElement.trailingComments() =
+        generateSequence(prevSibling) { it.prevSibling }
+            .takeWhile { it.isNonCodeElement() }
+
+    private fun PsiElement.hasLineBreakBefore() = trailingComments().any { it is PsiWhiteSpace && it.textContains('\n') }
+    private fun PsiElement.hasLineBreakAfter() = leadingComments().any { it is PsiWhiteSpace && it.textContains('\n') }
+
+    private fun PsiElement.collectComments(
+        takeTrailingComments: Boolean,
+        takeLeadingComments: Boolean
+    ): Pair<List<JKComment>, List<JKComment>> {
+        val leftInnerTokens = children.asSequence().toComments().asReversed()
+        val rightInnerTokens = when {
+            children.isEmpty() -> emptyList()
+            else -> generateSequence(children.last()) { it.prevSibling }
+                .toComments()
+                .asReversed()
+        }
+
+        val leftComments = (leftInnerTokens + if (takeTrailingComments) trailingCommentsWithParent() else emptySequence()).asReversed()
+        val rightComments = rightInnerTokens + if (takeLeadingComments) leadingCommentsWithParent() else emptySequence()
+        return leftComments to rightComments
+    }
+}

nj2k/src/org/jetbrains/kotlin/nj2k/JavaToJKTreeBuilder.kt

@@ -35,13 +35,11 @@ import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.asJava.elements.KtLightField
 import org.jetbrains.kotlin.asJava.elements.KtLightMethod
 import org.jetbrains.kotlin.idea.caches.lightClasses.KtLightClassForDecompiledDeclaration
-import org.jetbrains.kotlin.idea.j2k.IdeaDocCommentConverter
 import org.jetbrains.kotlin.idea.j2k.content
 import org.jetbrains.kotlin.idea.refactoring.fqName.getKotlinFqName
 import org.jetbrains.kotlin.j2k.ReferenceSearcher
 import org.jetbrains.kotlin.j2k.ast.Nullability
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.nj2k.symbols.*
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.JKLiteralExpression.LiteralType.*
@@ -52,7 +50,6 @@ import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isExtensionDeclaration
 import org.jetbrains.kotlin.utils.KotlinExceptionWithAttachments
-import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 
@@ -68,11 +65,36 @@ class JavaToJKTreeBuilder constructor(
     }
 
     private val expressionTreeMapper = ExpressionTreeMapper()
-
     private val referenceSearcher: ReferenceSearcher = converterServices.oldServices.referenceSearcher
-
     private val declarationMapper = DeclarationMapper(expressionTreeMapper)
 
+    private val formattingCollector = FormattingCollector()
+
+    // we don't want to capture comments of previous declaration/statement
+    private fun PsiElement.takeLeadingCommentsNeeded() =
+        this !is PsiMember && this !is PsiStatement
+
+    private fun <T : JKFormattingOwner> T.withFormattingFrom(
+        psi: PsiElement?,
+        assignLineBreaks: Boolean = false,
+        takeTrailingComments: Boolean = true,
+        takeLeadingComments: Boolean = psi?.takeLeadingCommentsNeeded() ?: false
+    ): T = with(formattingCollector) {
+        takeFormattingFrom(this@withFormattingFrom, psi, assignLineBreaks, takeTrailingComments, takeLeadingComments)
+        this@withFormattingFrom
+    }
+
+    private fun <O : JKFormattingOwner> O.withLineBreaksFrom(psi: PsiElement?) = with(formattingCollector) {
+        takeLineBreaksFrom(this@withLineBreaksFrom, psi)
+        this@withLineBreaksFrom
+    }
+
+    private fun <O : JKFormattingOwner> O.withLeadingCommentsWithParent(psi: PsiElement?) = with(formattingCollector) {
+        if (psi == null) return@with this@withLeadingCommentsWithParent
+        this@withLeadingCommentsWithParent.leadingComments += psi.leadingCommentsWithParent()
+        return this@withLeadingCommentsWithParent
+    }
+
     private fun PsiJavaFile.toJK(): JKFile =
         JKFile(
             packageStatement?.toJK() ?: JKPackageDeclaration(JKNameIdentifier("")),
@@ -83,15 +105,15 @@ class JavaToJKTreeBuilder constructor(
     private fun PsiImportList?.toJK(saveImports: Boolean): JKImportList =
         JKImportList(this?.allImportStatements?.mapNotNull { it.toJK(saveImports) }.orEmpty()).also { importList ->
             val innerComments = this?.collectDescendantsOfType<PsiComment>()?.map { comment ->
-                JKCommentElement(comment.text)
+                JKComment(comment.text)
             }.orEmpty()
-            importList.leftNonCodeElements += innerComments
+            importList.trailingComments += innerComments
         }
 
     private fun PsiPackageStatement.toJK(): JKPackageDeclaration =
         JKPackageDeclaration(JKNameIdentifier(packageName))
             .also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
                 symbolProvider.provideUniverseSymbol(this, it)
             }
 
@@ -118,14 +140,14 @@ class JavaToJKTreeBuilder constructor(
 
         return JKImportStatement(JKNameIdentifier(name))
             .also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
     }
 
     private fun PsiIdentifier?.toJK(): JKNameIdentifier =
         this?.let {
             JKNameIdentifier(it.text).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         } ?: JKNameIdentifier("")
 
@@ -165,7 +187,7 @@ class JavaToJKTreeBuilder constructor(
                 is PsiPolyadicExpression -> {
                     val token = JKOperatorToken.fromElementType(operationTokenType)
                     val type = type?.toJK() ?: typeFactory.types.nullableAny
-                    val jkOperands = operands.map { it.toJK().parenthesizeIfBinaryExpression() }
+                    val jkOperands = operands.map { it.toJK().withLineBreaksFrom(it).parenthesizeIfBinaryExpression() }
                     jkOperands.reduce { acc, operand ->
                         JKBinaryExpression(acc, operand, JKKtOperatorImpl(token, type))
                     }.let { folded ->
@@ -180,7 +202,7 @@ class JavaToJKTreeBuilder constructor(
             }.also {
                 if (this != null) {
                     (it as PsiOwner).psi = this
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
             }
         }
@@ -195,14 +217,14 @@ class JavaToJKTreeBuilder constructor(
                     else -> JKClassLiteralExpression.ClassLiteralType.JAVA_CLASS
                 }
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
         fun PsiInstanceOfExpression.toJK(): JKIsExpression =
             JKIsExpression(operand.toJK(), JKTypeElement(checkType?.type?.toJK() ?: JKNoType))
                 .also {
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
 
         fun PsiAssignmentExpression.toJK(): JKJavaAssignmentExpression {
@@ -211,7 +233,7 @@ class JavaToJKTreeBuilder constructor(
                 rExpression.toJK(),
                 createOperator(operationSign.tokenType, type)
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -226,14 +248,14 @@ class JavaToJKTreeBuilder constructor(
                 else -> JKOperatorToken.fromElementType(operationSign.tokenType)
             }
             return JKBinaryExpression(
-                lOperand.toJK(),
+                lOperand.toJK().withLineBreaksFrom(lOperand),
-                rOperand.toJK(),
+                rOperand.toJK().withLineBreaksFrom(rOperand),
                 JKKtOperatorImpl(
                     token,
                     type?.toJK() ?: typeFactory.types.nullableAny
                 )
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -252,7 +274,7 @@ class JavaToJKTreeBuilder constructor(
                 JavaTokenType.DOUBLE_LITERAL -> JKLiteralExpression(text, DOUBLE)
                 else -> throwCanNotConvertError("Unknown literal element type: $literalElementType")
             }.also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -266,12 +288,12 @@ class JavaToJKTreeBuilder constructor(
             JavaTokenType.TILDE -> operand.toJK().callOn(symbolProvider.provideMethodSymbol("kotlin.Int.inv"))
             else -> JKPrefixExpression(operand.toJK(), createOperator(operationSign.tokenType, type))
         }.also {
-            it.assignNonCodeElements(this)
+            it.withFormattingFrom(this)
         }
 
         fun PsiPostfixExpression.toJK(): JKExpression =
             JKPostfixExpression(operand.toJK(), createOperator(operationSign.tokenType, type)).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
         fun PsiLambdaExpression.toJK(): JKExpression {
@@ -286,7 +308,7 @@ class JavaToJKTreeBuilder constructor(
                 with(declarationMapper) { parameterList.parameters.map { it.toJK() } },
                 functionalType()
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -309,7 +331,7 @@ class JavaToJKTreeBuilder constructor(
         fun PsiMethodCallExpression.toJK(): JKExpression {
             val arguments = argumentList
             val typeArguments = getExplicitTypeArguments().toJK()
-            val qualifier = methodExpression.qualifierExpression?.toJK()
+            val qualifier = methodExpression.qualifierExpression?.toJK()?.withLineBreaksFrom(methodExpression.qualifierExpression)
             val target = methodExpression.resolve()
             val symbol = target?.let {
                 symbolProvider.provideDirectSymbol(it)
@@ -397,7 +419,7 @@ class JavaToJKTreeBuilder constructor(
                     JKFieldAccessExpression(symbol).qualified(qualifier)
                 else -> throwCanNotConvertError("unexpected symbol ${symbol::class}")
             }.also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -447,16 +469,16 @@ class JavaToJKTreeBuilder constructor(
                 is JKPackageSymbol -> JKPackageAccessExpression(symbol)
                 else -> throwCanNotConvertError("unexpected symbol ${symbol::class}")
             }.qualified(qualifierExpression?.toJK()).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
         fun PsiArrayInitializerExpression.toJK(): JKExpression {
             return JKJavaNewArray(
-                initializers.map { it.toJK() },
+                initializers.map { it.toJK().withLineBreaksFrom(it) },
                 JKTypeElement(type?.toJK().safeAs<JKJavaArrayType>()?.type ?: JKContextType)
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -479,7 +501,7 @@ class JavaToJKTreeBuilder constructor(
                             child = child.nextSibling
                         }
                         JKJavaNewEmptyArray(
-                            dimensions.map { it?.toJK() ?: JKStubExpression() },
+                            dimensions.map { it?.toJK()?.withLineBreaksFrom(it) ?: JKStubExpression() },
                             JKTypeElement(generateSequence(type?.toJK()) { it.safeAs<JKJavaArrayType>()?.type }.last())
                         ).also {
                             it.psi = this
@@ -515,7 +537,7 @@ class JavaToJKTreeBuilder constructor(
         fun PsiReferenceParameterList.toJK(): JKTypeArgumentList =
             JKTypeArgumentList(typeArguments.map { JKTypeElement(it.toJK()) })
                 .also {
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
 
 
@@ -525,7 +547,7 @@ class JavaToJKTreeBuilder constructor(
                     symbolProvider.provideMethodSymbol("kotlin.Array.get"),
                     arguments = listOf(indexExpression?.toJK() ?: JKStubExpression())
                 ).also {
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
 
 
@@ -534,19 +556,19 @@ class JavaToJKTreeBuilder constructor(
                 operand?.toJK() ?: throwCanNotConvertError(),
                 (castType?.type?.toJK() ?: JKNoType).asTypeElement()
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
         fun PsiParenthesizedExpression.toJK(): JKExpression {
             return JKParenthesizedExpression(expression.toJK())
                 .also {
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
         }
 
         fun PsiExpressionList.toJK(): JKArgumentList {
-            val jkExpressions = expressions.map { it.toJK() }
+            val jkExpressions = expressions.map { it.toJK().withLineBreaksFrom(it) }
             return ((parent as? PsiCall)?.resolveMethod()
                 ?.let { method ->
                     val lastExpressionType = expressions.lastOrNull()?.type
@@ -558,13 +580,17 @@ class JavaToJKTreeBuilder constructor(
                             JKPrefixExpression(
                                 jkExpressions.last(),
                                 JKKtSpreadOperator(lastExpressionType.toJK())
-                            ).withNonCodeElementsFrom(jkExpressions.last())
+                            ).withFormattingFrom(jkExpressions.last())
+                        staredExpression.expression.also {
+                            it.hasLeadingLineBreak = false
+                            it.hasTrailingLineBreak = false
+                        }
                         jkExpressions.dropLast(1) + staredExpression
                     } else jkExpressions
                 } ?: jkExpressions)
                 .toArgumentList()
                 .also {
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
         }
 
@@ -575,7 +601,7 @@ class JavaToJKTreeBuilder constructor(
         fun PsiTypeParameterList.toJK(): JKTypeParameterList =
             JKTypeParameterList(typeParameters.map { it.toJK() })
                 .also {
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
 
         fun PsiTypeParameter.toJK(): JKTypeParameter =
@@ -584,7 +610,7 @@ class JavaToJKTreeBuilder constructor(
                 extendsListTypes.map { JKTypeElement(it.toJK()) }
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
         fun PsiClass.toJK(): JKClass =
@@ -601,7 +627,7 @@ class JavaToJKTreeBuilder constructor(
             ).also { klass ->
                 klass.psi = this
                 symbolProvider.provideUniverseSymbol(this, klass)
-                klass.assignNonCodeElements(this)
+                klass.withFormattingFrom(this)
             }
 
 
@@ -611,7 +637,7 @@ class JavaToJKTreeBuilder constructor(
             return JKInheritanceInfo(extensionType, implTypes)
                 .also {
                     if (implementsList != null) {
-                        it.assignNonCodeElements(implementsList!!)
+                        it.withFormattingFrom(implementsList!!)
                     }
                 }
         }
@@ -630,15 +656,21 @@ class JavaToJKTreeBuilder constructor(
                     }
                 }
             ).also {
-                it.leftBrace.assignNonCodeElements(lBrace)
+                it.leftBrace.withFormattingFrom(
-                it.rightBrace.assignNonCodeElements(rBrace)
+                    lBrace,
+                    takeLeadingComments = false
+                ) // do not capture comments which belongs to following declarations
+                it.rightBrace.withFormattingFrom(rBrace)
+                it.declarations.lastOrNull()?.let { lastMember ->
+                    lastMember.withLeadingCommentsWithParent(lastMember.psi)
+                }
             }
 
         fun PsiClassInitializer.toJK(): JKDeclaration = when {
             hasModifier(JvmModifier.STATIC) -> JKJavaStaticInitDeclaration(body.toJK())
             else -> JKKtInitDeclaration(body.toJK())
         }.also {
-            it.assignNonCodeElements(this)
+            it.withFormattingFrom(this)
         }
 
 
@@ -658,12 +690,12 @@ class JavaToJKTreeBuilder constructor(
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
                 it.psi = this
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
 
         fun PsiMember.modality() =
-            modality { ast, psi -> ast.assignNonCodeElements(psi) }
+            modality { ast, psi -> ast.withFormattingFrom(psi) }
 
         fun PsiMember.otherModifiers() =
             modifierList?.children?.mapNotNull { child ->
@@ -678,17 +710,17 @@ class JavaToJKTreeBuilder constructor(
 
                     else -> null
                 }?.let {
-                    JKOtherModifierElement(it).withAssignedNonCodeElements(child)
+                    JKOtherModifierElement(it).withFormattingFrom(child)
                 }
             }.orEmpty()
 
 
         private fun PsiMember.visibility(): JKVisibilityModifierElement =
-            visibility(referenceSearcher) { ast, psi -> ast.assignNonCodeElements(psi) }
+            visibility(referenceSearcher) { ast, psi -> ast.withFormattingFrom(psi) }
 
         fun PsiField.toJK(): JKField {
             return JKField(
-                JKTypeElement(type.toJK()).withAssignedNonCodeElements(typeElement),
+                JKTypeElement(type.toJK()).withFormattingFrom(typeElement),
                 nameIdentifier.toJK(),
                 with(expressionTreeMapper) { initializer.toJK() },
                 annotationList(this),
@@ -699,7 +731,7 @@ class JavaToJKTreeBuilder constructor(
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
                 it.psi = this
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
@@ -735,7 +767,7 @@ class JavaToJKTreeBuilder constructor(
                     }
                 }
             ).also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
         fun PsiDocComment.deprecatedAnnotation(): JKAnnotation? =
@@ -756,7 +788,7 @@ class JavaToJKTreeBuilder constructor(
                     JKKtAnnotationArrayInitializerExpression(initializers.map { it.toJK() })
                 else -> throwCanNotConvertError()
             }.also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
         fun PsiAnnotationMethod.toJK(): JKJavaAnnotationMethod =
@@ -773,7 +805,7 @@ class JavaToJKTreeBuilder constructor(
             ).also {
                 it.psi = this
                 symbolProvider.provideUniverseSymbol(this, it)
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
 
@@ -784,7 +816,7 @@ class JavaToJKTreeBuilder constructor(
                         ?: JKJavaVoidType.takeIf { isConstructor }
                         ?: JKNoType),
                 nameIdentifier.toJK(),
-                parameterList.parameters.map { it.toJK() },
+                parameterList.parameters.map { it.toJK().withLineBreaksFrom(it) },
                 body?.toJK() ?: JKBodyStub,
                 typeParameterList?.toJK() ?: JKTypeParameterList(),
                 annotationList(this),
@@ -798,10 +830,10 @@ class JavaToJKTreeBuilder constructor(
                 parameterList.node
                     ?.safeAs<CompositeElement>()
                     ?.also {
-                        jkMethod.leftParen.assignNonCodeElements(it.findChildByRoleAsPsiElement(ChildRole.LPARENTH))
+                        jkMethod.leftParen.withFormattingFrom(it.findChildByRoleAsPsiElement(ChildRole.LPARENTH))
-                        jkMethod.rightParen.assignNonCodeElements(it.findChildByRoleAsPsiElement(ChildRole.RPARENTH))
+                        jkMethod.rightParen.withFormattingFrom(it.findChildByRoleAsPsiElement(ChildRole.RPARENTH))
                     }
-            }.withAssignedNonCodeElements(this)
+            }.withFormattingFrom(this)
         }
 
         fun PsiParameter.toJK(): JKParameter {
@@ -817,22 +849,22 @@ class JavaToJKTreeBuilder constructor(
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
                 it.psi = this
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
         }
 
         fun PsiCodeBlock.toJK(): JKBlock =
             JKBlockImpl(statements.map { it.toJK() })
-                .withAssignedNonCodeElements(this)
+                .withFormattingFrom(this)
                 .also {
-                    it.leftBrace.assignNonCodeElements(lBrace)
+                    it.leftBrace.withFormattingFrom(lBrace)
-                    it.rightBrace.assignNonCodeElements(rBrace)
+                    it.rightBrace.withFormattingFrom(rBrace)
                 }
 
 
         fun PsiLocalVariable.toJK(): JKLocalVariable =
             JKLocalVariable(
-                JKTypeElement(type.toJK()).withAssignedNonCodeElements(typeElement),
+                JKTypeElement(type.toJK()).withFormattingFrom(typeElement),
                 nameIdentifier.toJK(),
                 with(expressionTreeMapper) { initializer.toJK() },
                 JKMutabilityModifierElement(
@@ -844,7 +876,7 @@ class JavaToJKTreeBuilder constructor(
                 symbolProvider.provideUniverseSymbol(this, i)
                 i.psi = this
             }.also {
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
 
         fun PsiStatement?.toJK(): JKStatement {
@@ -903,7 +935,7 @@ class JavaToJKTreeBuilder constructor(
                                         with(expressionTreeMapper) { statement.caseValue.toJK() },
                                         emptyList()
                                     )
-                                }.withAssignedNonCodeElements(statement)
+                                }.withFormattingFrom(statement)
                             else ->
                                 cases.lastOrNull()?.also { it.statements = it.statements + statement.toJK() }
                                     ?: run {
@@ -947,7 +979,7 @@ class JavaToJKTreeBuilder constructor(
             }.also {
                 if (this != null) {
                     (it as PsiOwner).psi = this
-                    it.assignNonCodeElements(this)
+                    it.withFormattingFrom(this)
                 }
             }
         }
@@ -958,7 +990,7 @@ class JavaToJKTreeBuilder constructor(
                 catchBlock?.toJK() ?: JKBodyStub
             ).also {
                 it.psi = this
-                it.assignNonCodeElements(this)
+                it.withFormattingFrom(this)
             }
     }
 
@@ -989,77 +1021,6 @@ class JavaToJKTreeBuilder constructor(
             else -> null
         }?.let { JKTreeRoot(it) }
 
-    private val tokenCache = mutableMapOf<PsiElement, JKNonCodeElement>()
-
-    private fun PsiElement.collectNonCodeElements(): Pair<List<JKNonCodeElement>, List<JKNonCodeElement>> {
-        fun PsiElement.toToken(): JKNonCodeElement? {
-            if (this in tokenCache) return tokenCache.getValue(this)
-            val token = when {
-                this is PsiDocComment ->
-                    JKCommentElement(IdeaDocCommentConverter.convertDocComment(this))
-                this is PsiComment -> JKCommentElement(text)
-                this is PsiWhiteSpace -> JKSpaceElement(text)
-                text == ";" -> null
-                text == "" -> null
-                else -> null
-            } ?: return null
-            tokenCache[this] = token
-            return token
-        }
-
-        fun Sequence<PsiElement>.toNonCodeElements(): List<JKNonCodeElement> =
-            takeWhile { it is PsiComment || it is PsiWhiteSpace || it.text == ";" }
-                .mapNotNull { it.toToken() }
-                .toList()
-
-        fun PsiElement.isNonCodeElement() =
-            this is PsiComment || this is PsiWhiteSpace || text == ";" || text == ""
-
-        fun PsiElement.nextNonCodeElements() =
-            generateSequence(nextSibling) { it.nextSibling }
-                .takeWhile { it.isNonCodeElement() }
-
-        fun PsiElement.prevNonCodeElements() =
-            generateSequence(prevSibling) { it.prevSibling }
-                .takeWhile { it.isNonCodeElement() }
-
-
-        fun PsiElement.nextNonCodeElementsWithParent(): Sequence<JKNonCodeElement> {
-            val innerElements = nextNonCodeElements()
-            return (if (innerElements.lastOrNull()?.nextSibling == null && this is PsiKeyword)
-                innerElements + parent?.nextNonCodeElements().orEmpty()
-            else innerElements).mapNotNull { it.toToken() }
-        }
-
-        fun PsiElement.prevNonCodeElementsWithParent(): Sequence<JKNonCodeElement> {
-            val innerElements = prevNonCodeElements()
-            return (if (innerElements.firstOrNull()?.prevSibling == null && this is PsiKeyword)
-                innerElements + parent?.prevNonCodeElements().orEmpty()
-            else innerElements).mapNotNull { it.toToken() }
-        }
-
-
-        val leftInnerTokens = children.asSequence().toNonCodeElements().reversed()
-        val rightInnerTokens =
-            if (children.isEmpty()) emptyList()
-            else generateSequence(children.last()) { it.prevSibling }
-                .toNonCodeElements()
-                .reversed()
-
-
-        return (leftInnerTokens + prevNonCodeElementsWithParent()).reversed() to
-                (rightInnerTokens + nextNonCodeElementsWithParent())
-    }
-
-    private fun JKNonCodeElementsListOwner.assignNonCodeElements(psi: PsiElement?) {
-        if (psi == null) return
-        val (leftTokens, rightTokens) = psi.collectNonCodeElements()
-        this.leftNonCodeElements += leftTokens
-        this.rightNonCodeElements += rightTokens
-    }
-
-    private inline fun <reified T : JKNonCodeElementsListOwner> T.withAssignedNonCodeElements(psi: PsiElement?): T =
-        also { it.assignNonCodeElements(psi) }
 
     private fun PsiElement.throwCanNotConvertError(message: String? = null): Nothing {
         throw KotlinExceptionWithAttachments("Cannot convert the following Java element ${this::class}" + message?.let { " due to `$it`" })

nj2k/src/org/jetbrains/kotlin/nj2k/NewCodeBuilder.kt

@@ -29,49 +29,50 @@ class NewCodeBuilder(context: NewJ2kConverterContext) {
     }
 
     inner class Visitor : JKVisitorWithCommentsPrinting() {
-        private val printedTokens = mutableSetOf<JKNonCodeElement>()
+        private val printedTokens = mutableSetOf<JKComment>()
 
         //TODO move to ast transformation phase
-        private fun JKNonCodeElement.shouldBeDropped(): Boolean =
+        private fun JKComment.shouldBeDropped(): Boolean =
-            this is JKCommentElement && text.startsWith("//noinspection")
+            text.startsWith("//noinspection")
 
-        private fun JKNonCodeElement.createText() =
+        private fun JKComment.createText() =
             if (this !in printedTokens) {
                 printedTokens += this
                 text
-            } else ""
+            } else null
 
 
-        private fun List<JKNonCodeElement>.createText(): String {
+        private fun List<JKComment>.createText(): String = buildString {
-            val text = filterNot { it.shouldBeDropped() }.joinToString("") { token -> token.createText() }
+            var needNewLine = false
-            val needNewLine = text.lastIndexOf('\n') < text.lastIndexOf("//")
+            for (comment in this@createText) {
-            return text + "\n".takeIf { needNewLine }.orEmpty()
+                if (comment.shouldBeDropped()) continue
+                val text = comment.createText() ?: continue
+                if (needNewLine) appendln() else append(' ')
+                append(text)
+                needNewLine = text.startsWith("//") || '\n' in text
+            }
         }
 
-        private fun JKNonCodeElementsListOwner.needPreserveSpacesAfterLastComment() =
+        private fun String.hasNoLineBreakAfterSingleLineComment() = lastIndexOf('\n') < lastIndexOf("//")
-            this is JKArgument
-                    || this is JKParameter
-                    || safeAs<JKTreeElement>()?.parent is JKArgument
-                    || safeAs<JKTreeElement>()?.parent is JKBinaryExpression
-                    || safeAs<JKTreeElement>()?.parent is JKQualifiedExpression
 
-        override fun printLeftNonCodeElements(element: JKNonCodeElementsListOwner) {
+        override fun printLeftNonCodeElements(element: JKFormattingOwner) {
-            val text = element.leftNonCodeElements
+            val text = element.trailingComments.createText()
-                .let {
-                    if (element.needPreserveSpacesAfterLastComment()) it
-                    else it.dropWhile { it is JKSpaceElement }
-                }.createText()
             printer.printWithNoIndent(text)
+
+            val addNewLine = element.hasTrailingLineBreak
+                    || element is JKDeclaration && element.trailingComments.isNotEmpty() // add new line between comment & declaration
+                    || text.hasNoLineBreakAfterSingleLineComment()
+
+            if (addNewLine) printer.println()
         }
 
 
-        override fun printRightNonCodeElements(element: JKNonCodeElementsListOwner) {
+        override fun printRightNonCodeElements(element: JKFormattingOwner) {
-            val text = element.rightNonCodeElements
+            val text = element.leadingComments.createText()
-                .let {
-                    if (element.needPreserveSpacesAfterLastComment()) it
-                    else it.dropLastWhile { it is JKSpaceElement }
-                }.createText()
             printer.printWithNoIndent(text)
+
+            val addNewLine = element.hasLeadingLineBreak || text.hasNoLineBreakAfterSingleLineComment()
+            if (addNewLine) printer.println()
         }
 
         private fun renderTokenElement(tokenElement: JKTokenElement) {

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/AddParenthesisForLineBreaksInBinaryExpression.kt

@@ -13,15 +13,13 @@ class AddParenthesisForLineBreaksInBinaryExpression(override val context: NewJ2k
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         if (element !is JKBinaryExpression) return recurse(element)
         if (element.parent is JKBinaryExpression) return recurse(element)
-        if (element.left.rightNonCodeElements.any {
+        if (element.left.hasLeadingLineBreak) {
-                it is JKSpaceElement && '\n' in it.text
-            }) {
             return JKParenthesizedExpression(
                 JKBinaryExpression(
                     element::left.detached(),
                     element::right.detached(),
                     element.operator
-                ).withNonCodeElementsFrom(element)
+                ).withFormattingFrom(element)
             )
         }
         return recurse(element)

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/AnnotationClassConversion.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.j2k.ast.Nullability
 import org.jetbrains.kotlin.nj2k.NewJ2kConverterContext
-import org.jetbrains.kotlin.nj2k.modality
 import org.jetbrains.kotlin.nj2k.toExpression
 import org.jetbrains.kotlin.nj2k.tree.*
 
@@ -58,11 +57,11 @@ class AnnotationClassConversion(context: NewJ2kConverterContext) : RecursiveAppl
                 JKKtAnnotationArrayInitializerExpression(initializer)
             } else initializer
         ).also { parameter ->
-            if (leftNonCodeElements.any { it is JKCommentElement }) {
+            if (trailingComments.any { it is JKComment }) {
-                parameter.leftNonCodeElements += leftNonCodeElements
+                parameter.trailingComments += trailingComments
             }
-            if (rightNonCodeElements.any { it is JKCommentElement }) {
+            if (leadingComments.any { it is JKComment }) {
-                parameter.rightNonCodeElements += rightNonCodeElements
+                parameter.leadingComments += leadingComments
             }
 
         }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/AssignmentExpressionUnfoldingConversion.kt

@@ -34,14 +34,14 @@ class AssignmentExpressionUnfoldingConversion(context: NewJ2kConverterContext) :
                     val assignment = statement.expression as JKJavaAssignmentExpression
                     newStatements += assignment
                         .unfoldToStatementsList(assignmentTarget = null)
-                        .withNonCodeElementsFrom(statement)
+                        .withFormattingFrom(statement)
                 }
                 statement is JKDeclarationStatement && statement.containsAssignment() -> {
                     val variable = statement.declaredStatements.single() as JKVariable
                     val assignment = variable.initializer as JKJavaAssignmentExpression
                     newStatements += assignment
                         .unfoldToStatementsList(variable.detached(statement))
-                        .withNonCodeElementsFrom(statement)
+                        .withFormattingFrom(statement)
                 }
                 else -> {
                     newStatements += statement
@@ -62,7 +62,7 @@ class AssignmentExpressionUnfoldingConversion(context: NewJ2kConverterContext) :
                 assignment::expression.detached(),
                 assignment.operator
             )
-        }.withNonCodeElementsFrom(this)
+        }.withFormattingFrom(this)
     }
 
     private fun JKExpressionStatement.canBeConvertedToBlock() = when (val parent = parent) {
@@ -73,7 +73,7 @@ class AssignmentExpressionUnfoldingConversion(context: NewJ2kConverterContext) :
     }
 
     private fun JKJavaAssignmentExpression.convertAssignments() =
-        unfoldToExpressionsChain().withNonCodeElementsFrom(this)
+        unfoldToExpressionsChain().withFormattingFrom(this)
 
     private fun JKDeclarationStatement.containsAssignment() =
         declaredStatements.singleOrNull()?.safeAs<JKVariable>()?.initializer is JKJavaAssignmentExpression
@@ -118,7 +118,7 @@ class AssignmentExpressionUnfoldingConversion(context: NewJ2kConverterContext) :
             this::field.detached(),
             JKKtItExpression(operator.returnType),
             operator
-        ).withNonCodeElementsFrom(this)
+        ).withFormattingFrom(this)
         return when {
             operator.isSimpleToken() ->
                 JKAssignmentChainAlsoLink(receiver, assignment, field.copyTreeAndDetach())
@@ -129,7 +129,7 @@ class AssignmentExpressionUnfoldingConversion(context: NewJ2kConverterContext) :
 
     private fun JKJavaAssignmentExpression.toDeclarationChainLink(expression: JKExpression) =
         createKtAssignmentStatement(this::field.detached(), expression, this.operator)
-            .withNonCodeElementsFrom(this)
+            .withFormattingFrom(this)
 
     private fun createKtAssignmentStatement(
         field: JKExpression,

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/BlockToRunConversion.kt

@@ -27,7 +27,7 @@ class BlockToRunConversion(context: NewJ2kConverterContext) : RecursiveApplicabl
             symbolProvider.provideMethodSymbol("kotlin.run"),
             JKArgumentList(lambda)
         )
-        return recurse(JKExpressionStatement(call).withNonCodeElementsFrom(element))
+        return recurse(JKExpressionStatement(call).withFormattingFrom(element))
     }
 
 }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/BoxedTypeOperationsConversion.kt

@@ -43,7 +43,7 @@ class BoxedTypeOperationsConversion(context: NewJ2kConverterContext) : Recursive
                 "kotlin.${primitiveTypeName.capitalize()}.to${operationType.capitalize()}"
             ),
             JKArgumentList()
-        ).withNonCodeElementsFrom(methodCallExpression)
+        ).withFormattingFrom(methodCallExpression)
     }
 
     companion object {

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/BuiltinMembersConversion.kt

@@ -33,11 +33,11 @@ class BuiltinMembersConversion(context: NewJ2kConverterContext) : RecursiveAppli
         val newSelector = conversion.createBuilder().build(selector)
 
         if (this is JKQualifiedExpression && conversion.replaceType == ReplaceType.REPLACE_WITH_QUALIFIER) {
-            newSelector.rightNonCodeElements += receiver.leftNonCodeElements
+            newSelector.leadingComments += receiver.trailingComments
-            newSelector.rightNonCodeElements += receiver.leftNonCodeElements
+            newSelector.leadingComments += receiver.trailingComments
-            newSelector.rightNonCodeElements += receiver.rightNonCodeElements
+            newSelector.leadingComments += receiver.leadingComments
-            newSelector.rightNonCodeElements += selector.leftNonCodeElements
+            newSelector.leadingComments += selector.trailingComments
-            newSelector.rightNonCodeElements += selector.rightNonCodeElements
+            newSelector.leadingComments += selector.leadingComments
         }
 
         return when (conversion.replaceType) {
@@ -93,19 +93,19 @@ class BuiltinMembersConversion(context: NewJ2kConverterContext) : RecursiveAppli
                         symbolProvider.provideMethodSymbol(fqName),
                         argumentsProvider(from::arguments.detached()),
                         from::typeArgumentList.detached()
-                    ).withNonCodeElementsFrom(from)
+                    ).withFormattingFrom(from)
                 is JKFieldAccessExpression ->
                     JKCallExpressionImpl(
                         symbolProvider.provideMethodSymbol(fqName),
                         JKArgumentList(),
                         JKTypeArgumentList()
-                    ).withNonCodeElementsFrom(from)
+                    ).withFormattingFrom(from)
                 is JKNewExpression ->
                     JKCallExpressionImpl(
                         symbolProvider.provideMethodSymbol(fqName),
                         argumentsProvider(from::arguments.detached()),
                         JKTypeArgumentList()
-                    ).withNonCodeElementsFrom(from)
+                    ).withFormattingFrom(from)
                 else -> error("Bad conversion")
             }
     }
@@ -118,11 +118,11 @@ class BuiltinMembersConversion(context: NewJ2kConverterContext) : RecursiveAppli
                 is JKCallExpression ->
                     JKFieldAccessExpression(
                         symbolProvider.provideFieldSymbol(fqName)
-                    ).withNonCodeElementsFrom(from)
+                    ).withFormattingFrom(from)
                 is JKFieldAccessExpression ->
                     JKFieldAccessExpression(
                         symbolProvider.provideFieldSymbol(fqName)
-                    ).withNonCodeElementsFrom(from)
+                    ).withFormattingFrom(from)
                 else -> error("Bad conversion")
             }
     }
@@ -141,7 +141,7 @@ class BuiltinMembersConversion(context: NewJ2kConverterContext) : RecursiveAppli
                             JKArgumentList(arguments.drop(1)),
                             from::typeArgumentList.detached()
                         )
-                    ).withNonCodeElementsFrom(from)
+                    ).withFormattingFrom(from)
                 }
                 else -> error("Bad conversion")
             }
@@ -354,7 +354,7 @@ class BuiltinMembersConversion(context: NewJ2kConverterContext) : RecursiveAppli
                 JKArgumentList(
                     detachedArguments.drop(1) + JKNamedArgument(
                         detachedArguments.first()::value.detached().also {
-                            it.clearNonCodeElements()
+                            it.clearFormatting()
                         },
                         JKNameIdentifier("ignoreCase")
                     )

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/ClassToObjectPromotionConversion.kt

@@ -53,7 +53,7 @@ class ClassToObjectPromotionConversion(context: NewJ2kConverterContext) : Recurs
                         element.otherModifierElements,
                         element.visibilityElement,
                         JKModalityModifierElement(Modality.FINAL)
-                    ).withNonCodeElementsFrom(element)
+                    ).withFormattingFrom(element)
                 )
             }
         }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/ConstructorConversion.kt

@@ -29,7 +29,7 @@ class ConstructorConversion(context: NewJ2kConverterContext) : RecursiveApplicab
             element.modalityElement
         ).also {
             symbolProvider.transferSymbol(it, element)
-        }.withNonCodeElementsFrom(element)
+        }.withFormattingFrom(element)
     }
 
     private fun lookupDelegationCall(block: JKBlock): JKDelegationConstructorCall? {

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/DefaultArgumentsConversion.kt

@@ -108,7 +108,7 @@ class DefaultArgumentsConversion(context: NewJ2kConverterContext) : RecursiveApp
                 parameter.initializer = remapParameterSymbol(defaultValue) as JKExpression
             }
             element.declarations -= method
-            calledMethod.withNonCodeElementsFrom(method)
+            calledMethod.withFormattingFrom(method)
         }
         if (element.parentOfType<JKClass>()?.classKind != JKClass.ClassKind.ANNOTATION) {
             for (method in element.declarations) {

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/ForConversion.kt

@@ -27,8 +27,8 @@ class ForConversion(context: NewJ2kConverterContext) : RecursiveApplicableConver
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         if (element !is JKJavaForLoopStatement) return recurse(element)
 
-        convertToForeach(element)?.also { return recurse(it.withNonCodeElementsFrom(element)) }
+        convertToForeach(element)?.also { return recurse(it.withFormattingFrom(element)) }
-        convertToWhile(element)?.also { return recurse(it.withNonCodeElementsFrom(element)) }
+        convertToWhile(element)?.also { return recurse(it.withFormattingFrom(element)) }
 
         return recurse(element)
     }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/JavaModifiersConversion.kt

@@ -23,25 +23,25 @@ class JavaModifiersConversion(context: NewJ2kConverterContext) : RecursiveApplic
             element.elementByModifier(OtherModifier.VOLATILE)?.also { modifierElement ->
                 element.otherModifierElements -= modifierElement
                 element.annotationList.annotations +=
-                    jvmAnnotation("Volatile", symbolProvider).withNonCodeElementsFrom(modifierElement)
+                    jvmAnnotation("Volatile", symbolProvider).withFormattingFrom(modifierElement)
             }
 
             element.elementByModifier(OtherModifier.TRANSIENT)?.also { modifierElement ->
                 element.otherModifierElements -= modifierElement
                 element.annotationList.annotations +=
-                    jvmAnnotation("Transient", symbolProvider).withNonCodeElementsFrom(modifierElement)
+                    jvmAnnotation("Transient", symbolProvider).withFormattingFrom(modifierElement)
             }
 
             element.elementByModifier(OtherModifier.STRICTFP)?.also { modifierElement ->
                 element.otherModifierElements -= modifierElement
                 element.annotationList.annotations +=
-                    jvmAnnotation("Strictfp", symbolProvider).withNonCodeElementsFrom(modifierElement)
+                    jvmAnnotation("Strictfp", symbolProvider).withFormattingFrom(modifierElement)
             }
 
             element.elementByModifier(OtherModifier.SYNCHRONIZED)?.also { modifierElement ->
                 element.otherModifierElements -= modifierElement
                 element.annotationList.annotations +=
-                    jvmAnnotation("Synchronized", symbolProvider).withNonCodeElementsFrom(modifierElement)
+                    jvmAnnotation("Synchronized", symbolProvider).withFormattingFrom(modifierElement)
             }
 
             element.elementByModifier(OtherModifier.NATIVE)?.also { modifierElement ->

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/NonCodeElementsConversion.kt

@@ -14,8 +14,8 @@ class NonCodeElementsConversion(context: NewJ2kConverterContext) : RecursiveAppl
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         when (element) {
             is JKClass -> {
-                element.name.rightNonCodeElements += element.inheritance.leftNonCodeElements
+                element.name.leadingComments += element.inheritance.trailingComments
-                element.inheritance.leftNonCodeElements.clear()
+                element.inheritance.trailingComments.clear()
             }
         }
         return recurse(element)

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/PrimaryConstructorDetectConversion.kt

@@ -32,10 +32,10 @@ class PrimaryConstructorDetectConversion(context: NewJ2kConverterContext) : Recu
         primaryConstructorCandidate.invalidate()
         if (primaryConstructorCandidate.block.statements.isNotEmpty()) {
             val initDeclaration = JKKtInitDeclaration(primaryConstructorCandidate.block)
-                .withNonCodeElementsFrom(primaryConstructorCandidate)
+                .withFormattingFrom(primaryConstructorCandidate)
-            primaryConstructorCandidate.clearNonCodeElements()
+            primaryConstructorCandidate.clearFormatting()
             primaryConstructorCandidate.forEachModifier { modifierElement ->
-                modifierElement.clearNonCodeElements()
+                modifierElement.clearFormatting()
             }
             element.classBody.declarations =
                 element.classBody.declarations.replace(primaryConstructorCandidate, initDeclaration)
@@ -52,7 +52,7 @@ class PrimaryConstructorDetectConversion(context: NewJ2kConverterContext) : Recu
                 primaryConstructorCandidate.otherModifierElements,
                 primaryConstructorCandidate.visibilityElement,
                 primaryConstructorCandidate.modalityElement
-            ).withNonCodeElementsFrom(primaryConstructorCandidate)
+            ).withFormattingFrom(primaryConstructorCandidate)
 
         symbolProvider.transferSymbol(primaryConstructor, primaryConstructorCandidate)
 

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/RemoveRedundantQualifiersForCallsConversion.kt

@@ -24,7 +24,7 @@ class RemoveRedundantQualifiersForCallsConversion(context: NewJ2kConverterContex
         }
         if (needRemoveQualifier) {
             element.invalidate()
-            return recurse(element.selector.withNonCodeElementsFrom(element.receiver).withNonCodeElementsFrom(element))
+            return recurse(element.selector.withFormattingFrom(element.receiver).withFormattingFrom(element))
         }
         return recurse(element)
     }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/StaticInitDeclarationConversion.kt

@@ -34,5 +34,5 @@ class StaticInitDeclarationConversion(context : NewJ2kConverterContext) : Recurs
     }
 
     private fun JKJavaStaticInitDeclaration.toKtInitDeclaration() =
-        JKKtInitDeclaration(::block.detached()).withNonCodeElementsFrom(this)
+        JKKtInitDeclaration(::block.detached()).withFormattingFrom(this)
 }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/SwitchStatementConversion.kt

@@ -48,7 +48,7 @@ class SwitchStatementConversion(context: NewJ2kConverterContext) : RecursiveAppl
                                 statement.block.statements
                                     .takeWhile { !isSwitchBreak(it) }
                                     .map { it.copyTreeAndDetach() }
-                            ).withNonCodeElementsFrom(statement)
+                            ).withFormattingFrom(statement)
                         isSwitchBreak(statement) -> null
                         else -> statement.copyTreeAndDetach()
                     }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/SynchronizedStatementConversion.kt

@@ -24,7 +24,7 @@ class SynchronizedStatementConversion(context: NewJ2kConverterContext) : Recursi
                     element.lockExpression,
                     lambdaBody
                 )
-            ).withNonCodeElementsFrom(element)
+            ).withFormattingFrom(element)
         return recurse(JKExpressionStatement(synchronizedCall))
     }
 

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/TryStatementConversion.kt

@@ -29,7 +29,7 @@ class TryStatementConversion(context: NewJ2kConverterContext) : RecursiveApplica
                 tryStatement::finallyBlock.detached(),
                 tryStatement.catchSections.flatMap(::convertCatchSection)
             )
-        ).withNonCodeElementsFrom(tryStatement)
+        ).withFormattingFrom(tryStatement)
 
     private fun convertTryStatementWithResources(tryStatement: JKJavaTryStatement): JKStatement {
         val body =
@@ -44,7 +44,7 @@ class TryStatementConversion(context: NewJ2kConverterContext) : RecursiveApplica
                     tryStatement::finallyBlock.detached(),
                     tryStatement.catchSections.flatMap(::convertCatchSection)
                 )
-            ).withNonCodeElementsFrom(tryStatement)
+            ).withFormattingFrom(tryStatement)
         } else body
     }
 
@@ -77,7 +77,7 @@ class TryStatementConversion(context: NewJ2kConverterContext) : RecursiveApplica
             JKKtTryCatchSection(
                 parameter,
                 javaCatchSection.block.copyTreeAndDetach()
-            ).withNonCodeElementsFrom(javaCatchSection)
+            ).withFormattingFrom(javaCatchSection)
         }
     }
 }

nj2k/src/org/jetbrains/kotlin/nj2k/conversions/TypeMappingConversion.kt

@@ -37,7 +37,7 @@ class TypeMappingConversion(
                         element::typeArgumentList.detached().fixTypeArguments(newClassSymbol),
                         element::classBody.detached(),
                         element.isAnonymousClass
-                    ).withNonCodeElementsFrom(element)
+                    ).withFormattingFrom(element)
                 )
             }
         }

nj2k/src/org/jetbrains/kotlin/nj2k/psiUtils.kt

@@ -45,7 +45,7 @@ fun canKeepEqEq(left: PsiExpression, right: PsiExpression?): Boolean {
 
 internal fun PsiMember.visibility(
     referenceSearcher: ReferenceSearcher,
-    assignNonCodeElements: ((JKNonCodeElementsListOwner, PsiElement) -> Unit)?
+    assignNonCodeElements: ((JKFormattingOwner, PsiElement) -> Unit)?
 ): JKVisibilityModifierElement =
     modifierList?.children?.mapNotNull { child ->
         if (child !is PsiKeyword) return@mapNotNull null
@@ -64,7 +64,7 @@ internal fun PsiMember.visibility(
     }?.firstOrNull() ?: JKVisibilityModifierElement(Visibility.INTERNAL)
 
 
-fun PsiMember.modality(assignNonCodeElements: ((JKNonCodeElementsListOwner, PsiElement) -> Unit)?) =
+fun PsiMember.modality(assignNonCodeElements: ((JKFormattingOwner, PsiElement) -> Unit)?) =
     modifierList?.children?.mapNotNull { child ->
         if (child !is PsiKeyword) return@mapNotNull null
         when (child.text) {

nj2k/src/org/jetbrains/kotlin/nj2k/tree/base.kt

@@ -46,9 +46,11 @@ private class JKListChild<T : JKElement>(val value: Int) : ReadWriteProperty<JKT
 }
 
 
-abstract class JKTreeElement : JKElement, JKNonCodeElementsListOwner, Cloneable {
+abstract class JKTreeElement : JKElement, JKFormattingOwner, Cloneable {
-    override val leftNonCodeElements: MutableList<JKNonCodeElement> = mutableListOf()
+    override val trailingComments: MutableList<JKComment> = mutableListOf()
-    override val rightNonCodeElements: MutableList<JKNonCodeElement> = mutableListOf()
+    override val leadingComments: MutableList<JKComment> = mutableListOf()
+    override var hasTrailingLineBreak = false
+    override var hasLeadingLineBreak = false
 
     override var parent: JKElement? = null
 
@@ -149,11 +151,11 @@ interface PsiOwner {
 
 class PsiOwnerImpl(override var psi: PsiElement? = null) : PsiOwner
 
-interface JKTypeArgumentListOwner : JKNonCodeElementsListOwner {
+interface JKTypeArgumentListOwner : JKFormattingOwner {
     var typeArgumentList: JKTypeArgumentList
 }
 
-interface JKTypeParameterListOwner : JKNonCodeElementsListOwner {
+interface JKTypeParameterListOwner : JKFormattingOwner {
     var typeParameterList: JKTypeParameterList
 }
 

nj2k/src/org/jetbrains/kotlin/nj2k/tree/elements.kt

@@ -40,8 +40,10 @@ abstract class JKBlock : JKTreeElement() {
 
 
 object JKBodyStub : JKBlock() {
-    override val leftNonCodeElements: MutableList<JKNonCodeElement> = mutableListOf()
+    override val trailingComments: MutableList<JKComment> = mutableListOf()
-    override val rightNonCodeElements: MutableList<JKNonCodeElement> = mutableListOf()
+    override val leadingComments: MutableList<JKComment> = mutableListOf()
+    override var hasTrailingLineBreak = false
+    override var hasLeadingLineBreak = false
 
     override fun copy(): JKTreeElement = this
 
@@ -172,7 +174,7 @@ class JKNameIdentifier(val value: String) : JKTreeElement() {
 }
 
 
-interface JKAnnotationListOwner : JKNonCodeElementsListOwner {
+interface JKAnnotationListOwner : JKFormattingOwner {
     var annotationList: JKAnnotationList
 }
 

nj2k/src/org/jetbrains/kotlin/nj2k/tree/expressions.kt

@@ -26,8 +26,8 @@ class JKBinaryExpression(
     right: JKExpression,
     override var operator: JKOperator
 ) : JKOperatorExpression() {
-    var right by child(right)
     var left by child(left)
+    var right by child(right)
 
     override fun accept(visitor: JKVisitor) = visitor.visitBinaryExpression(this)
 }

nj2k/src/org/jetbrains/kotlin/nj2k/tree/formatting.kt

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.nj2k.tree
+
+import org.jetbrains.kotlin.utils.SmartList
+
+
+class JKComment(val text: String) {
+    val isSingleline
+        get() = text.startsWith("//")
+}
+
+class JKTokenElementImpl(override val text: String) : JKTokenElement {
+    override val trailingComments: MutableList<JKComment> = SmartList()
+    override val leadingComments: MutableList<JKComment> = SmartList()
+    override var hasTrailingLineBreak: Boolean = false
+    override var hasLeadingLineBreak: Boolean = false
+}
+
+interface JKFormattingOwner {
+    val trailingComments: MutableList<JKComment>
+    val leadingComments: MutableList<JKComment>
+    var hasTrailingLineBreak: Boolean
+    var hasLeadingLineBreak: Boolean
+}
+
+inline fun <reified T : JKFormattingOwner> T.withFormattingFrom(other: JKFormattingOwner): T = also {
+    trailingComments += other.trailingComments
+    leadingComments += other.leadingComments
+    hasTrailingLineBreak = other.hasTrailingLineBreak
+    hasLeadingLineBreak = other.hasLeadingLineBreak
+}
+
+inline fun <reified T : JKFormattingOwner> List<T>.withFormattingFrom(other: JKFormattingOwner): List<T> = also {
+    if (isNotEmpty()) {
+        it.first().trailingComments += other.trailingComments
+        it.first().hasTrailingLineBreak = other.hasTrailingLineBreak
+        it.last().leadingComments += other.leadingComments
+        it.last().hasLeadingLineBreak = other.hasLeadingLineBreak
+    }
+}
+
+fun JKFormattingOwner.clearFormatting() {
+    trailingComments.clear()
+    leadingComments.clear()
+    hasTrailingLineBreak = false
+    hasLeadingLineBreak = false
+}
+
+interface JKTokenElement : JKFormattingOwner {
+    val text: String
+}
+
+fun JKFormattingOwner.containsNewLine(): Boolean =
+    hasTrailingLineBreak || hasLeadingLineBreak

nj2k/src/org/jetbrains/kotlin/nj2k/tree/modifiers.kt

@@ -98,7 +98,7 @@ enum class Mutability(override val text: String) : Modifier {
     UNKNOWN("var")
 }
 
-interface JKModifiersListOwner : JKNonCodeElementsListOwner
+interface JKModifiersListOwner : JKFormattingOwner
 
 fun JKOtherModifiersOwner.elementByModifier(modifier: OtherModifier): JKOtherModifierElement? =
     otherModifierElements.firstOrNull { it.otherModifier == modifier }

nj2k/src/org/jetbrains/kotlin/nj2k/tree/tokenElements.kt

@@ -1,56 +0,0 @@
-/*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.nj2k.tree
-
-import org.jetbrains.kotlin.utils.SmartList
-
-
-interface JKNonCodeElement {
-    val text: String
-}
-
-class JKSpaceElement(override val text: String) : JKNonCodeElement
-class JKCommentElement(override val text: String) : JKNonCodeElement
-
-class JKTokenElementImpl(override val text: String) : JKTokenElement {
-    override val leftNonCodeElements: MutableList<JKNonCodeElement> = SmartList()
-    override val rightNonCodeElements: MutableList<JKNonCodeElement> = SmartList()
-}
-
-interface JKNonCodeElementsListOwner {
-    val leftNonCodeElements: MutableList<JKNonCodeElement>
-    val rightNonCodeElements: MutableList<JKNonCodeElement>
-}
-
-fun JKNonCodeElementsListOwner.takeNonCodeElementsFrom(other: JKNonCodeElementsListOwner) {
-    leftNonCodeElements += other.leftNonCodeElements
-    rightNonCodeElements += other.rightNonCodeElements
-}
-
-inline fun <reified T : JKNonCodeElementsListOwner> T.withNonCodeElementsFrom(other: JKNonCodeElementsListOwner): T =
-    also { it.takeNonCodeElementsFrom(other) }
-
-inline fun <reified T : JKNonCodeElementsListOwner> List<T>.withNonCodeElementsFrom(other: JKNonCodeElementsListOwner): List<T> =
-    also {
-        if (isNotEmpty()) {
-            it.first().leftNonCodeElements += other.leftNonCodeElements
-            it.last().rightNonCodeElements += other.rightNonCodeElements
-        }
-    }
-
-fun JKNonCodeElementsListOwner.clearNonCodeElements() {
-    leftNonCodeElements.clear()
-    rightNonCodeElements.clear()
-}
-
-interface JKTokenElement : JKNonCodeElementsListOwner {
-    val text: String
-}
-
-fun JKNonCodeElementsListOwner.containsNewLine(): Boolean =
-    (leftNonCodeElements + rightNonCodeElements).any {
-        it is JKSpaceElement && '\n' in it.text
-    }

nj2k/src/org/jetbrains/kotlin/nj2k/tree/treeUtils.kt

@@ -102,7 +102,7 @@ fun <R : JKTreeElement> applyRecursive(
 
 
 inline fun <reified T : JKTreeElement> T.copyTree(): T =
-    copy().withNonCodeElementsFrom(this) as T
+    copy().withFormattingFrom(this) as T
 
 inline fun <reified T : JKTreeElement> T.copyTreeAndDetach(): T =
     copyTree().also {

nj2k/src/org/jetbrains/kotlin/nj2k/tree/visitors/JKVisitorWithCommentsPrinting.kt

@@ -3,8 +3,8 @@ package org.jetbrains.kotlin.nj2k.tree.visitors
 import org.jetbrains.kotlin.nj2k.tree.*
 
 abstract class JKVisitorWithCommentsPrinting : JKVisitor() {
-    abstract fun printLeftNonCodeElements(element: JKNonCodeElementsListOwner)
+    abstract fun printLeftNonCodeElements(element: JKFormattingOwner)
-    abstract fun printRightNonCodeElements(element: JKNonCodeElementsListOwner)
+    abstract fun printRightNonCodeElements(element: JKFormattingOwner)
 
     override fun visitTreeElement(treeElement: JKTreeElement) {
         printLeftNonCodeElements(treeElement)

nj2k/testData/newJ2k/assignmentExpression/bitOperationPriorityComments.kt

@@ -1 +1 @@
-/*pre*/ 5 /*operand '5'*/ shl  /*left shift*/ 16 /*operand '16'*/ or (/*or*/ 1 /*operand '1'*/ shr  /*right shift*/ 8 /*operand '8'*/) or  /*or*/ 0 /*operand '0'*/ /*post comment*/
+/*pre*/5 /*operand '5'*/ shl  /*left shift*/16 /*operand '16'*/ or ( /*or*/1 /*operand '1'*/ shr  /*right shift*/8 /*operand '8'*/) or  /*or*/0 /*operand '0'*/ /*post comment*/

nj2k/testData/newJ2k/comments/comments2.kt

@@ -3,19 +3,17 @@
 package foo
 
 internal class A {
-    fun /* nothing to return */ foo(/* no parameters at all */) {
+    fun  /* nothing to return */foo( /* no parameters at all */) { // let declare a variable
-        // let declare a variable
+// with 2 comments before
-        // with 2 comments before
+        val   /*int*/a /* it's a */ = 2 /* it's 2 */ + 1 /* it's 1 */ // variable a declared
-
-        val /*int*/ a /* it's a */ = 2 /* it's 2 */ + 1 /* it's 1 */ // variable a declared
     } // end of foo
 
-    fun /* we return int*/ foo(  /*int*/ p: Int/* parameter p */): Int { /* body is empty */
+    fun  /* we return int*/foo(   /*int*/p: Int /* parameter p */): Int { /* body is empty */
     }
 
-    private/*it's private*/  val field = 0
+    private /*it's private*/  val field = 0
-    /*it's public*/  fun foo(s: String?): Char {}
+    /*it's public*/   fun foo(s: String?): Char {}
-    protected/*it's protected*/  fun foo(c: Char) {}
+    protected /*it's protected*/   fun foo(c: Char) {}
     /**
      * Method description.
      * Multi-line method description.
@@ -28,7 +26,7 @@ internal class A {
      */
     fun foo(param1: String?, param2: String?, param3: String?) {}
 
-    companion object {
+    companion object { /*it's public*/  /*and static*/ /*and final*/
-        /*it's public*/ /*and static*/ /*and final*/ const val C = 1
+        const val C = 1
     }
 }

nj2k/testData/newJ2k/comments/commentsForConstructors.kt

@@ -5,10 +5,7 @@ internal class A // end of primary constructor body
     // this is a secondary constructor 2
     constructor(s: String) : this(s.length) {} // end of secondary constructor 2 body
 
-
+    // this is a primary constructor
-// this is a primary constructor
-
-
     // this is a secondary constructor 1
 
 
@@ -18,7 +15,7 @@ internal class A // end of primary constructor body
 internal class B // end of constructor body
 (private val x: Int) {
     fun foo() {}
-// this constructor will disappear
+    // this constructor will disappear
 
 }
 

nj2k/testData/newJ2k/comments/fieldsInitializedFromParams.kt

@@ -3,4 +3,5 @@ internal class C(// field p1
         /**
          * Field myP2
          */
-        private val myP2: Int, /* Field p3 */var p3: Int)
+        private val myP2: Int, /* Field p3 */
+        var p3: Int)

nj2k/testData/newJ2k/detectProperties/Comments.kt

@@ -1,14 +1,12 @@
 internal class A(// comment for field2 setter
         // comment for field2 getter
         var field2 // comment for field2
-        : Int) {
+        : Int) { // Comment for field1
     /**
      * Comment for field1 setter
      */
     // Comment for field1 getter
-    // Comment for field1
     var field1 = 0
-
     // comment for field3 setter
     // comment for field3 getter
     // comment before field3
@@ -22,4 +20,6 @@ internal class A(// comment for field2 setter
     var property: Int
         get() = 1
         set(value) {}
+
+
 }

nj2k/testData/newJ2k/detectProperties/commentInInitStatement.kt

@@ -1,3 +1,2 @@
-class Nya /*comment before*/
+class Nya /*comment before*/ /* comment after*/ // and after again
-/* comment after*/ // and after again
+(/*1*/ /*3*/  /*4*/private val i: Int)
-(/*1*//*3*/ /*4*/private val i: Int)

nj2k/testData/newJ2k/formatting/staticAndNonStaticMembersWithComments.kt

@@ -1,29 +1,16 @@
-internal class F {
+internal class F { //c3
-
-    //c1
-
-    /*c2*///c3
-
-
     //c4
-
     fun f2() {}
+
     fun f3() {}
     fun f4() {}
 
-    //c5
+    companion object { //c1
-
+        /*c2*/
-    companion object {
-
         fun f1() {}
 
-
         var i = 0
-
+        //c5
-
+        fun f5() {} //c6
-        fun f5() {}
-
-        //c6
     }
-
 }

nj2k/testData/newJ2k/issues/comments.kt

@@ -1,17 +1,12 @@
 // This is an end-of-line comment
-
 /*
 This is a block comment
 */
-
-
 /*doc comment of class*/
 //one line comment of class
 //another one
 /*another doc*/
-internal class C {
+internal class C { // This is a class comment
-    // This is a class comment
-
 
     /**
      * This is a field doc comment.
@@ -21,8 +16,7 @@ internal class C {
     /**
      * This is a function doc comment.
      */
-    fun foo() {
+    fun foo() { /* This is a function comment */
-        /* This is a function comment */
     }
 
     //simple one line comment for function
@@ -43,8 +37,6 @@ internal class C {
     /** of
      */
     //
-
-
     /**
      * different
      */
@@ -55,13 +47,13 @@ internal class C {
     /** of
      */
     //
-
-
     /**
      * different
      */
     //comments
     var l = 0
-    /*two*/ /*comments*//*line*/
+    /*two*/
+    /*comments*/
+    /*line*/
     var z = 0
 }

nj2k/testData/newJ2k/issues/kt-19348.kt

@@ -5,10 +5,6 @@ class TestMutltipleCtorsWithJavadoc
  */(private val x: String?) {
     private var y: String? = null
 
-    // ---
-    // Constructors
-    //
-
     /**
      * Javadoc for 2nd ctor
      * @param x
@@ -17,6 +13,7 @@ class TestMutltipleCtorsWithJavadoc
     constructor(x: String?, y: String?) : this(x) {
         this.y = y
     }
-
+    // ---
-
+// Constructors
+//
 }

nj2k/testData/newJ2k/javaStreamsApi/formatting.kt

@@ -3,12 +3,12 @@ import java.util.stream.Collectors
 
 internal class Test {
     fun main(lst: List<Int>) {
-        val newLst = /*before list*/ lst/*after list*/.stream/*before stream*/()/* after stream*/
+        val newLst =  /*before list*/lst /*after list*/.stream /*before stream*/() /* after stream*/
                 .filter { x: Int -> x > 10 }
-                .map { x: Int -> x + 2 }/*some comment*/.distinct/*another comment*/()/* one more comment */.sorted()/*another one comment*/
+                .map { x: Int -> x + 2 } /*some comment*/.distinct /*another comment*/() /* one more comment */.sorted() /*another one comment*/
                 .sorted(Comparator.naturalOrder())
                 .peek { x: Int? -> println(x) }.limit(1)
-                .skip(42)/*skipped*/
+                .skip(42) /*skipped*/ /*collecting one*/
-                /*collecting one*/.collect/*collecting two */(Collectors.toList())/* cool */
+                .collect /*collecting two */(Collectors.toList()) /* cool */
     }
 }

nj2k/testData/newJ2k/methodCallExpression/stringMethods.java

@@ -87,7 +87,6 @@ class A {
         int limit = 5;
         useSplit(s.split("\\s+", limit));
         */
-         */
         s.trim();
         s.concat(" another");
 

nj2k/testData/newJ2k/methodCallExpression/stringMethods.kt

@@ -8,8 +8,6 @@ internal class A {
     fun constructors() {
         String()
         // TODO: new String("original");
-
-
         String(charArrayOf('a', 'b', 'c'))
         String(charArrayOf('b', 'd'), 1, 1)
         String(intArrayOf(32, 65, 127), 0, 3)
@@ -88,8 +86,7 @@ internal class A {
         useSplit(s.split("\\s+", 2));
         int limit = 5;
         useSplit(s.split("\\s+", limit));
-        */
+        */s.trim { it <= ' ' }
-        s.trim { it <= ' ' }
         "$s another"
         s.toByteArray()
         s.toByteArray(Charset.forName("utf-8"))
@@ -121,8 +118,7 @@ internal class A {
     }
 
     fun unsupportedMethods() {
-        val s = "test string"
+        val s = "test string" /* TODO:
-        /* TODO:
         s.indexOf(32);
         s.indexOf(32, 2);
         s.lastIndexOf(32);