New J2K: Preserve comments and line breaks

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

@@ -24,6 +24,7 @@ object ConversionsRunner {
     private fun createRootConversion(context: ConversionContext) =
         batchPipe {
             //Java --> Kotlin conversions
+            +NonCodeElementsConversion()
             +JavaModifiersConversion(context)
             +JavaAnnotationsConversion(context)
             +AnnotationClassConversion(context)
@@ -75,7 +76,7 @@ object ConversionsRunner {
             +ImplicitCastsConversion(context)
             +DefaultNullabilityToNullableConversion()
             +NullableQualifierConversion(context)
-            +InternalClassConversion()
+            +InternalDeclarationConversion()
 
             //Kotlin --> Kotlin conversions
             +InnerClassConversion()

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

@@ -22,6 +22,7 @@ import com.intellij.psi.*
 import com.intellij.psi.JavaTokenType.SUPER_KEYWORD
 import com.intellij.psi.JavaTokenType.THIS_KEYWORD
 import com.intellij.psi.impl.source.tree.ChildRole
+import com.intellij.psi.impl.source.tree.CompositeElement
 import com.intellij.psi.impl.source.tree.java.PsiClassObjectAccessExpressionImpl
 import com.intellij.psi.impl.source.tree.java.PsiLabeledStatementImpl
 import com.intellij.psi.impl.source.tree.java.PsiLiteralExpressionImpl
@@ -34,6 +35,7 @@ import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
 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.j2k.ReferenceSearcher
 import org.jetbrains.kotlin.j2k.ast.Nullability
@@ -69,6 +71,10 @@ class JavaToJKTreeBuilder(
 
     private fun PsiPackageStatement.toJK(): JKPackageDeclaration =
         JKPackageDeclarationImpl(JKNameIdentifierImpl(packageName))
+            .also {
+                it.assignNonCodeElements(this)
+            }
+
 
     private fun PsiImportStatement.toJK(): JKImportStatementImpl {
         val target = resolve()
@@ -77,8 +83,18 @@ class JavaToJKTreeBuilder(
             if (target is KtLightClassForFacade) rawName.replaceAfterLast('.', "*")
             else rawName
         return JKImportStatementImpl(JKNameIdentifierImpl(name))
+            .also {
+                it.assignNonCodeElements(this)
+            }
     }
 
+    private fun PsiIdentifier?.toJK(): JKNameIdentifier =
+        this?.let {
+            JKNameIdentifierImpl(it.text).also {
+                it.assignNonCodeElements(this)
+            }
+        } ?: JKNameIdentifierImpl("")
+
 
     private inner class ExpressionTreeMapper {
         fun PsiExpression?.toJK(): JKExpression {
@@ -118,7 +134,10 @@ class JavaToJKTreeBuilder(
                     throw RuntimeException("Not supported: ${this::class}")
                 }
             }.also {
-                if (this != null) (it as PsiOwner).psi = this
+                if (this != null) {
+                    (it as PsiOwner).psi = this
+                    it.assignNonCodeElements(this)
+                }
             }
         }
 
@@ -131,19 +150,25 @@ class JavaToJKTreeBuilder(
                     is JKJavaVoidType -> JKClassLiteralExpression.LiteralType.JAVA_VOID_TYPE
                     else -> JKClassLiteralExpression.LiteralType.JAVA_CLASS
                 }
-            )
+            ).also {
+                it.assignNonCodeElements(this)
+            }
         }
 
         fun PsiInstanceOfExpression.toJK(): JKKtIsExpression =
             JKKtIsExpressionImpl(operand.toJK(), JKTypeElementImpl(checkType?.type?.toJK(symbolProvider) ?: JKNoTypeImpl))
-
+                .also {
+                    it.assignNonCodeElements(this)
+                }
 
         fun PsiAssignmentExpression.toJK(): JKJavaAssignmentExpression {
             return JKJavaAssignmentExpressionImpl(
                 lExpression.toJK() as? JKAssignableExpression ?: error("Its possible? ${lExpression.toJK().prettyDebugPrintTree()}"),
                 rExpression.toJK(),
                 operationSign.tokenType.toJK()
-            )
+            ).also {
+                it.assignNonCodeElements(this)
+            }
         }
 
         fun PsiBinaryExpression.toJK(): JKExpression {
@@ -157,6 +182,9 @@ class JavaToJKTreeBuilder(
                 else -> operationSign.tokenType
             }
             return JKBinaryExpressionImpl(lOperand.toJK(), rOperand.toJK(), token.toJK())
+                .also {
+                    it.assignNonCodeElements(this)
+                }
         }
 
         fun PsiLiteralExpression.toJK(): JKLiteralExpression {
@@ -173,6 +201,8 @@ class JavaToJKTreeBuilder(
                 JavaTokenType.FLOAT_LITERAL -> JKJavaLiteralExpressionImpl(text, FLOAT)
                 JavaTokenType.DOUBLE_LITERAL -> JKJavaLiteralExpressionImpl(text, DOUBLE)
                 else -> error("Unknown literal element type: ${this.literalElementType}")
+            }.also {
+                it.assignNonCodeElements(this)
             }
         }
 
@@ -188,11 +218,15 @@ class JavaToJKTreeBuilder(
                         (it.expression as JKLiteralExpression).type
                     )
                 } else it
+            }.also {
+                it.assignNonCodeElements(this)
             }
         }
 
         fun PsiPostfixExpression.toJK(): JKExpression {
-            return JKPostfixExpressionImpl(operand.toJK(), operationSign.tokenType.toJK())
+            return JKPostfixExpressionImpl(operand.toJK(), operationSign.tokenType.toJK()).also {
+                it.assignNonCodeElements(this)
+            }
         }
 
         fun PsiLambdaExpression.toJK(): JKExpression {
@@ -201,11 +235,13 @@ class JavaToJKTreeBuilder(
                     when (it) {
                         is PsiExpression -> JKExpressionStatementImpl(it.toJK())
                         is PsiCodeBlock -> JKBlockStatementImpl(with(declarationMapper) { it.toJK() })
-                        else -> JKBlockStatementImpl(JKBodyStub)
+                        else -> JKBlockStatementImpl(JKBodyStubImpl)
                     }
                 },
                 with(declarationMapper) { parameterList.parameters.map { it.toJK() } }
-            )
+            ).also {
+                it.assignNonCodeElements(this)
+            }
         }
 
         private fun JKExpression.qualified(qualifier: JKExpression?) =
@@ -296,6 +332,8 @@ class JavaToJKTreeBuilder(
                 symbol is JKFieldSymbol ->
                     JKFieldAccessExpressionImpl(symbol).qualified(qualifier)
                 else -> TODO(text)
+            }.also {
+                it.assignNonCodeElements(this)
             }
         }
 
@@ -316,14 +354,18 @@ class JavaToJKTreeBuilder(
                 is JKClassSymbol -> JKClassAccessExpressionImpl(symbol)
                 is JKFieldSymbol -> JKFieldAccessExpressionImpl(symbol)
                 else -> TODO()
-            }.qualified(qualifierExpression?.toJK())
+            }.qualified(qualifierExpression?.toJK()).also {
+                it.assignNonCodeElements(this)
+            }
         }
 
         fun PsiArrayInitializerExpression.toJK(): JKExpression {
             return JKJavaNewArrayImpl(
                 initializers.map { it.toJK() },
                 JKTypeElementImpl(type?.toJK(symbolProvider).safeAs<JKJavaArrayType>()?.type ?: JKContextType)
-            )
+            ).also {
+                it.assignNonCodeElements(this)
+            }
         }
 
         fun PsiNewExpression.toJK(): JKExpression {
@@ -371,23 +413,37 @@ class JavaToJKTreeBuilder(
                         with(declarationMapper) { anonymousClass?.createClassBody() } ?: JKEmptyClassBodyImpl()
                     )
                 }
-            return qualifier?.let { JKQualifiedExpressionImpl(it.toJK(), JKJavaQualifierImpl.DOT, newExpression) } ?: newExpression
+            return qualifier?.let {
+                JKQualifiedExpressionImpl(it.toJK(), JKJavaQualifierImpl.DOT, newExpression)
+            } ?: newExpression
         }
 
         fun PsiReferenceParameterList.toJK(): JKTypeArgumentList =
             JKTypeArgumentListImpl(this.typeArguments.map { JKTypeElementImpl(it.toJK(symbolProvider)) })
+                .also {
+                    it.assignNonCodeElements(this)
+                }
 
 
         fun PsiArrayAccessExpression.toJK(): JKExpression {
             return JKArrayAccessExpressionImpl(arrayExpression.toJK(), indexExpression?.toJK() ?: TODO())
+                .also {
+                    it.assignNonCodeElements(this)
+                }
         }
 
         fun PsiTypeCastExpression.toJK(): JKExpression {
             return JKTypeCastExpressionImpl(operand?.toJK() ?: TODO(), castType?.type?.toJK(symbolProvider)?.asTypeElement() ?: TODO())
+                .also {
+                    it.assignNonCodeElements(this)
+                }
         }
 
         fun PsiParenthesizedExpression.toJK(): JKExpression {
             return JKParenthesizedExpressionImpl(expression?.toJK() ?: TODO())
+                .also {
+                    it.assignNonCodeElements(this)
+                }
         }
 
         fun PsiExpressionList.toJK(): JKArgumentList {
@@ -402,10 +458,14 @@ class JavaToJKTreeBuilder(
                             JKPrefixExpressionImpl(
                                 jkExpressions.last(),
                                 JKKtSpreadOperator
-                            )
+                            ).withNonCodeElementsFrom(jkExpressions.last())
                         jkExpressions.dropLast(1) + staredExpression
                     } else jkExpressions
-                } ?: jkExpressions).let { it.toArgumentList() }
+                } ?: jkExpressions)
+                .toArgumentList()
+                .also {
+                    it.assignNonCodeElements(this)
+                }
         }
 
     }
@@ -413,10 +473,17 @@ class JavaToJKTreeBuilder(
     private inner class DeclarationMapper(val expressionTreeMapper: ExpressionTreeMapper) {
         fun PsiTypeParameterList.toJK(): JKTypeParameterList =
             JKTypeParameterListImpl(typeParameters.map { it.toJK() })
+                .also {
+                    it.assignNonCodeElements(this)
+                }
 
         fun PsiTypeParameter.toJK(): JKTypeParameter =
-            JKTypeParameterImpl(JKNameIdentifierImpl(name!!),
+            JKTypeParameterImpl(
-                                extendsListTypes.map { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.Default)) })
+                nameIdentifier.toJK(),
+                extendsListTypes.map { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.Default)) }
+            ).also {
+                it.assignNonCodeElements(this)
+            }
 
         fun PsiClass.toJK(): JKClass {
             val classKind: JKClass.ClassKind = when {
@@ -427,26 +494,39 @@ class JavaToJKTreeBuilder(
             }
 
             fun PsiReferenceList.mapTypes() =
-                this.referencedTypes.map { with(expressionTreeMapper) { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.Default)) } }
+                this.referencedTypes.map { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.Default)) }
 
-            val implTypes = this.implementsList?.mapTypes().orEmpty()
-            val extensionType = this.extendsList?.mapTypes().orEmpty()
             return JKClassImpl(
-                JKNameIdentifierImpl(name!!),
+                nameIdentifier.toJK(),
-                JKInheritanceInfoImpl(extensionType, implTypes),
+                inheritanceInfo(),
                 classKind,
                 typeParameterList?.toJK() ?: JKTypeParameterListImpl(),
                 createClassBody(),
                 annotationList(this),
-                modifiers(),
+                extraModifiers(),
                 visibility(),
                 modality()
             ).also { jkClassImpl ->
                 jkClassImpl.psi = this
                 symbolProvider.provideUniverseSymbol(this, jkClassImpl)
+            }.also {
+                it.assignNonCodeElements(this)
             }
         }
 
+        fun PsiClass.inheritanceInfo(): JKInheritanceInfo {
+            val implTypes =
+                implementsList?.referencedTypes?.map { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.Default)) }.orEmpty()
+            val extensionType =
+                extendsList?.referencedTypes?.map { JKTypeElementImpl(it.toJK(symbolProvider, Nullability.Default)) }.orEmpty()
+            return JKInheritanceInfoImpl(extensionType, implTypes)
+                .also {
+                    if (implementsList != null) {
+                        it.assignNonCodeElements(implementsList!!)
+                    }
+                }
+        }
+
         fun PsiClass.createClassBody() =
             JKClassBodyImpl(
                 children.mapNotNull {
@@ -460,42 +540,80 @@ class JavaToJKTreeBuilder(
                         else -> null
                     }
                 }
-            )
+            ).also {
+                it.leftBrace.assignNonCodeElements(lBrace)
+                it.rightBrace.assignNonCodeElements(rBrace)
+            }
 
         fun PsiClassInitializer.toJK(): JKDeclaration =
-            if (hasModifier(JvmModifier.STATIC))
+            (if (hasModifier(JvmModifier.STATIC))
                 JKJavaStaticInitDeclarationImpl(body.toJK())
-            else JKKtInitDeclarationImpl(body.toJK())
+            else JKKtInitDeclarationImpl(body.toJK()))
+                .also {
+                    it.assignNonCodeElements(this)
+                }
 
 
         fun PsiEnumConstant.toJK(): JKEnumConstant =
             JKEnumConstantImpl(
-                JKNameIdentifierImpl(name),
+                nameIdentifier.toJK(),
                 with(expressionTreeMapper) { argumentList?.toJK() ?: JKArgumentListImpl() },
                 initializingClass?.createClassBody() ?: JKEmptyClassBodyImpl(),
                 JKTypeElementImpl(JKClassTypeImpl(symbolProvider.provideDirectSymbol(containingClass!!) as JKClassSymbol, emptyList()))
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
                 it.psi = this
+            }.also {
+                it.assignNonCodeElements(this)
             }
 
+
         fun PsiMember.modality() =
-            when {
+            modifierList?.children?.mapNotNull { child ->
-                modifierList == null -> Modality.OPEN
+                if (child !is PsiKeyword) return@mapNotNull null
-                hasModifierProperty(PsiModifier.FINAL) -> Modality.FINAL
+                when (child.text) {
-                hasModifierProperty(PsiModifier.ABSTRACT) -> Modality.ABSTRACT
+                    PsiModifier.FINAL -> Modality.FINAL
-                else -> Modality.OPEN
+                    PsiModifier.ABSTRACT -> Modality.ABSTRACT
-            }
+
+                    else -> null
+                }?.let {
+                    JKModalityModifierElementImpl(it).withAssignedNonCodeElements(child)
+                }
+            }?.firstOrNull() ?: JKModalityModifierElementImpl(Modality.OPEN)
+
+
+        fun PsiMember.extraModifiers() =
+            modifierList?.children?.mapNotNull { child ->
+                if (child !is PsiKeyword) return@mapNotNull null
+                when (child.text) {
+                    PsiModifier.NATIVE -> ExtraModifier.NATIVE
+                    PsiModifier.STATIC -> ExtraModifier.STATIC
+                    PsiModifier.STRICTFP -> ExtraModifier.STRICTFP
+                    PsiModifier.SYNCHRONIZED -> ExtraModifier.SYNCHRONIZED
+                    PsiModifier.TRANSIENT -> ExtraModifier.TRANSIENT
+                    PsiModifier.VOLATILE -> ExtraModifier.VOLATILE
+
+                    else -> null
+                }?.let {
+                    JKExtraModifierElementImpl(it).withAssignedNonCodeElements(child)
+                }
+            }.orEmpty()
 
         fun PsiMember.visibility() =
-            when {
+            modifierList?.children?.mapNotNull { child ->
-                modifierList == null -> Visibility.PACKAGE_PRIVATE
+                if (child !is PsiKeyword) return@mapNotNull null
-                hasModifierProperty(PsiModifier.PACKAGE_LOCAL) -> Visibility.PACKAGE_PRIVATE
+                when (child.text) {
-                hasModifierProperty(PsiModifier.PRIVATE) -> Visibility.PRIVATE
+                    PsiModifier.PACKAGE_LOCAL -> Visibility.PACKAGE_PRIVATE
-                hasModifierProperty(PsiModifier.PROTECTED) -> handleProtectedVisibility()
+                    PsiModifier.PRIVATE -> Visibility.PRIVATE
-                hasModifierProperty(PsiModifier.PUBLIC) -> Visibility.PUBLIC
+                    PsiModifier.PROTECTED -> handleProtectedVisibility()
-                else -> Visibility.PACKAGE_PRIVATE
+                    PsiModifier.PUBLIC -> Visibility.PUBLIC
-            }
+
+                    else -> null
+                }?.let {
+                    JKVisibilityModifierElementImpl(it).withAssignedNonCodeElements(child)
+                }
+            }?.firstOrNull() ?: JKVisibilityModifierElementImpl(Visibility.PACKAGE_PRIVATE)
+
 
         private fun PsiMember.handleProtectedVisibility(): Visibility {
             val originalClass = containingClass ?: return Visibility.PROTECTED
@@ -532,44 +650,21 @@ class JavaToJKTreeBuilder(
         }
 
 
-        fun PsiMember.modifiers() =
-            if (modifierList == null) emptyList()
-            else
-                PsiModifier.MODIFIERS
-                    .filter { this.modifierList!!.hasExplicitModifier(it) }
-                    .mapNotNull {
-                        when (it) {
-                            PsiModifier.NATIVE -> ExtraModifier.NATIVE
-                            PsiModifier.STATIC -> ExtraModifier.STATIC
-                            PsiModifier.STRICTFP -> ExtraModifier.STRICTFP
-                            PsiModifier.SYNCHRONIZED -> ExtraModifier.SYNCHRONIZED
-                            PsiModifier.TRANSIENT -> ExtraModifier.TRANSIENT
-                            PsiModifier.VOLATILE -> ExtraModifier.VOLATILE
-
-                            PsiModifier.PROTECTED -> null
-                            PsiModifier.PUBLIC -> null
-                            PsiModifier.PRIVATE -> null
-                            PsiModifier.FINAL -> null
-                            PsiModifier.ABSTRACT -> null
-
-                            else -> TODO("Not yet supported")
-                        }
-                    }
-
-
         fun PsiField.toJK(): JKJavaField {
             return JKJavaFieldImpl(
-                JKTypeElementImpl(type.toJK(symbolProvider)),
+                JKTypeElementImpl(type.toJK(symbolProvider)).withAssignedNonCodeElements(typeElement),
-                JKNameIdentifierImpl(name),
+                nameIdentifier.toJK(),
                 with(expressionTreeMapper) { initializer.toJK() },
                 annotationList(this),
-                modifiers(),
+                extraModifiers(),
                 visibility(),
                 modality(),
-                Mutability.UNKNOWN
+                JKMutabilityModifierElementImpl(Mutability.UNKNOWN)
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
                 it.psi = this
+            }.also {
+                it.assignNonCodeElements(this)
             }
         }
 
@@ -598,7 +693,9 @@ class JavaToJKTreeBuilder(
                         )
                     }
                 }
-            )
+            ).also {
+                it.assignNonCodeElements(this)
+            }
 
         fun PsiDocComment.deprecatedAnnotation(): JKAnnotation? =
             findTagByName("deprecated")?.let { tag ->
@@ -619,6 +716,8 @@ class JavaToJKTreeBuilder(
                         initializers.map { it.toJK() }
                     )
                 else -> TODO(this::class.toString())
+            }.also {
+                it.assignNonCodeElements(this)
             }
 
         fun PsiAnnotationMethod.toJK(): JKJavaAnnotationMethod =
@@ -626,31 +725,40 @@ class JavaToJKTreeBuilder(
                 returnType?.toJK(symbolProvider)?.asTypeElement()
                     ?: JKTypeElementImpl(JKJavaVoidType).takeIf { isConstructor }
                     ?: TODO(),
-                JKNameIdentifierImpl(name),
+                nameIdentifier.toJK(),
                 defaultValue?.toJK() ?: JKStubExpressionImpl()
             ).also {
                 it.psi = this
                 symbolProvider.provideUniverseSymbol(this, it)
+            }.also {
+                it.assignNonCodeElements(this)
             }
 
+
         fun PsiMethod.toJK(): JKJavaMethod {
             return JKJavaMethodImpl(
                 returnType?.toJK(symbolProvider)?.asTypeElement()
                     ?: JKTypeElementImpl(JKJavaVoidType).takeIf { isConstructor }
                     ?: TODO(),
-                JKNameIdentifierImpl(name),
+                nameIdentifier.toJK(),
                 parameterList.parameters.map { it.toJK() },
-                body?.toJK() ?: JKBodyStub,
+                body?.toJK() ?: JKBodyStubImpl,
                 typeParameterList?.toJK() ?: JKTypeParameterListImpl(),
                 annotationList(this),
                 throwsList.referencedTypes.map { JKTypeElementImpl(it.toJK(symbolProvider)) },
-                modifiers(),
+                extraModifiers(),
                 visibility(),
                 modality()
-            ).also {
+            ).also { jkMethod ->
-                it.psi = this
+                jkMethod.psi = this
-                symbolProvider.provideUniverseSymbol(this, it)
+                symbolProvider.provideUniverseSymbol(this, jkMethod)
-            }
+                parameterList.node
+                    ?.safeAs<CompositeElement>()
+                    ?.also {
+                        jkMethod.leftParen.assignNonCodeElements(it.findChildByRoleAsPsiElement(ChildRole.LPARENTH))
+                        jkMethod.rightParen.assignNonCodeElements(it.findChildByRoleAsPsiElement(ChildRole.RPARENTH))
+                    }
+            }.withAssignedNonCodeElements(this)
         }
 
         fun PsiParameter.toJK(): JKParameter {
@@ -660,28 +768,40 @@ class JavaToJKTreeBuilder(
                 else rawType.asTypeElement()
             return JKParameterImpl(
                 type,
-                JKNameIdentifierImpl(name!!),
+                nameIdentifier.toJK(),
                 isVarArgs
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
                 it.psi = this
+            }.also {
+                it.assignNonCodeElements(this)
             }
         }
 
-        fun PsiCodeBlock.toJK(): JKBlock {
+        fun PsiCodeBlock.toJK(): JKBlock =
-            return JKBlockImpl(statements.map { it.toJK() })
+            JKBlockImpl(statements.map { it.toJK() })
-        }
+                .withAssignedNonCodeElements(this)
+                .also {
+                    it.leftBrace.assignNonCodeElements(lBrace)
+                    it.rightBrace.assignNonCodeElements(rBrace)
+                }
+
 
         fun PsiLocalVariable.toJK(): JKLocalVariable =
             JKLocalVariableImpl(
-                JKTypeElementImpl(type.toJK(symbolProvider)),
+                JKTypeElementImpl(type.toJK(symbolProvider)).withAssignedNonCodeElements(typeElement),
-                JKNameIdentifierImpl(this.name ?: TODO()),
+                nameIdentifier.toJK(),
                 with(expressionTreeMapper) { initializer.toJK() },
-                if (hasModifierProperty(PsiModifier.FINAL)) Mutability.IMMUTABLE else Mutability.UNKNOWN,
+                JKMutabilityModifierElementImpl(
+                    if (hasModifierProperty(PsiModifier.FINAL)) Mutability.IMMUTABLE
+                    else Mutability.UNKNOWN
+                ),
                 annotationList(null)
             ).also { i ->
                 symbolProvider.provideUniverseSymbol(this, i)
                 i.psi = this
+            }.also {
+                it.assignNonCodeElements(this)
             }
 
         fun PsiStatement?.toJK(): JKStatement {
@@ -735,12 +855,12 @@ class JavaToJKTreeBuilder(
                         when (statement) {
                             is PsiSwitchLabelStatement ->
                                 cases += if (statement.isDefaultCase)
-                                    JKJavaDefaultSwitchCaseImpl(emptyList())
+                                    JKJavaDefaultSwitchCaseImpl(emptyList()).withAssignedNonCodeElements(statement)
                                 else
                                     JKJavaLabelSwitchCaseImpl(
                                         with(expressionTreeMapper) { statement.caseValue.toJK() },
                                         emptyList()
-                                    )
+                                    ).withAssignedNonCodeElements(statement)
                             else ->
                                 //TODO Handle case then there is no last case
                                 cases.lastOrNull()?.also { it.statements = it.statements + statement.toJK() }
@@ -762,7 +882,7 @@ class JavaToJKTreeBuilder(
                 }
                 is PsiLabeledStatement -> {
                     val (labels, statement) = collectLabels()
-                    JKLabeledStatementImpl(statement.toJK(), labels.map { JKNameIdentifierImpl(it.text) })
+                    JKLabeledStatementImpl(statement.toJK(), labels.map { JKNameIdentifierImpl(it.text) }).asStatement()
                 }
                 is PsiEmptyStatement -> JKEmptyStatementImpl()
                 is PsiThrowStatement ->
@@ -770,24 +890,31 @@ class JavaToJKTreeBuilder(
                 is PsiTryStatement ->
                     JKJavaTryStatementImpl(
                         resourceList?.toList()?.map { (it as PsiLocalVariable).toJK() }.orEmpty(),
-                        tryBlock?.toJK() ?: JKBodyStub,
+                        tryBlock?.toJK() ?: JKBodyStubImpl,
-                        finallyBlock?.toJK() ?: JKBodyStub,
+                        finallyBlock?.toJK() ?: JKBodyStubImpl,
                         catchSections.map { it.toJK() }
                     )
                 is PsiSynchronizedStatement ->
                     JKJavaSynchronizedStatementImpl(
                         with(expressionTreeMapper) { lockExpression?.toJK() } ?: JKStubExpressionImpl(),
-                        body?.toJK() ?: JKBodyStub
+                        body?.toJK() ?: JKBodyStubImpl
                     )
                 else -> TODO("for ${this::class}")
             }.also {
-                if (this != null) (it as PsiOwner).psi = this
+                if (this != null) {
+                    (it as PsiOwner).psi = this
+                    it.assignNonCodeElements(this)
+                }
             }
         }
 
         fun PsiCatchSection.toJK(): JKJavaTryCatchSection =
-            JKJavaTryCatchSectionImpl(parameter?.toJK()!!, catchBlock?.toJK() ?: JKBodyStub)
+            JKJavaTryCatchSectionImpl(parameter?.toJK()!!, catchBlock?.toJK() ?: JKBodyStubImpl)
-                .also { it.psi = this }
+                .also {
+                    it.psi = this
+                }.also {
+                    it.assignNonCodeElements(this)
+                }
     }
 
     //TODO better way than generateSequence.last??
@@ -804,5 +931,79 @@ class JavaToJKTreeBuilder(
             else -> error("Cannot convert non-java file")
         }
 
+    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[this]!!
+            val token = when {
+                this is PsiDocComment ->
+                    JKCommentElementImpl(IdeaDocCommentConverter.convertDocComment(this))
+                this is PsiComment -> JKCommentElementImpl(text)
+                this is PsiWhiteSpace -> JKSpaceElementImpl(text)
+                text == ";" -> null
+                text == "" -> null
+                else -> error("Token should be either token or whitespace")
+            } ?: 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()
+        val rightInnerTokens =
+            if (children.isEmpty()) emptyList()
+            else generateSequence(children.last()) { it.prevSibling }
+                .toNonCodeElements()
+                .reversed()
+
+
+        val leftTokens = generateSequence(prevSibling) { it.prevSibling }
+            .toNonCodeElements()
+            .reversed()
+
+        return (leftInnerTokens + prevNonCodeElementsWithParent()) 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) }
 }
 

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

@@ -24,8 +24,9 @@ import org.jetbrains.kotlin.nj2k.NewCodeBuilder.ParenthesisKind.*
 import org.jetbrains.kotlin.nj2k.conversions.parentOfType
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.impl.*
-import org.jetbrains.kotlin.nj2k.tree.visitors.JKVisitorVoid
+import org.jetbrains.kotlin.nj2k.tree.visitors.JKVisitorWithCommentsPrinting
 import org.jetbrains.kotlin.utils.Printer
+import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class NewCodeBuilder {
 
@@ -41,7 +42,57 @@ class NewCodeBuilder {
         JKClass.ClassKind.COMPANION -> "companion object"
     }
 
-    inner class Visitor : JKVisitorVoid {
+    inner class Visitor : JKVisitorWithCommentsPrinting {
+        private val printedTokens = mutableSetOf<JKNonCodeElement>()
+
+        private fun JKNonCodeElement.createText() =
+            if (this !in printedTokens) {
+                printedTokens += this
+                text
+            } else ""
+
+
+        private fun List<JKNonCodeElement>.createText(): String {
+            val text = joinToString("") { token -> token.createText() }
+            val needNewLine = text.lastIndexOf('\n') < text.lastIndexOf("//")
+            return text +"\n".takeIf { needNewLine }.orEmpty()
+        }
+
+        private fun JKNonCodeElementsListOwner.needPreserveSpacesAfterLastCommit() =
+            this is JKArgument
+                    || this is JKParameter
+                    || safeAs<JKTreeElement>()?.parent is JKArgument
+                    || safeAs<JKTreeElement>()?.parent is JKBinaryExpression
+
+        override fun printLeftNonCodeElements(element: JKNonCodeElementsListOwner) {
+            val text = element.leftNonCodeElements
+                .let {
+                    if (element.needPreserveSpacesAfterLastCommit()) it
+                    else it.dropWhile { it is JKSpaceElement }
+                }.createText()
+            printer.printWithNoIndent(text)
+        }
+
+
+        override fun printRightNonCodeElements(element: JKNonCodeElementsListOwner) {
+            val text = element.rightNonCodeElements
+                .let {
+                    if (element.needPreserveSpacesAfterLastCommit()) it
+                    else it.dropLastWhile { it is JKSpaceElement }
+                }.createText()
+            printer.printWithNoIndent(text)
+        }
+
+        private fun renderTokenElement(tokenElement: JKTokenElement) {
+            printLeftNonCodeElements(tokenElement)
+            printer.printWithNoIndent(tokenElement.text)
+            printRightNonCodeElements(tokenElement)
+        }
+
+        override fun visitModifierElementRaw(modifierElement: JKModifierElement) {
+            printer.printWithNoIndent(modifierElement.modifier.text)
+        }
+
         private fun renderExtraTypeParametersUpperBounds(typeParameterList: JKTypeParameterList) {
             val extraTypeBounds = typeParameterList.typeParameters
                 .filter { it.upperBounds.size > 1 }
@@ -61,33 +112,29 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitTreeElement(treeElement: JKTreeElement) {
+        override fun visitTreeElementRaw(treeElement: JKTreeElement) {
             printer.print("/* !!! Hit visitElement for element type: ${treeElement::class} !!! */")
         }
 
-        override fun visitKtTryExpression(ktTryExpression: JKKtTryExpression) {
+        override fun visitKtTryExpressionRaw(ktTryExpression: JKKtTryExpression) {
             printer.printWithNoIndent("try ")
-            if (ktTryExpression.tryBlock != JKBodyStub) {
+            ktTryExpression.tryBlock.accept(this)
-                printer.block { ktTryExpression.tryBlock.accept(this) }
-            }
             ktTryExpression.catchSections.forEach { it.accept(this) }
-            if (ktTryExpression.finallyBlock != JKBodyStub) {
+            if (ktTryExpression.finallyBlock != JKBodyStubImpl) {
                 printer.printWithNoIndent("finally ")
-                printer.block { ktTryExpression.finallyBlock.accept(this) }
+                ktTryExpression.finallyBlock.accept(this)
             }
         }
 
-        override fun visitKtTryCatchSection(ktTryCatchSection: JKKtTryCatchSection) {
+        override fun visitKtTryCatchSectionRaw(ktTryCatchSection: JKKtTryCatchSection) {
             printer.printWithNoIndent("catch ")
             printer.par {
                 ktTryCatchSection.parameter.accept(this)
             }
-            if (ktTryCatchSection.block != JKBodyStub) {
+            ktTryCatchSection.block.accept(this)
-                printer.block { ktTryCatchSection.block.accept(this) }
-            }
         }
 
-        override fun visitForInStatement(forInStatement: JKForInStatement) {
+        override fun visitForInStatementRaw(forInStatement: JKForInStatement) {
             printer.printWithNoIndent("for (")
             forInStatement.declaration.accept(this)
             printer.printWithNoIndent(" in ")
@@ -100,12 +147,12 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitKtThrowExpression(ktThrowExpression: JKKtThrowExpression) {
+        override fun visitKtThrowExpressionRaw(ktThrowExpression: JKKtThrowExpression) {
             printer.printWithNoIndent("throw ")
             ktThrowExpression.exception.accept(this)
         }
 
-        override fun visitDoWhileStatement(doWhileStatement: JKDoWhileStatement) {
+        override fun visitDoWhileStatementRaw(doWhileStatement: JKDoWhileStatement) {
             printer.printWithNoIndent("do ")
             doWhileStatement.body.accept(this)
             printer.printWithNoIndent(" while (")
@@ -113,11 +160,11 @@ class NewCodeBuilder {
             printer.printWithNoIndent(")")
         }
 
-        override fun visitClassAccessExpression(classAccessExpression: JKClassAccessExpression) {
+        override fun visitClassAccessExpressionRaw(classAccessExpression: JKClassAccessExpression) {
             printer.printWithNoIndent(classAccessExpression.identifier.displayName().escaped())
         }
 
-        override fun visitFile(file: JKFile) {
+        override fun visitFileRaw(file: JKFile) {
             if (file.packageDeclaration.packageName.value.isNotEmpty()) {
                 file.packageDeclaration.accept(this)
             }
@@ -130,32 +177,32 @@ class NewCodeBuilder {
                 .map { it.escaped() }
                 .joinToString(".") { it }
 
-        override fun visitPackageDeclaration(packageDeclaration: JKPackageDeclaration) {
+        override fun visitPackageDeclarationRaw(packageDeclaration: JKPackageDeclaration) {
             printer.printWithNoIndent("package ")
             val packageNameEscaped =
                 packageDeclaration.packageName.value.escapedAsQualifiedName()
             printer.printlnWithNoIndent(packageNameEscaped)
         }
 
-        override fun visitImportStatement(importStatement: JKImportStatement) {
+        override fun visitImportStatementRaw(importStatement: JKImportStatement) {
             printer.printWithNoIndent("import ")
             val importNameEscaped =
                 importStatement.name.value.escapedAsQualifiedName()
             printer.printlnWithNoIndent(importNameEscaped)
         }
 
-        override fun visitBreakStatement(breakStatement: JKBreakStatement) {
+        override fun visitBreakStatementRaw(breakStatement: JKBreakStatement) {
             printer.printWithNoIndent("break")
         }
 
-        override fun visitBreakWithLabelStatement(breakWithLabelStatement: JKBreakWithLabelStatement) {
+        override fun visitBreakWithLabelStatementRaw(breakWithLabelStatement: JKBreakWithLabelStatement) {
             printer.printWithNoIndent("break@")
             breakWithLabelStatement.label.accept(this)
         }
 
         private fun renderModifiersList(modifiersList: JKModifiersListOwner) {
-            renderList(modifiersList.modifiers(), " ") {
+            renderList(modifiersList.modifierElements(), " ") {
-                printer.printWithNoIndent(it.text)
+                it.accept(this)
             }
         }
 
@@ -171,7 +218,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitClass(klass: JKClass) {
+        override fun visitClassRaw(klass: JKClass) {
             klass.annotationList.accept(this)
             if (klass.annotationList.annotations.isNotEmpty()) {
                 printer.println()
@@ -199,7 +246,7 @@ class NewCodeBuilder {
             klass.classBody.accept(this)
         }
 
-        override fun visitInheritanceInfo(inheritanceInfo: JKInheritanceInfo) {
+        override fun visitInheritanceInfoRaw(inheritanceInfo: JKInheritanceInfo) {
             val parentClass = inheritanceInfo.parentOfType<JKClass>()!!
             val isInInterface = parentClass.classKind == JKClass.ClassKind.INTERFACE
             val extendTypes = inheritanceInfo.extends.map { it.type.updateNullability(Nullability.NotNull) }
@@ -243,7 +290,7 @@ class NewCodeBuilder {
         }
 
 
-        override fun visitKtProperty(ktProperty: JKKtProperty) {
+        override fun visitKtPropertyRaw(ktProperty: JKKtProperty) {
             ktProperty.annotationList.accept(this)
             if (ktProperty.annotationList.annotations.isNotEmpty()) {
                 printer.println()
@@ -270,7 +317,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitEnumConstant(enumConstant: JKEnumConstant) {
+        override fun visitEnumConstantRaw(enumConstant: JKEnumConstant) {
             enumConstant.name.accept(this)
             if (enumConstant.arguments.arguments.isNotEmpty()) {
                 printer.par {
@@ -282,22 +329,20 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitKtInitDeclaration(ktInitDeclaration: JKKtInitDeclaration) {
+        override fun visitKtInitDeclarationRaw(ktInitDeclaration: JKKtInitDeclaration) {
             if (ktInitDeclaration.block.statements.isNotEmpty()) {
                 printer.print("init ")
-                printer.block(multiline = true) {
+                ktInitDeclaration.block.accept(this)
-                    ktInitDeclaration.block.accept(this)
-                }
             }
         }
 
-        override fun visitKtIsExpression(ktIsExpression: JKKtIsExpression) {
+        override fun visitKtIsExpressionRaw(ktIsExpression: JKKtIsExpression) {
             ktIsExpression.expression.accept(this)
             printer.printWithNoIndent(" is ")
             ktIsExpression.type.accept(this)
         }
 
-        override fun visitParameter(parameter: JKParameter) {
+        override fun visitParameterRaw(parameter: JKParameter) {
             renderModifiersList(parameter)
             printer.printWithNoIndent(" ")
             parameter.annotationList.accept(this)
@@ -321,7 +366,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitKtAnnotationArrayInitializerExpression(ktAnnotationArrayInitializerExpression: JKKtAnnotationArrayInitializerExpression) {
+        override fun visitKtAnnotationArrayInitializerExpressionRaw(ktAnnotationArrayInitializerExpression: JKKtAnnotationArrayInitializerExpression) {
             printer.print("[")
             renderList(ktAnnotationArrayInitializerExpression.initializers) {
                 it.accept(this)
@@ -329,7 +374,7 @@ class NewCodeBuilder {
             printer.print("]")
         }
 
-        override fun visitForLoopVariable(forLoopVariable: JKForLoopVariable) {
+        override fun visitForLoopVariableRaw(forLoopVariable: JKForLoopVariable) {
             forLoopVariable.name.accept(this)
             if (forLoopVariable.type.present() && forLoopVariable.type.type !is JKContextType) {
                 printer.printWithNoIndent(": ")
@@ -337,7 +382,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitKtFunction(ktFunction: JKKtFunction) {
+        override fun visitKtFunctionRaw(ktFunction: JKKtFunction) {
             printer.printIndent()
             if (ktFunction.annotationList.annotations.isNotEmpty()) {
                 ktFunction.annotationList.accept(this)
@@ -347,21 +392,18 @@ class NewCodeBuilder {
             printer.printWithNoIndent(" fun ")
             ktFunction.typeParameterList.accept(this)
             ktFunction.name.accept(this)
-            printer.printWithNoIndent("(")
+            renderTokenElement(ktFunction.leftParen)
             renderList(ktFunction.parameters) {
                 it.accept(this)
             }
-            printer.printWithNoIndent(")", ": ")
+            renderTokenElement(ktFunction.rightParen)
+            printer.printWithNoIndent(": ")
             ktFunction.returnType.accept(this)
             renderExtraTypeParametersUpperBounds(ktFunction.typeParameterList)
-            if (ktFunction.block !== JKBodyStub) {
+            ktFunction.block.accept(this)
-                printer.block(multiline = ktFunction.block.statements.isNotEmpty()) {
-                    ktFunction.block.accept(this)
-                }
-            }
         }
 
-        override fun visitIfElseExpression(ifElseExpression: JKIfElseExpression) {
+        override fun visitIfElseExpressionRaw(ifElseExpression: JKIfElseExpression) {
             printer.printWithNoIndent("if (")
             ifElseExpression.condition.accept(this)
             printer.printWithNoIndent(")")
@@ -370,34 +412,24 @@ class NewCodeBuilder {
             ifElseExpression.elseBranch.accept(this)
         }
 
-        override fun visitIfStatement(ifStatement: JKIfStatement) {
+        override fun visitIfStatementRaw(ifStatement: JKIfStatement) {
             printer.printWithNoIndent("if (")
             ifStatement.condition.accept(this)
             printer.printWithNoIndent(")")
             if (ifStatement.thenBranch.isEmpty()) {
                 printer.printWithNoIndent(";")
             } else {
-                renderStatementOrBlock(ifStatement.thenBranch)
+                ifStatement.thenBranch.accept(this)
             }
         }
 
-        override fun visitIfElseStatement(ifElseStatement: JKIfElseStatement) {
+        override fun visitIfElseStatementRaw(ifElseStatement: JKIfElseStatement) {
             visitIfStatement(ifElseStatement)
             printer.printWithNoIndent(" else ")
-            renderStatementOrBlock(ifElseStatement.elseBranch)
+            ifElseStatement.elseBranch.accept(this)
         }
 
-        private fun renderStatementOrBlock(statement: JKStatement, multiline: Boolean = false) {
+        override fun visitKtGetterOrSetterRaw(ktGetterOrSetter: JKKtGetterOrSetter) {
-            if (statement is JKBlockStatement) {
-                printer.block(multiline) {
-                    statement.block.accept(this)
-                }
-            } else {
-                statement.accept(this)
-            }
-        }
-
-        override fun visitKtGetterOrSetter(ktGetterOrSetter: JKKtGetterOrSetter) {
             printer.indented {
                 renderModifiersList(ktGetterOrSetter)
                 printer.printWithNoIndent(" ")
@@ -416,11 +448,11 @@ class NewCodeBuilder {
             printer.printlnWithNoIndent()
         }
 
-        override fun visitKtEmptyGetterOrSetter(ktEmptyGetterOrSetter: JKKtEmptyGetterOrSetter) {
+        override fun visitKtEmptyGetterOrSetterRaw(ktEmptyGetterOrSetter: JKKtEmptyGetterOrSetter) {
 
         }
 
-        override fun visitBinaryExpression(binaryExpression: JKBinaryExpression) {
+        override fun visitBinaryExpressionRaw(binaryExpression: JKBinaryExpression) {
             binaryExpression.left.accept(this)
             printer.printWithNoIndent(" ")
             printer.printWithNoIndent(binaryExpression.operator.token.text)
@@ -428,7 +460,7 @@ class NewCodeBuilder {
             binaryExpression.right.accept(this)
         }
 
-        override fun visitTypeParameterList(typeParameterList: JKTypeParameterList) {
+        override fun visitTypeParameterListRaw(typeParameterList: JKTypeParameterList) {
             if (typeParameterList.typeParameters.isNotEmpty()) {
                 printer.par(ANGLE) {
                     renderList(typeParameterList.typeParameters) {
@@ -438,7 +470,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitTypeParameter(typeParameter: JKTypeParameter) {
+        override fun visitTypeParameterRaw(typeParameter: JKTypeParameter) {
             typeParameter.name.accept(this)
             if (typeParameter.upperBounds.size == 1) {
                 printer.printWithNoIndent(" : ")
@@ -446,42 +478,42 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitLiteralExpression(literalExpression: JKLiteralExpression) {
+        override fun visitLiteralExpressionRaw(literalExpression: JKLiteralExpression) {
             printer.printWithNoIndent(literalExpression.literal)
         }
 
-        override fun visitPrefixExpression(prefixExpression: JKPrefixExpression) {
+        override fun visitPrefixExpressionRaw(prefixExpression: JKPrefixExpression) {
             printer.printWithNoIndent(prefixExpression.operator.token.text)
             prefixExpression.expression.accept(this)
         }
 
-        override fun visitThisExpression(thisExpression: JKThisExpression) {
+        override fun visitThisExpressionRaw(thisExpression: JKThisExpression) {
             printer.printWithNoIndent("this")
             thisExpression.qualifierLabel.accept(this)
         }
 
-        override fun visitSuperExpression(superExpression: JKSuperExpression) {
+        override fun visitSuperExpressionRaw(superExpression: JKSuperExpression) {
             printer.printWithNoIndent("super")
             superExpression.qualifierLabel.accept(this)
         }
 
-        override fun visitContinueStatement(continueStatement: JKContinueStatement) {
+        override fun visitContinueStatementRaw(continueStatement: JKContinueStatement) {
             printer.printWithNoIndent("continue")
             continueStatement.label.accept(this)
-            printer.printlnWithNoIndent(" ")
+            printer.printWithNoIndent(" ")
         }
 
-        override fun visitLabelEmpty(labelEmpty: JKLabelEmpty) {
+        override fun visitLabelEmptyRaw(labelEmpty: JKLabelEmpty) {
 
         }
 
-        override fun visitLabelText(labelText: JKLabelText) {
+        override fun visitLabelTextRaw(labelText: JKLabelText) {
             printer.printWithNoIndent("@")
             labelText.label.accept(this)
             printer.printWithNoIndent(" ")
         }
 
-        override fun visitLabeledStatement(labeledStatement: JKLabeledStatement) {
+        override fun visitLabeledStatementRaw(labeledStatement: JKLabeledStatement) {
             for (label in labeledStatement.labels) {
                 label.accept(this)
                 printer.printWithNoIndent("@")
@@ -489,16 +521,16 @@ class NewCodeBuilder {
             labeledStatement.statement.accept(this)
         }
 
-        override fun visitNameIdentifier(nameIdentifier: JKNameIdentifier) {
+        override fun visitNameIdentifierRaw(nameIdentifier: JKNameIdentifier) {
             printer.printWithNoIndent(nameIdentifier.value.escaped())
         }
 
-        override fun visitPostfixExpression(postfixExpression: JKPostfixExpression) {
+        override fun visitPostfixExpressionRaw(postfixExpression: JKPostfixExpression) {
             postfixExpression.expression.accept(this)
             printer.printWithNoIndent(postfixExpression.operator.token.text)
         }
 
-        override fun visitQualifiedExpression(qualifiedExpression: JKQualifiedExpression) {
+        override fun visitQualifiedExpressionRaw(qualifiedExpression: JKQualifiedExpression) {
             qualifiedExpression.receiver.accept(this)
             printer.printWithNoIndent(
                 when (qualifiedExpression.operator) {
@@ -510,26 +542,26 @@ class NewCodeBuilder {
             qualifiedExpression.selector.accept(this)
         }
 
-        override fun visitExpressionList(expressionList: JKExpressionList) {
+        override fun visitExpressionListRaw(expressionList: JKExpressionList) {
             renderList(expressionList.expressions) { it.accept(this) }
         }
 
 
-        override fun visitArgumentList(argumentList: JKArgumentList) {
+        override fun visitArgumentListRaw(argumentList: JKArgumentList) {
             renderList(argumentList.arguments) { it.accept(this) }
         }
 
-        override fun visitArgument(argument: JKArgument) {
+        override fun visitArgumentRaw(argument: JKArgument) {
             argument.value.accept(this)
         }
 
-        override fun visitNamedArgument(namedArgument: JKNamedArgument) {
+        override fun visitNamedArgumentRaw(namedArgument: JKNamedArgument) {
             namedArgument.name.accept(this)
             printer.printWithNoIndent(" = ")
             namedArgument.value.accept(this)
         }
 
-        override fun visitMethodCallExpression(methodCallExpression: JKMethodCallExpression) {
+        override fun visitMethodCallExpressionRaw(methodCallExpression: JKMethodCallExpression) {
             printer.printWithNoIndent(FqName(methodCallExpression.identifier.fqName).shortName().asString().escaped())
             methodCallExpression.typeArgumentList.accept(this)
             printer.par {
@@ -537,7 +569,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitTypeArgumentList(typeArgumentList: JKTypeArgumentList) {
+        override fun visitTypeArgumentListRaw(typeArgumentList: JKTypeArgumentList) {
             if (typeArgumentList.typeArguments.isNotEmpty()) {
                 printer.par(ANGLE) {
                     renderList(typeArgumentList.typeArguments) {
@@ -547,36 +579,36 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitParenthesizedExpression(parenthesizedExpression: JKParenthesizedExpression) {
+        override fun visitParenthesizedExpressionRaw(parenthesizedExpression: JKParenthesizedExpression) {
             printer.par {
                 parenthesizedExpression.expression.accept(this)
             }
         }
 
-        override fun visitDeclarationStatement(declarationStatement: JKDeclarationStatement) {
+        override fun visitDeclarationStatementRaw(declarationStatement: JKDeclarationStatement) {
             renderList(declarationStatement.declaredStatements, { printer.println() }) {
                 it.accept(this)
             }
         }
 
-        override fun visitTypeCastExpression(typeCastExpression: JKTypeCastExpression) {
+        override fun visitTypeCastExpressionRaw(typeCastExpression: JKTypeCastExpression) {
             typeCastExpression.expression.accept(this)
             printer.printWithNoIndent(" as ")
             typeCastExpression.type.accept(this)
         }
 
-        override fun visitWhileStatement(whileStatement: JKWhileStatement) {
+        override fun visitWhileStatementRaw(whileStatement: JKWhileStatement) {
             printer.print("while(")
             whileStatement.condition.accept(this)
             printer.printWithNoIndent(")")
             if (whileStatement.body.isEmpty()) {
                 printer.printWithNoIndent(";")
             } else {
-                renderStatementOrBlock(whileStatement.body, multiline = true)
+                whileStatement.body.accept(this)
             }
         }
 
-        override fun visitLocalVariable(localVariable: JKLocalVariable) {
+        override fun visitLocalVariableRaw(localVariable: JKLocalVariable) {
             printer.printWithNoIndent(" ")
             localVariable.annotationList.accept(this)
             printer.printWithNoIndent(" ")
@@ -593,13 +625,13 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitEmptyStatement(emptyStatement: JKEmptyStatement) {
+        override fun visitEmptyStatementRaw(emptyStatement: JKEmptyStatement) {
         }
 
-        override fun visitStubExpression(stubExpression: JKStubExpression) {
+        override fun visitStubExpressionRaw(stubExpression: JKStubExpression) {
         }
 
-        override fun visitKtConvertedFromForLoopSyntheticWhileStatement(
+        override fun visitKtConvertedFromForLoopSyntheticWhileStatementRaw(
             ktConvertedFromForLoopSyntheticWhileStatement: JKKtConvertedFromForLoopSyntheticWhileStatement
         ) {
             ktConvertedFromForLoopSyntheticWhileStatement.variableDeclaration.accept(this)
@@ -638,7 +670,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitJavaNewExpression(javaNewExpression: JKJavaNewExpression) {
+        override fun visitJavaNewExpressionRaw(javaNewExpression: JKJavaNewExpression) {
             if (javaNewExpression.isAnonymousClass()) {
                 printer.printWithNoIndent("object : ")
             }
@@ -655,62 +687,65 @@ class NewCodeBuilder {
         }
 
 
-        override fun visitClassBody(classBody: JKClassBody) {
+        override fun visitClassBodyRaw(classBody: JKClassBody) {
             val declarationsToPrint = classBody.declarations.filterNot { it is JKKtPrimaryConstructor }
-            printer.block(multiline = true) {
+            renderTokenElement(classBody.leftBrace)
-                val enumConstants = declarationsToPrint.filterIsInstance<JKEnumConstant>()
+            val enumConstants = declarationsToPrint.filterIsInstance<JKEnumConstant>()
-                val otherDeclarations = declarationsToPrint.filterNot { it is JKEnumConstant }
+            val otherDeclarations = declarationsToPrint.filterNot { it is JKEnumConstant }
-                renderEnumConstants(enumConstants)
+            renderEnumConstants(enumConstants)
-                if ((classBody.parent as? JKClass)?.classKind == JKClass.ClassKind.ENUM
+            if ((classBody.parent as? JKClass)?.classKind == JKClass.ClassKind.ENUM
-                    && otherDeclarations.isNotEmpty()
+                && otherDeclarations.isNotEmpty()
-                ) {
+            ) {
-                    printer.printlnWithNoIndent(";")
+                printer.printlnWithNoIndent(";")
-                }
-                if (enumConstants.isNotEmpty() && otherDeclarations.isNotEmpty()) {
-                    printer.println()
-                }
-                renderNonEnumClassDeclarations(otherDeclarations)
             }
+            if (enumConstants.isNotEmpty() && otherDeclarations.isNotEmpty()) {
+                printer.println()
+            }
+            renderNonEnumClassDeclarations(otherDeclarations)
+            renderTokenElement(classBody.rightBrace)
         }
 
-        override fun visitEmptyClassBody(emptyClassBody: JKEmptyClassBody) {}
+        override fun visitEmptyClassBodyRaw(emptyClassBody: JKEmptyClassBody) {}
 
-        override fun visitTypeElement(typeElement: JKTypeElement) {
+        override fun visitTypeElementRaw(typeElement: JKTypeElement) {
             renderType(typeElement.type)
         }
 
-        override fun visitBlock(block: JKBlock) {
+        override fun visitBlockRaw(block: JKBlock) {
+            renderTokenElement(block.leftBrace)
             renderList(block.statements, { printer.println() }) {
                 it.accept(this)
             }
+            renderTokenElement(block.rightBrace)
+        }
+
+        override fun visitBlockStatementWithoutBracketsRaw(blockStatementWithoutBrackets: JKBlockStatementWithoutBrackets) {
+            renderList(blockStatementWithoutBrackets.statements, { printer.println() }) {
+                it.accept(this)
+            }
         }
 
-        override fun visitBlockStatementWithoutBrackets(blockStatementWithoutBrackets: JKBlockStatementWithoutBrackets) {
+        override fun visitExpressionStatementRaw(expressionStatement: JKExpressionStatement) {
-            blockStatementWithoutBrackets.block.accept(this)
-        }
-
-        override fun visitExpressionStatement(expressionStatement: JKExpressionStatement) {
-            printer.printIndent()
             expressionStatement.expression.accept(this)
         }
 
-        override fun visitReturnStatement(returnStatement: JKReturnStatement) {
+        override fun visitReturnStatementRaw(returnStatement: JKReturnStatement) {
             printer.print("return")
             returnStatement.label.accept(this)
             printer.print(" ")
             returnStatement.expression.accept(this)
         }
 
-        override fun visitFieldAccessExpression(fieldAccessExpression: JKFieldAccessExpression) {
+        override fun visitFieldAccessExpressionRaw(fieldAccessExpression: JKFieldAccessExpression) {
             printer.printWithNoIndent(fieldAccessExpression.identifier.name.escaped())
         }
 
-        override fun visitArrayAccessExpression(arrayAccessExpression: JKArrayAccessExpression) {
+        override fun visitArrayAccessExpressionRaw(arrayAccessExpression: JKArrayAccessExpression) {
             arrayAccessExpression.expression.accept(this)
             printer.par(SQUARE) { arrayAccessExpression.indexExpression.accept(this) }
         }
 
-        override fun visitDelegationConstructorCall(delegationConstructorCall: JKDelegationConstructorCall) {
+        override fun visitDelegationConstructorCallRaw(delegationConstructorCall: JKDelegationConstructorCall) {
             delegationConstructorCall.expression.accept(this)
             printer.par {
                 delegationConstructorCall.arguments.accept(this)
@@ -725,7 +760,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitKtConstructor(ktConstructor: JKKtConstructor) {
+        override fun visitKtConstructorRaw(ktConstructor: JKKtConstructor) {
             ktConstructor.annotationList.accept(this)
             if (ktConstructor.annotationList.annotations.isNotEmpty()) {
                 printer.println()
@@ -737,14 +772,10 @@ class NewCodeBuilder {
                 builder.append(" : ")
                 ktConstructor.delegationCall.accept(this)
             }
-            if (ktConstructor.block !== JKBodyStub) {
+            ktConstructor.block.accept(this)
-                printer.block(multiline = true) {
-                    ktConstructor.block.accept(this)
-                }
-            }
         }
 
-        override fun visitKtPrimaryConstructor(ktPrimaryConstructor: JKKtPrimaryConstructor) {
+        override fun visitKtPrimaryConstructorRaw(ktPrimaryConstructor: JKKtPrimaryConstructor) {
             ktPrimaryConstructor.annotationList.accept(this)
             printer.printWithNoIndent(" ")
             renderModifiersList(ktPrimaryConstructor)
@@ -756,6 +787,9 @@ class NewCodeBuilder {
             }
         }
 
+        override fun visitBodyStub(bodyStub: JKBodyStub) {
+        }
+
         private inline fun Printer.indented(block: () -> Unit) {
             this.pushIndent()
             block()
@@ -764,7 +798,6 @@ class NewCodeBuilder {
 
         private inline fun Printer.block(multiline: Boolean = false, crossinline body: () -> Unit) {
             par(CURVED) {
-                if (multiline) printer.printWithNoIndent("\n")
                 indented(body)
             }
         }
@@ -775,7 +808,7 @@ class NewCodeBuilder {
             this.printWithNoIndent(kind.close)
         }
 
-        override fun visitLambdaExpression(lambdaExpression: JKLambdaExpression) {
+        override fun visitLambdaExpressionRaw(lambdaExpression: JKLambdaExpression) {
             printer.par(CURVED) {
                 if (lambdaExpression.statement.statements.size > 1)
                     printer.println()
@@ -787,7 +820,7 @@ class NewCodeBuilder {
 
                 val statement = lambdaExpression.statement
                 if (statement is JKBlockStatement) {
-                    statement.block.accept(this)
+                    renderList(statement.block.statements, { printer.println() }) { it.accept(this) }
                 } else {
                     statement.accept(this)
                 }
@@ -796,13 +829,11 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitBlockStatement(blockStatement: JKBlockStatement) {
+        override fun visitBlockStatementRaw(blockStatement: JKBlockStatement) {
-            printer.par(CURVED) {
+            blockStatement.block.accept(this)
-                blockStatement.block.accept(this)
-            }
         }
 
-        override fun visitKtAssignmentStatement(ktAssignmentStatement: JKKtAssignmentStatement) {
+        override fun visitKtAssignmentStatementRaw(ktAssignmentStatement: JKKtAssignmentStatement) {
             ktAssignmentStatement.field.accept(this)
             printer.printWithNoIndent(" ")
             printer.printWithNoIndent(ktAssignmentStatement.operator.token.text)
@@ -810,7 +841,7 @@ class NewCodeBuilder {
             ktAssignmentStatement.expression.accept(this)
         }
 
-        override fun visitKtWhenStatement(ktWhenStatement: JKKtWhenStatement) {
+        override fun visitKtWhenStatementRaw(ktWhenStatement: JKKtWhenStatement) {
             printer.printWithNoIndent("when(")
             ktWhenStatement.expression.accept(this)
             printer.printWithNoIndent(")")
@@ -821,13 +852,13 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitAnnotationList(annotationList: JKAnnotationList) {
+        override fun visitAnnotationListRaw(annotationList: JKAnnotationList) {
             renderList(annotationList.annotations, " ") {
                 it.accept(this)
             }
         }
 
-        override fun visitAnnotation(annotation: JKAnnotation) {
+        override fun visitAnnotationRaw(annotation: JKAnnotation) {
             printer.printWithNoIndent("@")
             printer.printWithNoIndent(annotation.classSymbol.displayName().escapedAsQualifiedName())
             if (annotation.arguments.isNotEmpty()) {
@@ -837,17 +868,17 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitAnnotationNameParameter(annotationNameParameter: JKAnnotationNameParameter) {
+        override fun visitAnnotationNameParameterRaw(annotationNameParameter: JKAnnotationNameParameter) {
             annotationNameParameter.name.accept(this)
             printer.print(" = ")
             annotationNameParameter.value.accept(this)
         }
 
-        override fun visitAnnotationParameter(annotationParameter: JKAnnotationParameter) {
+        override fun visitAnnotationParameterRaw(annotationParameter: JKAnnotationParameter) {
             annotationParameter.value.accept(this)
         }
 
-        override fun visitClassLiteralExpression(classLiteralExpression: JKClassLiteralExpression) {
+        override fun visitClassLiteralExpressionRaw(classLiteralExpression: JKClassLiteralExpression) {
             if (classLiteralExpression.literalType == JKClassLiteralExpression.LiteralType.JAVA_VOID_TYPE) {
                 printer.printWithNoIndent("Void.TYPE")
             } else {
@@ -861,7 +892,7 @@ class NewCodeBuilder {
             }
         }
 
-        override fun visitKtWhenCase(ktWhenCase: JKKtWhenCase) {
+        override fun visitKtWhenCaseRaw(ktWhenCase: JKKtWhenCase) {
             renderList(ktWhenCase.labels, ", ") {
                 it.accept(this)
             }
@@ -869,11 +900,11 @@ class NewCodeBuilder {
             ktWhenCase.statement.accept(this)
         }
 
-        override fun visitKtElseWhenLabel(ktElseWhenLabel: JKKtElseWhenLabel) {
+        override fun visitKtElseWhenLabelRaw(ktElseWhenLabel: JKKtElseWhenLabel) {
             printer.printWithNoIndent("else")
         }
 
-        override fun visitKtValueWhenLabel(ktValueWhenLabel: JKKtValueWhenLabel) {
+        override fun visitKtValueWhenLabelRaw(ktValueWhenLabel: JKKtValueWhenLabel) {
             ktValueWhenLabel.expression.accept(this)
         }
     }
@@ -889,7 +920,6 @@ class NewCodeBuilder {
     }
 }
 
-
 private inline fun <T> List<T>.headTail(): Pair<T?, List<T>?> {
     val head = this.firstOrNull()
     val tail = if (size <= 1) null else subList(1, size)

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

@@ -7,11 +7,12 @@ package org.jetbrains.kotlin.nj2k
 
 import org.jetbrains.kotlin.nj2k.tree.JKElement
 import org.jetbrains.kotlin.nj2k.tree.impl.JKElementBase
+import org.jetbrains.kotlin.nj2k.tree.withNonCodeElementsFrom
 
 fun <T : JKElement> T.copyTree(): T =
     when (this) {
         is JKElementBase ->
-            this.copy() as T
+            this.copy().withNonCodeElementsFrom(this) as T
         else -> TODO("Not supported+$this.toString()")
     }
 

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.conversions
+
+import org.jetbrains.kotlin.nj2k.tree.*
+import org.jetbrains.kotlin.nj2k.tree.impl.JKBinaryExpressionImpl
+import org.jetbrains.kotlin.nj2k.tree.impl.JKParenthesizedExpressionImpl
+
+
+class AddParenthesisForLineBreaksInBinaryExpression : RecursiveApplicableConversionBase() {
+    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 {
+                it is JKSpaceElement && '\n' in it.text
+            }) {
+            return JKParenthesizedExpressionImpl(
+                JKBinaryExpressionImpl(
+                    element::left.detached(),
+                    element::right.detached(),
+                    element.operator
+                ).withNonCodeElementsFrom(element)
+            )
+        }
+        return recurse(element)
+    }
+
+}

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

@@ -5,8 +5,8 @@
 
 package org.jetbrains.kotlin.nj2k.conversions
 
-import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.j2k.ast.Nullability
+import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.nj2k.toExpression
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.impl.*
@@ -24,8 +24,8 @@ class AnnotationClassConversion(private val context: ConversionContext) : Recurs
             JKStubExpressionImpl(),
             JKAnnotationListImpl(),
             emptyList(),
-            Visibility.PUBLIC,
+            JKVisibilityModifierElementImpl(Visibility.PUBLIC),
-            Modality.FINAL
+            JKModalityModifierElementImpl(Modality.FINAL)
         )
         element.modality = Modality.FINAL
         element.classBody.declarations += constructor

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

@@ -28,7 +28,7 @@ class AssignmentStatementSplitAlsoConversion : RecursiveApplicableConversionBase
                     if (statements.size == 1) statements.map { it.detached(this) }
                     else statements
                 }
-            )
+            ).withNonCodeElementsFrom(this)
         } else null
     }
 
@@ -58,8 +58,8 @@ class AssignmentStatementSplitAlsoConversion : RecursiveApplicableConversionBase
                             variable::type.detached(),
                             variable::name.detached(),
                             innerAssignmentStatement.field.copyTreeAndDetach(),
-                            variable.mutability
+                            variable::mutabilityElement.detached()
-                        )
+                        ).withNonCodeElementsFrom(variable)
                     )
                 )
             }
@@ -72,7 +72,7 @@ class AssignmentStatementSplitAlsoConversion : RecursiveApplicableConversionBase
                 qualifiedExpression::receiver.detached(),
                 innerAssignmentStatement.operator
             ),
-            secondAssignment
+            secondAssignment.withNonCodeElementsFrom(this)
         )
     }
 }

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

@@ -27,14 +27,14 @@ class AssignmentStatementValCreationConversion(val context: ConversionContext) :
                                 JKTypeElementImpl(JKJavaVoidType/*TODO*/),
                                 JKNameIdentifierImpl("arr"),
                                 ex1,
-                                Mutability.IMMUTABLE
+                                JKMutabilityModifierElementImpl(Mutability.IMMUTABLE)
                             ).also {
                                 assignableExpr.expression = JKFieldAccessExpressionImpl(context.symbolProvider.provideUniverseSymbol(it))
                             }, JKLocalVariableImpl(
                                 JKTypeElementImpl(JKJavaVoidType/*TODO*/),
                                 JKNameIdentifierImpl("i"),
                                 ex2,
-                                Mutability.UNKNOWN
+                                JKMutabilityModifierElementImpl(Mutability.UNKNOWN)
                             ).also {
                                 assignableExpr.indexExpression =
                                         JKFieldAccessExpressionImpl(context.symbolProvider.provideUniverseSymbol(it))
@@ -53,7 +53,7 @@ class AssignmentStatementValCreationConversion(val context: ConversionContext) :
                                 JKTypeElementImpl(JKJavaVoidType/*TODO*/),
                                 JKNameIdentifierImpl("arg"),
                                 ex,
-                                Mutability.UNKNOWN
+                                JKMutabilityModifierElementImpl(Mutability.UNKNOWN)
                             ).also {
                                 assignableExpr.receiver = JKFieldAccessExpressionImpl(context.symbolProvider.provideUniverseSymbol(it))
                             }

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

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.nj2k.ConversionContext
-import org.jetbrains.kotlin.nj2k.tree.JKBlock
+import org.jetbrains.kotlin.nj2k.tree.*
-import org.jetbrains.kotlin.nj2k.tree.JKBlockStatement
-import org.jetbrains.kotlin.nj2k.tree.JKDeclaration
-import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.impl.*
 
 class BlockToRunConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
@@ -28,7 +25,7 @@ class BlockToRunConversion(private val context: ConversionContext) : RecursiveAp
             emptyList()
         )
         val call = JKKtCallExpressionImpl(context.symbolProvider.provideByFqName("kotlin.run", true), JKArgumentListImpl(lambda))
-        return recurse(JKExpressionStatementImpl(call))
+        return recurse(JKExpressionStatementImpl(call).withNonCodeElementsFrom(element))
     }
 
 }

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

@@ -43,7 +43,7 @@ class BoxedTypeOperationsConversion(private val context: ConversionContext) : Re
                 "kotlin.${primitiveTypeName.capitalize()}.to${operationType.capitalize()}"
             ),
             JKArgumentListImpl()
-        )
+        ).withNonCodeElementsFrom(methodCallExpression)
     }
 
     companion object {

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

@@ -66,13 +66,13 @@ class BuiltinMembersConversion(private val context: ConversionContext) : Recursi
                         context.symbolProvider.provideByFqNameMulti(fqName),
                         argumentsProvider(from::arguments.detached()),
                         from::typeArgumentList.detached()
-                    )
+                    ).withNonCodeElementsFrom(from)
                 is JKFieldAccessExpression ->
                     JKKtCallExpressionImpl(
                         context.symbolProvider.provideByFqNameMulti(fqName),
                         JKArgumentListImpl(),
                         JKTypeArgumentListImpl()
-                    )
+                    ).withNonCodeElementsFrom(from)
                 else -> error("Bad conversion")
             }
     }
@@ -85,11 +85,11 @@ class BuiltinMembersConversion(private val context: ConversionContext) : Recursi
                 is JKMethodCallExpression ->
                     JKFieldAccessExpressionImpl(
                         context.symbolProvider.provideByFqNameMulti(fqName)
-                    )
+                    ).withNonCodeElementsFrom(from)
                 is JKFieldAccessExpression ->
                     JKFieldAccessExpressionImpl(
                         context.symbolProvider.provideByFqNameMulti(fqName)
-                    )
+                    ).withNonCodeElementsFrom(from)
                 else -> error("Bad conversion")
             }
     }
@@ -109,7 +109,7 @@ class BuiltinMembersConversion(private val context: ConversionContext) : Recursi
                             JKArgumentListImpl(arguments.drop(1)),
                             from::typeArgumentList.detached()
                         )
-                    )
+                    ).withNonCodeElementsFrom(from)
                 }
                 else -> error("Bad conversion")
             }

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

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.nj2k.getCompanion
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.impl.JKAnnotationListImpl
 import org.jetbrains.kotlin.nj2k.tree.impl.JKClassImpl
+import org.jetbrains.kotlin.nj2k.tree.impl.JKModalityModifierElementImpl
 import org.jetbrains.kotlin.nj2k.tree.impl.psi
 
 class ClassToObjectPromotionConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
@@ -49,10 +50,10 @@ class ClassToObjectPromotionConversion(private val context: ConversionContext) :
                             }.map { it.detached(element.classBody) }
                         },
                         JKAnnotationListImpl(),
-                        element.extraModifiers,
+                        element.extraModifierElements,
-                        element.visibility,
+                        element.visibilityElement,
-                        Modality.FINAL
+                        JKModalityModifierElementImpl(Modality.FINAL)
-                    )
+                    ).withNonCodeElementsFrom(element)
                 )
             }
         }

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

@@ -7,7 +7,8 @@ package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.nj2k.tree.*
-import org.jetbrains.kotlin.nj2k.tree.impl.*
+import org.jetbrains.kotlin.nj2k.tree.impl.JKKtConstructorImpl
+import org.jetbrains.kotlin.nj2k.tree.impl.JKStubExpressionImpl
 
 class ConstructorConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
@@ -24,12 +25,12 @@ class ConstructorConversion(private val context: ConversionContext) : RecursiveA
             element.block,
             delegationCall,
             element.annotationList,
-            element.extraModifiers,
+            element.extraModifierElements,
-            element.visibility,
+            element.visibilityElement,
-            element.modality
+            element.modalityElement
         ).also {
             context.symbolProvider.transferSymbol(it, element)
-        }
+        }.withNonCodeElementsFrom(element)
     }
 
     private fun lookupDelegationCall(block: JKBlock): JKDelegationConstructorCall? {

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

@@ -16,10 +16,12 @@ import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 class DefaultArgumentsConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
 
     private fun JKMethod.canNotBeMerged(): Boolean =
-        modality == Modality.ABSTRACT ||
+        modality == Modality.ABSTRACT
-                modality == Modality.OVERRIDE ||
+                || modality == Modality.OVERRIDE
-                context.converter.converterServices.oldServices.referenceSearcher.hasOverrides(psi()!!) ||
+                || hasExtraModifier(ExtraModifier.NATIVE)
-                annotationList.annotations.isNotEmpty()
+                || hasExtraModifier(ExtraModifier.SYNCHRONIZED)
+                || context.converter.converterServices.oldServices.referenceSearcher.hasOverrides(psi()!!)
+                || annotationList.annotations.isNotEmpty()
 
 
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
@@ -39,13 +41,13 @@ class DefaultArgumentsConversion(private val context: ConversionContext) : Recur
             if (calledMethod.parent != method.parent
                 || callee.name != method.name.value
                 || calledMethod.returnType.type != method.returnType.type
-                || call.arguments.expressions.size <= method.parameters.size
+                || call.arguments.arguments.size <= method.parameters.size
             ) {
                 continue
             }
 
 
-            // TODO: Filter by annotations, visibility, modality, modifiers like synchronized
+            // TODO: Filter by annotations, visibility, modality, extraModifiers like synchronized
             if (calledMethod.visibility != method.visibility) continue@checkMethod
             if (calledMethod.canNotBeMerged()) continue
 
@@ -54,8 +56,8 @@ class DefaultArgumentsConversion(private val context: ConversionContext) : Recur
                 val targetParameter = calledMethod.parameters[i]
                 val argument = call.arguments.arguments[i].value
                 if (parameter.name.value != targetParameter.name.value) continue@checkMethod
-//                if (parameter.type.type != targetParameter.type.type) continue@checkMethod
+                if (parameter.type.type != targetParameter.type.type) continue@checkMethod
-//                if (argument !is JKFieldAccessExpression || argument.identifier.target != parameter) continue@checkMethod
+                if (argument !is JKFieldAccessExpression || argument.identifier.target != parameter) continue@checkMethod
             }
 
 

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

@@ -23,12 +23,13 @@ class FieldToPropertyConversion : RecursiveApplicableConversionBase() {
                 JKKtEmptyGetterOrSetterImpl(),
                 JKKtEmptyGetterOrSetterImpl(),
                 element.annotationList,
-                element.extraModifiers,
+                element.extraModifierElements,
-                element.visibility,
+                element.visibilityElement,
-                Modality.FINAL,
+                JKModalityModifierElementImpl(Modality.FINAL),
-                mutability
+                JKMutabilityModifierElementImpl(mutability).withNonCodeElementsFrom(element.modalityElement)
             ).also {
                 it.psi = element.psi
+                it.takeNonCodeElementsFrom(element)
             }
         )
     }

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

@@ -23,8 +23,8 @@ class ForConversion(private val context: ConversionContext) : RecursiveApplicabl
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         if (element !is JKJavaForLoopStatement) return recurse(element)
 
-        convertToForeach(element)?.also { return recurse(it) }
+        convertToForeach(element)?.also { return recurse(it.withNonCodeElementsFrom(element)) }
-        convertToWhile(element)?.also { return recurse(it) }
+        convertToWhile(element)?.also { return recurse(it.withNonCodeElementsFrom(element)) }
 
         return recurse(element)
     }
@@ -69,7 +69,7 @@ class ForConversion(private val context: ConversionContext) : RecursiveApplicabl
                 if (elementPsi.findContinuedStatement()?.toContinuedLoop() != loopStatement.psi<PsiForStatement>()) return recurse(element)
                 val statements = loopStatement.updaters.map { it.copyTreeAndDetach() } + element.copyTreeAndDetach()
                 return if (element.parent is JKBlock)
-                    JKBlockStatementWithoutBracketsImpl(JKBlockImpl(statements))
+                    JKBlockStatementWithoutBracketsImpl(statements)
                 else JKBlockStatementImpl(JKBlockImpl(statements))
             }
         }

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.nj2k.conversions
 import org.jetbrains.kotlin.nj2k.tree.ExtraModifier
 import org.jetbrains.kotlin.nj2k.tree.JKClass
 import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
+import org.jetbrains.kotlin.nj2k.tree.impl.JKExtraModifierElementImpl
 import org.jetbrains.kotlin.nj2k.tree.isLocalClass
 
 class InnerClassConversion : RecursiveApplicableConversionBase() {
@@ -25,14 +26,14 @@ class InnerClassConversion : RecursiveApplicableConversionBase() {
         if (element.classKind == JKClass.ClassKind.COMPANION) return recurseArmed(element, outer)
         if (element.isLocalClass()) return recurseArmed(element, outer)
 
-        val static = element.extraModifiers.find { it == ExtraModifier.STATIC }
+        val static = element.extraModifierElements.find { it.extraModifier == ExtraModifier.STATIC }
         if (static != null) {
-            element.extraModifiers -= static
+            element.extraModifierElements -= static
         } else if (element.classKind != JKClass.ClassKind.INTERFACE &&
             outer.classKind != JKClass.ClassKind.INTERFACE &&
             element.classKind != JKClass.ClassKind.ENUM
         ) {
-            element.extraModifiers += ExtraModifier.INNER
+            element.extraModifierElements += JKExtraModifierElementImpl(ExtraModifier.INNER)
         }
         return recurseArmed(element, element)
     }

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

@@ -22,9 +22,8 @@ class InsertDefaultPrimaryConstructorConversion(private val context: ConversionC
             JKStubExpressionImpl(),
             JKAnnotationListImpl(),
             emptyList(),
-            Visibility.PUBLIC,
+            JKVisibilityModifierElementImpl(Visibility.PUBLIC),
-            Modality.FINAL
+            JKModalityModifierElementImpl(Modality.FINAL)
-
         )
 
         element.classBody.declarations += constructor

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

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.nj2k.tree.*
 
-class InternalClassConversion : RecursiveApplicableConversionBase() {
+class InternalDeclarationConversion : RecursiveApplicableConversionBase() {
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         if (element !is JKVisibilityOwner || element !is JKModalityOwner) return recurse(element)
         val containingClass = element.parentOfType<JKClass>() ?: return recurse(element)
@@ -19,6 +19,16 @@ class InternalClassConversion : RecursiveApplicableConversionBase() {
         ) {
             element.visibility = Visibility.PUBLIC
         }
+
+        if (containingClass.classKind == JKClass.ClassKind.INTERFACE) {
+            element.visibility = Visibility.PUBLIC
+        }
+
+        if (containingClass.classKind == JKClass.ClassKind.ENUM
+            && element is JKKtPrimaryConstructor
+        ) {
+            element.visibility = Visibility.PRIVATE
+        }
         return recurse(element)
     }
 }

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

@@ -19,12 +19,9 @@ package org.jetbrains.kotlin.nj2k.conversions
 import org.jetbrains.kotlin.j2k.ast.Nullability
 import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.nj2k.throwAnnotation
-import org.jetbrains.kotlin.nj2k.tree.JKClassBody
+import org.jetbrains.kotlin.nj2k.tree.*
-import org.jetbrains.kotlin.nj2k.tree.JKJavaMethod
-import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.impl.JKKtFunctionImpl
 import org.jetbrains.kotlin.nj2k.tree.impl.psi
-import org.jetbrains.kotlin.nj2k.tree.updateNullabilityRecursively
 
 class JavaMethodToKotlinFunctionConversion(private val context: ConversionContext) : TransformerBasedConversion() {
     override fun visitTreeElement(element: JKTreeElement) {
@@ -53,13 +50,15 @@ class JavaMethodToKotlinFunctionConversion(private val context: ConversionContex
                                 )
                         }
                     },
-                    declaration.extraModifiers,
+                    declaration.extraModifierElements,
-                    declaration.visibility,
+                    declaration.visibilityElement,
-                    declaration.modality
+                    declaration.modalityElement
                 ).also {
                     it.psi = declaration.psi
                     context.symbolProvider.transferSymbol(it, declaration)
-                }
+                    it.leftParen.takeNonCodeElementsFrom(declaration.leftParen)
+                    it.rightParen.takeNonCodeElementsFrom(declaration.rightParen)
+                }.withNonCodeElementsFrom(declaration)
             } else {
                 declaration
             }

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

@@ -29,25 +29,32 @@ class JavaModifiersConversion(private val context: ConversionContext) : Recursiv
             }
         }
         if (element is JKExtraModifiersOwner && element is JKAnnotationListOwner) {
-            if (ExtraModifier.VOLATILE in element.extraModifiers) {
+            element.elementByModifier(ExtraModifier.VOLATILE)?.also { modifierElement ->
-                element.extraModifiers -= ExtraModifier.VOLATILE
+                element.extraModifierElements -= modifierElement
-                element.annotationList.annotations += jvmAnnotation("Volatile", context.symbolProvider)
+                element.annotationList.annotations +=
+                    jvmAnnotation("Volatile", context.symbolProvider).withNonCodeElementsFrom(modifierElement)
             }
-            if (ExtraModifier.TRANSIENT in element.extraModifiers) {
+
-                element.extraModifiers -= ExtraModifier.TRANSIENT
+            element.elementByModifier(ExtraModifier.TRANSIENT)?.also { modifierElement ->
-                element.annotationList.annotations += jvmAnnotation("Transient", context.symbolProvider)
+                element.extraModifierElements -= modifierElement
+                element.annotationList.annotations +=
+                    jvmAnnotation("Transient", context.symbolProvider).withNonCodeElementsFrom(modifierElement)
             }
-            if (ExtraModifier.STRICTFP in element.extraModifiers) {
+
-                element.extraModifiers -= ExtraModifier.STRICTFP
+            element.elementByModifier(ExtraModifier.STRICTFP)?.also { modifierElement ->
-                element.annotationList.annotations += jvmAnnotation("Strictfp", context.symbolProvider)
+                element.extraModifierElements -= modifierElement
+                element.annotationList.annotations +=
+                    jvmAnnotation("Strictfp", context.symbolProvider).withNonCodeElementsFrom(modifierElement)
             }
-            if (ExtraModifier.SYNCHRONIZED in element.extraModifiers) {
+
-                element.extraModifiers -= ExtraModifier.SYNCHRONIZED
+            element.elementByModifier(ExtraModifier.SYNCHRONIZED)?.also { modifierElement ->
-                element.annotationList.annotations += jvmAnnotation("Synchronized", context.symbolProvider)
+                element.extraModifierElements -= modifierElement
+                element.annotationList.annotations +=
+                    jvmAnnotation("Synchronized", context.symbolProvider).withNonCodeElementsFrom(modifierElement)
             }
-            if (ExtraModifier.NATIVE in element.extraModifiers) {
+
-                element.extraModifiers -= ExtraModifier.NATIVE
+            element.elementByModifier(ExtraModifier.NATIVE)?.also { modifierElement ->
-                element.extraModifiers += ExtraModifier.EXTERNAL
+                modifierElement.extraModifier = ExtraModifier.EXTERNAL
             }
         }
         return recurse(element)

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

@@ -23,8 +23,7 @@ class ModalityConversion(private val context: ConversionContext) : RecursiveAppl
     private fun processClass(klass: JKClass) {
         klass.modality = when {
             klass.classKind == JKClass.ClassKind.ENUM -> Modality.FINAL
-            klass.modality == Modality.OPEN
+            klass.classKind == JKClass.ClassKind.INTERFACE -> Modality.OPEN
-                    && context.converter.settings.openByDefault -> Modality.OPEN
             klass.modality == Modality.OPEN
                     && context.converter.settings.openByDefault -> Modality.OPEN
             klass.modality == Modality.OPEN
@@ -37,16 +36,17 @@ class ModalityConversion(private val context: ConversionContext) : RecursiveAppl
 
     private fun processMethod(method: JKJavaMethod) {
         val psi = method.psi<PsiMethod>()!!
-        val containingClass = method.parentOfType<JKClass>()
+        val containingClass = method.parentOfType<JKClass>()!!
         method.modality = when {
             method.modality != Modality.ABSTRACT
-                    && (method.psi!! as PsiMethod).findSuperMethods().isNotEmpty() -> Modality.OVERRIDE
+                    && psi.findSuperMethods().isNotEmpty() -> Modality.OVERRIDE
             method.modality == Modality.OPEN
                     && context.converter.settings.openByDefault
-                    && containingClass?.modality == Modality.OPEN
+                    && containingClass.modality == Modality.OPEN
                     && method.visibility != Visibility.PRIVATE -> Modality.OPEN
 
             method.modality == Modality.OPEN
+                    && containingClass.classKind != JKClass.ClassKind.INTERFACE
                     && !context.converter.converterServices.oldServices.referenceSearcher.hasOverrides(psi) -> Modality.FINAL
             else -> method.modality
         }

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.conversions
+
+import org.jetbrains.kotlin.nj2k.tree.JKClass
+import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
+
+class NonCodeElementsConversion : RecursiveApplicableConversionBase() {
+    override fun applyToElement(element: JKTreeElement): JKTreeElement {
+        when (element) {
+            is JKClass -> {
+                element.name.rightNonCodeElements += element.inheritance.leftNonCodeElements
+                element.inheritance.leftNonCodeElements = emptyList()
+            }
+        }
+        return recurse(element)
+    }
+}

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

@@ -25,7 +25,7 @@ class OperatorExpressionConversion(private val context: ConversionContext) : Rec
                 val operatorToken = operator.token.toKtToken()
                 val left = applyToElement(element::left.detached()) as JKExpression
                 val right = applyToElement(element::right.detached()) as JKExpression
-                recurse(convertBinaryExpression(left, right, operatorToken))
+                recurse(convertBinaryExpression(left, right, operatorToken).withNonCodeElementsFrom(element))
             }
             is JKPrefixExpression -> {
                 val operand = applyToElement(element::expression.detached()) as JKExpression
@@ -37,7 +37,7 @@ class OperatorExpressionConversion(private val context: ConversionContext) : Rec
                 recurse(kotlinPostfixExpression(operand, operatorToken, context.symbolProvider))
             }
             else -> TODO(element.javaClass.toString())
-        } ?: recurse(element)
+        }.withNonCodeElementsFrom(element)
     }
 
 

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

@@ -6,10 +6,10 @@
 package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.nj2k.ConversionContext
-import org.jetbrains.kotlin.nj2k.asAssignmentFromTarget
-import org.jetbrains.kotlin.nj2k.findUsages
 import org.jetbrains.kotlin.nj2k.hasWritableUsages
-import org.jetbrains.kotlin.nj2k.tree.*
+import org.jetbrains.kotlin.nj2k.tree.JKMethod
+import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
+import org.jetbrains.kotlin.nj2k.tree.Mutability
 import org.jetbrains.kotlin.nj2k.tree.impl.*
 
 
@@ -23,11 +23,13 @@ class ParameterModificationInMethodCallsConversion(private val context: Conversi
                         JKTypeElementImpl(parameter.type.type),
                         JKNameIdentifierImpl(parameter.name.value),
                         JKFieldAccessExpressionImpl(context.symbolProvider.provideUniverseSymbol(parameter)),
-                        Mutability.MUTABLE
+                        JKMutabilityModifierElementImpl(Mutability.MUTABLE)
                     )
                 } else null
             }
-        element.block.statements = listOf(JKDeclarationStatementImpl(newVariables)) + element.block.statements
+        if (newVariables.isNotEmpty()) {
+            element.block.statements = listOf(JKDeclarationStatementImpl(newVariables)) + element.block.statements
+        }
         return recurse(element)
     }
 }

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

@@ -32,8 +32,9 @@ class PrimaryConstructorDetectConversion(private val context: ConversionContext)
         primaryConstructorCandidate.invalidate()
         if (primaryConstructorCandidate.block.statements.isNotEmpty()) {
             val initDeclaration = JKKtInitDeclarationImpl(primaryConstructorCandidate.block)
+                .withNonCodeElementsFrom(primaryConstructorCandidate)
             element.classBody.declarations =
-                    element.classBody.declarations.replace(primaryConstructorCandidate, initDeclaration)
+                element.classBody.declarations.replace(primaryConstructorCandidate, initDeclaration)
         } else {
             element.classBody.declarations -= primaryConstructorCandidate
         }
@@ -44,10 +45,10 @@ class PrimaryConstructorDetectConversion(private val context: ConversionContext)
                 primaryConstructorCandidate.parameters,
                 primaryConstructorCandidate.delegationCall,
                 primaryConstructorCandidate.annotationList,
-                primaryConstructorCandidate.extraModifiers,
+                primaryConstructorCandidate.extraModifierElements,
-                primaryConstructorCandidate.visibility,
+                primaryConstructorCandidate.visibilityElement,
-                primaryConstructorCandidate.modality
+                primaryConstructorCandidate.modalityElement
-            )
+            ).withNonCodeElementsFrom(primaryConstructorCandidate)
 
         context.symbolProvider.transferSymbol(primaryConstructor, primaryConstructorCandidate)
 

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

@@ -5,12 +5,12 @@
 
 package org.jetbrains.kotlin.nj2k.conversions
 
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.impl.JKKtCallExpressionImpl
 import org.jetbrains.kotlin.nj2k.tree.impl.JKMethodSymbol
 import org.jetbrains.kotlin.nj2k.tree.impl.psi
-import org.jetbrains.kotlin.name.ClassId
 
 // TODO: Full special methods conversion
 class PrintlnConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
@@ -41,7 +41,9 @@ class PrintlnConversion(private val context: ConversionContext) : RecursiveAppli
         return JKKtCallExpressionImpl(
             context.symbolProvider.provideDirectSymbol(targetElements.first()) as JKMethodSymbol,
             selector.arguments
-        )
+        ).withNonCodeElementsFrom(element).also {
+            it.leftNonCodeElements += element.receiver.commentsFromInside()
+        }
     }
 
 }

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

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.nj2k.conversions
 import org.jetbrains.kotlin.nj2k.createCompanion
 import org.jetbrains.kotlin.nj2k.getCompanion
 import org.jetbrains.kotlin.nj2k.replace
-import org.jetbrains.kotlin.nj2k.tree.JKClass
+import org.jetbrains.kotlin.nj2k.tree.*
-import org.jetbrains.kotlin.nj2k.tree.JKJavaStaticInitDeclaration
-import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
-import org.jetbrains.kotlin.nj2k.tree.detached
 import org.jetbrains.kotlin.nj2k.tree.impl.JKKtInitDeclarationImpl
 
 class StaticInitDeclarationConversion : RecursiveApplicableConversionBase() {
@@ -36,5 +33,5 @@ class StaticInitDeclarationConversion : RecursiveApplicableConversionBase() {
     }
 
     private fun JKJavaStaticInitDeclaration.toKtInitDeclaration() =
-        JKKtInitDeclarationImpl(::block.detached())
+        JKKtInitDeclarationImpl(::block.detached()).withNonCodeElementsFrom(this)
 }

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

@@ -53,7 +53,7 @@ class StaticMemberAccessConversion(private val context: ConversionContext) : Rec
                 ?.safeAs<KtObjectDeclaration>()
                 ?.isCompanion() == true
             is JKTreeElement ->
-                target.safeAs<JKModalityOwner>()?.modifiers()?.contains(ExtraModifier.STATIC) == true
+                target.safeAs<JKExtraModifiersOwner>()?.hasExtraModifier(ExtraModifier.STATIC) == true
                         || target.parentOfType<JKClass>()?.classKind == JKClass.ClassKind.OBJECT
             else -> false
         }

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

@@ -7,15 +7,13 @@ package org.jetbrains.kotlin.nj2k.conversions
 
 import org.jetbrains.kotlin.nj2k.getOrCreateCompainonObject
 import org.jetbrains.kotlin.nj2k.tree.*
-import org.jetbrains.kotlin.nj2k.tree.impl.*
 
 class StaticsToCompanionExtractConversion : RecursiveApplicableConversionBase() {
     override fun applyToElement(element: JKTreeElement): JKTreeElement {
         if (element !is JKClass) return recurse(element)
         if (element.classKind == JKClass.ClassKind.COMPANION || element.classKind == JKClass.ClassKind.OBJECT) return element
         val statics = element.declarationList.filter { declaration ->
-            declaration is JKExtraModifiersOwner &&
+            declaration is JKExtraModifiersOwner && declaration.hasExtraModifier(ExtraModifier.STATIC)
-                    declaration.extraModifiers.any { it == ExtraModifier.STATIC }
         }
         if (statics.isEmpty()) return recurse(element)
         val companion = element.getOrCreateCompainonObject()
@@ -23,7 +21,7 @@ class StaticsToCompanionExtractConversion : RecursiveApplicableConversionBase()
         element.classBody.declarations -= statics
         companion.classBody.declarations += statics.onEach { declaration ->
             (declaration as JKExtraModifiersOwner)
-            declaration.extraModifiers -= ExtraModifier.STATIC
+            declaration.extraModifierElements -= declaration.elementByModifier(ExtraModifier.STATIC)!!
         }
         return recurse(element)
     }

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

@@ -49,7 +49,7 @@ class SwitchStatementConversion(private val context: ConversionContext) : Recurs
                                 statement.block.statements
                                     .takeWhile { !isSwitchBreak(it) }
                                     .map { it.copyTreeAndDetach() }
-                            )
+                            ).withNonCodeElementsFrom(statement)
                         isSwitchBreak(statement) -> null
                         else -> statement.copyTreeAndDetach()
                     }

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.nj2k.ConversionContext
 import org.jetbrains.kotlin.nj2k.tree.JKJavaSynchronizedStatement
 import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.impl.*
+import org.jetbrains.kotlin.nj2k.tree.withNonCodeElementsFrom
 
 
 class SynchronizedStatementConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
@@ -26,7 +27,7 @@ class SynchronizedStatementConversion(private val context: ConversionContext) :
                     element.lockExpression,
                     lambdaBody
                 )
-            )
+            ).withNonCodeElementsFrom(element)
         return recurse(JKExpressionStatementImpl(synchronizedCall))
     }
 

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

@@ -5,10 +5,12 @@
 
 package org.jetbrains.kotlin.nj2k.conversions
 
-import org.jetbrains.kotlin.nj2k.*
 import org.jetbrains.kotlin.j2k.ast.Nullability
+import org.jetbrains.kotlin.nj2k.ConversionContext
+import org.jetbrains.kotlin.nj2k.copyTreeAndDetach
 import org.jetbrains.kotlin.nj2k.tree.*
 import org.jetbrains.kotlin.nj2k.tree.impl.*
+import org.jetbrains.kotlin.nj2k.useExpression
 
 
 class TryStatementConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
@@ -26,7 +28,7 @@ class TryStatementConversion(private val context: ConversionContext) : Recursive
                 tryStatement::finallyBlock.detached(),
                 tryStatement.catchSections.flatMap(::convertCatchSection)
             )
-        )
+        ).withNonCodeElementsFrom(tryStatement)
 
     private fun convertTryStatementWithResources(tryStatement: JKJavaTryStatement): JKStatement {
         val body =
@@ -41,7 +43,7 @@ class TryStatementConversion(private val context: ConversionContext) : Recursive
                     tryStatement::finallyBlock.detached(),
                     tryStatement.catchSections.flatMap(::convertCatchSection)
                 )
-            )
+            ).withNonCodeElementsFrom(tryStatement)
         } else body
     }
 
@@ -74,7 +76,7 @@ class TryStatementConversion(private val context: ConversionContext) : Recursive
             JKKtTryCatchSectionImpl(
                 parameter,
                 javaCatchSection.block.copyTreeAndDetach()
-            )
+            ).withNonCodeElementsFrom(javaCatchSection)
         }
     }
 }

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

@@ -39,7 +39,7 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
                     .fixRawType(element)
                     .mapType(element)
                     .refineNullability(element)
-                JKTypeElementImpl(newType)
+                JKTypeElementImpl(newType).withNonCodeElementsFrom(element)
             }
             is JKJavaNewExpression -> {
                 val newClassSymbol = element.classSymbol.mapClassSymbol(null)
@@ -49,7 +49,7 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
                         element::arguments.detached(),
                         element::typeArgumentList.detached().fixTypeArguments(newClassSymbol),
                         element::classBody.detached()
-                    )
+                    ).withNonCodeElementsFrom(element)
                 )
             }
             else -> recurse(element)

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

@@ -384,8 +384,7 @@ fun JKLiteralExpression.fixLiteral(expectedType: JKLiteralExpression.LiteralType
         JKLiteralExpression.LiteralType.CHAR -> convertCharLiteral(literal)
         JKLiteralExpression.LiteralType.STRING -> convertStringLiteral(literal)
         else -> this
-    }
+    }.withNonCodeElementsFrom(this)
-
 
 private fun convertDoubleLiteral(text: String): JKKtLiteralExpression {
     var newText =
@@ -400,6 +399,7 @@ private fun convertDoubleLiteral(text: String): JKKtLiteralExpression {
     if (newText.endsWith("."))
         newText += "0"
 
+
     return JKKtLiteralExpressionImpl(
         newText,
         JKLiteralExpression.LiteralType.DOUBLE
@@ -491,8 +491,8 @@ fun createCompanion(declarations: List<JKDeclaration>): JKClass =
         JKClassBodyImpl(declarations),
         JKAnnotationListImpl(),
         emptyList(),
-        Visibility.PUBLIC,
+        JKVisibilityModifierElementImpl(Visibility.PUBLIC),
-        Modality.FINAL
+        JKModalityModifierElementImpl(Modality.FINAL)
     )
 
 fun JKClass.getCompanion(): JKClass? =
@@ -508,8 +508,8 @@ fun JKClass.getOrCreateCompainonObject(): JKClass =
             JKClassBodyImpl(),
             JKAnnotationListImpl(),
             emptyList(),
-            Visibility.PUBLIC,
+            JKVisibilityModifierElementImpl(Visibility.PUBLIC),
-            Modality.FINAL
+            JKModalityModifierElementImpl(Modality.FINAL)
         ).also { classBody.declarations += it }
 
 fun runExpression(body: JKStatement, symbolProvider: JKSymbolProvider): JKExpression {
@@ -523,6 +523,9 @@ fun runExpression(body: JKStatement, symbolProvider: JKSymbolProvider): JKExpres
     )
 }
 
+fun JKTreeElement.asQualifierWithThisAsSelector(): JKQualifiedExpression? =
+    parent?.safeAs<JKQualifiedExpression>()
+        ?.takeIf { it.selector == this }
 
 fun JKAnnotationMemberValue.toExpression(symbolProvider: JKSymbolProvider): JKExpression {
     fun handleAnnotationParameter(element: JKTreeElement): JKTreeElement =
@@ -563,6 +566,17 @@ fun JKAnnotationMemberValue.toExpression(symbolProvider: JKSymbolProvider): JKEx
     ) as JKExpression
 }
 
+
+fun JKExpression.asLiteralTextWithPrefix(): String? =
+    when {
+        this is JKPrefixExpression
+                && (operator.token.text == "+" || operator.token.text == "-")
+                && expression is JKLiteralExpression
+        -> operator.token.text + expression.cast<JKLiteralExpression>().literal
+        this is JKLiteralExpression -> literal
+        else -> null
+    }
+
 inline fun JKClass.primaryConstructor(): JKKtPrimaryConstructor? =
     classBody.declarations.firstIsInstanceOrNull()
 

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.nj2k.tree.impl
 
 import org.jetbrains.kotlin.nj2k.tree.JKBranchElement
 import org.jetbrains.kotlin.nj2k.tree.JKElement
+import org.jetbrains.kotlin.nj2k.tree.JKNonCodeElement
 import org.jetbrains.kotlin.nj2k.tree.JKTreeElement
 import org.jetbrains.kotlin.nj2k.tree.visitors.JKVisitor
 import kotlin.properties.ReadWriteProperty
@@ -38,9 +39,12 @@ private class JKListChild<T : JKElement>(val value: Int) : ReadWriteProperty<JKB
 }
 
 abstract class JKElementBase : JKTreeElement, Cloneable {
+    override var leftNonCodeElements: List<JKNonCodeElement> = emptyList()
+    override var rightNonCodeElements: List<JKNonCodeElement> = emptyList()
+
     override var parent: JKElement? = null
 
-    final override fun detach(from: JKElement) {
+    override fun detach(from: JKElement) {
         val prevParent = parent
         require(from == prevParent)
         parent = null
@@ -51,7 +55,7 @@ abstract class JKElementBase : JKTreeElement, Cloneable {
 
     }
 
-    final override fun attach(to: JKElement) {
+    override fun attach(to: JKElement) {
         check(parent == null)
         parent = to
         onAttach()
@@ -71,13 +75,13 @@ abstract class JKElementBase : JKTreeElement, Cloneable {
 
 abstract class JKBranchElementBase : JKElementBase(), JKBranchElement {
     private var childNum = 0
+
     protected fun <T : JKTreeElement, U : T> child(v: U): ReadWriteProperty<JKBranchElementBase, T> {
         children.add(childNum, v)
         v.attach(this)
         return JKChild(childNum++)
     }
 
-
     protected inline fun <reified T : JKTreeElement> children(): ReadWriteProperty<JKBranchElementBase, List<T>> {
         return children(emptyList())
     }
@@ -88,7 +92,7 @@ abstract class JKBranchElementBase : JKElementBase(), JKBranchElement {
         return JKListChild(childNum++)
     }
 
-    final override fun <D> acceptChildren(visitor: JKVisitor<Unit, D>, data: D) {
+    override fun <D> acceptChildren(visitor: JKVisitor<Unit, D>, data: D) {
         forEachChild { it.accept(visitor, data) }
     }
 

nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/j.kt

@@ -31,18 +31,22 @@ class JKJavaFieldImpl(
     name: JKNameIdentifier,
     initializer: JKExpression,
     annotationList: JKAnnotationList,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality,
+    modalityElement: JKModalityModifierElement,
-    override var mutability: Mutability
+    mutabilityElement: JKMutabilityModifierElement
-) : JKJavaField,
+) : JKJavaField(), PsiOwner by PsiOwnerImpl() {
-    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaField(this, data)
 
     override var annotationList: JKAnnotationList by child(annotationList)
     override var initializer: JKExpression by child(initializer)
     override var type by child(type)
     override var name: JKNameIdentifier by child(name)
+
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
+    override var mutabilityElement by child(mutabilityElement)
 }
 
 class JKJavaMethodImpl(
@@ -53,10 +57,10 @@ class JKJavaMethodImpl(
     typeParameterList: JKTypeParameterList,
     annotationList: JKAnnotationList,
     throwsList: List<JKTypeElement>,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality
+    modalityElement: JKModalityModifierElement
-) : JKJavaMethod, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKJavaMethod(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaMethod(this, data)
 
     override var returnType: JKTypeElement by child(returnType)
@@ -66,6 +70,10 @@ class JKJavaMethodImpl(
     override var typeParameterList: JKTypeParameterList by child(typeParameterList)
     override var annotationList: JKAnnotationList by child(annotationList)
     override var throwsList: List<JKTypeElement> by children(throwsList)
+
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
 }
 
 class JKJavaLiteralExpressionImpl(
@@ -182,6 +190,15 @@ class JKJavaOperatorImpl private constructor(psiToken: IElementType) : JKOperato
             else -> 6 /* simple name */
         }
 
+//    override val priority: Int
+//        get() = when (token.psiToken) {
+//            JavaTokenType.ASTERISK, JavaTokenType.DIV, JavaTokenType.PERC -> 12
+//            JavaTokenType.PLUS, JavaTokenType.MINUS -> 11
+//            JavaTokenType.GTGTGT, JavaTokenType.GTGT, JavaTokenType.LTLT -> 10
+//            JavaTokenType.GT, JavaTokenType.LT, JavaTokenType.GE, JavaTokenType.LE -> 9
+//            JavaTokenType.EQEQ, JavaTokenType.NE, KtTokens.EQEQEQ, KtTokens.EXCLEQEQEQ ->
+//        }
+
     companion object {
         val tokenToOperator =
             (OPERATION_BIT_SET.types + arrayOf(KtTokens.EQEQEQ, KtTokens.EXCLEQEQEQ))
@@ -206,12 +223,12 @@ class JKJavaMethodCallExpressionImpl(
     override var typeArgumentList: JKTypeArgumentList by child(typeArgumentList)
 }
 
-class JKClassBodyImpl(declarations: List<JKDeclaration> = emptyList()) : JKClassBody, JKBranchElementBase() {
+class JKClassBodyImpl(declarations: List<JKDeclaration> = emptyList()) : JKClassBody() {
     override var declarations: List<JKDeclaration> by children(declarations)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitClassBody(this, data)
 }
 
-class JKEmptyClassBodyImpl : JKEmptyClassBody, JKBranchElementBase() {
+class JKEmptyClassBodyImpl : JKEmptyClassBody() {
     override var declarations: List<JKDeclaration> by children(emptyList())
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitEmptyClassBody(this, data)
@@ -234,13 +251,15 @@ class JKJavaDefaultNewExpressionImpl(
     override val classSymbol: JKClassSymbol
 ) : JKJavaDefaultNewExpression, JKElementBase(), PsiOwner by PsiOwnerImpl()
 
-class JKJavaNewEmptyArrayImpl(initializer: List<JKExpression>, type: JKTypeElement) : JKJavaNewEmptyArray, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKJavaNewEmptyArrayImpl(initializer: List<JKExpression>, type: JKTypeElement) : JKJavaNewEmptyArray, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override val type by child(type)
     override var initializer by children(initializer)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaNewEmptyArray(this, data)
 }
 
-class JKJavaNewArrayImpl(initializer: List<JKExpression>, type: JKTypeElement) : JKJavaNewArray, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKJavaNewArrayImpl(initializer: List<JKExpression>, type: JKTypeElement) : JKJavaNewArray, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override val type by child(type)
     override var initializer by children(initializer)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaNewArray(this, data)
@@ -280,14 +299,15 @@ class JKJavaDisjunctionTypeImpl(
 class JKReturnStatementImpl(
     expression: JKExpression,
     label: JKLabel = JKLabelEmptyImpl()
-) : JKBranchElementBase(), JKReturnStatement, PsiOwner by PsiOwnerImpl() {
+) : JKReturnStatement(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitReturnStatement(this, data)
 
     override val expression by child(expression)
     override var label by child(label)
 }
 
-class JKJavaAssertStatementImpl(condition: JKExpression, description: JKExpression) : JKJavaAssertStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKJavaAssertStatementImpl(condition: JKExpression, description: JKExpression) : JKJavaAssertStatement(),
+    PsiOwner by PsiOwnerImpl() {
     override val description by child(description)
     override val condition by child(condition)
 
@@ -295,7 +315,7 @@ class JKJavaAssertStatementImpl(condition: JKExpression, description: JKExpressi
 }
 
 class JKJavaForLoopStatementImpl(initializer: JKStatement, condition: JKExpression, updaters: List<JKStatement>, body: JKStatement) :
-    JKJavaForLoopStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+    JKJavaForLoopStatement(), PsiOwner by PsiOwnerImpl() {
     override var body by child(body)
     override var updaters by children(updaters)
     override var condition by child(condition)
@@ -330,13 +350,14 @@ class JKJavaAssignmentExpressionImpl(
 class JKJavaSwitchStatementImpl(
     expression: JKExpression,
     cases: List<JKJavaSwitchCase>
-) : JKJavaSwitchStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKJavaSwitchStatement(), PsiOwner by PsiOwnerImpl() {
     override var expression: JKExpression by child(expression)
     override var cases: List<JKJavaSwitchCase> by children(cases)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaSwitchStatement(this, data)
 }
 
-class JKJavaDefaultSwitchCaseImpl(statements: List<JKStatement>) : JKJavaDefaultSwitchCase, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKJavaDefaultSwitchCaseImpl(statements: List<JKStatement>) : JKJavaDefaultSwitchCase, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override var statements: List<JKStatement> by children(statements)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaDefaultSwitchCase(this, data)
 }
@@ -350,7 +371,7 @@ class JKJavaLabelSwitchCaseImpl(
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaLabelSwitchCase(this, data)
 }
 
-class JKJavaThrowStatementImpl(exception: JKExpression) : JKJavaThrowStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKJavaThrowStatementImpl(exception: JKExpression) : JKJavaThrowStatement(), PsiOwner by PsiOwnerImpl() {
     override var exception: JKExpression by child(exception)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaThrowStatement(this, data)
 }
@@ -360,8 +381,7 @@ class JKJavaTryStatementImpl(
     tryBlock: JKBlock,
     finallyBlock: JKBlock,
     catchSections: List<JKJavaTryCatchSection>
-) : JKJavaTryStatement,
+) : JKJavaTryStatement(), PsiOwner by PsiOwnerImpl() {
-    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var resourceDeclarations: List<JKDeclaration> by children(resourceDeclarations)
     override var tryBlock: JKBlock by child(tryBlock)
     override var finallyBlock: JKBlock by child(finallyBlock)
@@ -381,7 +401,7 @@ class JKJavaTryCatchSectionImpl(
 class JKJavaSynchronizedStatementImpl(
     lockExpression: JKExpression,
     body: JKBlock
-) : JKJavaSynchronizedStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKJavaSynchronizedStatement(), PsiOwner by PsiOwnerImpl() {
     override val lockExpression: JKExpression by child(lockExpression)
     override val body: JKBlock by child(body)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaSynchronizedStatement(this, data)
@@ -391,17 +411,14 @@ class JKJavaAnnotationMethodImpl(
     returnType: JKTypeElement,
     name: JKNameIdentifier,
     defaultValue: JKAnnotationMemberValue
-) : JKJavaAnnotationMethod, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKJavaAnnotationMethod(), PsiOwner by PsiOwnerImpl() {
     override var returnType: JKTypeElement by child(returnType)
     override var name: JKNameIdentifier by child(name)
     override var parameters: List<JKParameter> by children()
     override var defaultValue: JKAnnotationMemberValue by child(defaultValue)
-    override var block: JKBlock by child(JKBodyStub)
+    override var block: JKBlock by child(JKBodyStubImpl)
     override var typeParameterList: JKTypeParameterList by child(JKTypeParameterListImpl())
     override var annotationList: JKAnnotationList by child(JKAnnotationListImpl())
-    override var extraModifiers: List<ExtraModifier> = emptyList()
-    override var visibility: Visibility = Visibility.PUBLIC
-    override var modality: Modality = Modality.FINAL
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaAnnotationMethod(this, data)
 }
@@ -414,7 +431,8 @@ class JKKtAnnotationArrayInitializerExpressionImpl(initializers: List<JKAnnotati
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtAnnotationArrayInitializerExpression(this, data)
 }
-class JKJavaStaticInitDeclarationImpl(block: JKBlock) : JKJavaStaticInitDeclaration, JKBranchElementBase() {
+
+class JKJavaStaticInitDeclarationImpl(block: JKBlock) : JKJavaStaticInitDeclaration() {
     override var block: JKBlock by child(block)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaStaticInitDeclaration(this, data)
 }

nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/jk.kt

@@ -50,10 +50,10 @@ class JKClassImpl(
     typeParameterList: JKTypeParameterList,
     classBody: JKClassBody,
     annotationList: JKAnnotationList,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality
+    modalityElement: JKModalityModifierElement
-) : JKClass, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKClass(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitClass(this, data)
 
     override var name by child(name)
@@ -61,6 +61,10 @@ class JKClassImpl(
     override var typeParameterList: JKTypeParameterList by child(typeParameterList)
     override var classBody: JKClassBody by child(classBody)
     override var annotationList: JKAnnotationList by child(annotationList)
+
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
 }
 
 class JKNameIdentifierImpl(override val value: String) : JKNameIdentifier, JKElementBase(), PsiOwner by PsiOwnerImpl() {
@@ -72,7 +76,7 @@ class JKForLoopVariableImpl(
     name: JKNameIdentifier,
     initializer: JKExpression,
     annotationList: JKAnnotationList = JKAnnotationListImpl()
-) : JKForLoopVariable, JKBranchElementBase() {
+) : JKForLoopVariable() {
     override var initializer by child(initializer)
     override var name by child(name)
     override var type by child(type)
@@ -88,16 +92,18 @@ class JKParameterImpl(
     override var isVarArgs: Boolean = false,
     initializer: JKExpression = JKStubExpressionImpl(),
     annotationList: JKAnnotationList = JKAnnotationListImpl()
-) : JKParameter, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKParameter(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitParameter(this, data)
 
     override var initializer by child(initializer)
     override var name by child(name)
     override var type by child(type)
+    override var annotationList by child(annotationList)
 }
 
-class JKBlockImpl(statements: List<JKStatement> = emptyList()) : JKBlock, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKBlockImpl(statements: List<JKStatement> = emptyList()) : JKBlock(), PsiOwner by PsiOwnerImpl() {
     constructor(vararg statements: JKStatement) : this(statements.toList())
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBlock(this, data)
 
     override var statements by children(statements)
@@ -114,21 +120,24 @@ class JKBinaryExpressionImpl(
 }
 
 
-
+class JKPrefixExpressionImpl(expression: JKExpression, override var operator: JKOperator) : JKPrefixExpression, JKBranchElementBase(),
-class JKPrefixExpressionImpl(expression: JKExpression, override var operator: JKOperator) : JKPrefixExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitPrefixExpression(this, data)
 
     override var expression by child(expression)
 }
 
-class JKPostfixExpressionImpl(expression: JKExpression, override var operator: JKOperator) : JKPostfixExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKPostfixExpressionImpl(expression: JKExpression, override var operator: JKOperator) : JKPostfixExpression, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitPostfixExpression(this, data)
 
     override var expression by child(expression)
 }
 
-class JKExpressionListImpl(expressions: List<JKExpression> = emptyList()) : JKExpressionList, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKExpressionListImpl(expressions: List<JKExpression> = emptyList()) : JKExpressionList, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     constructor(vararg expresions: JKExpression) : this(expresions.asList())
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitExpressionList(this, data)
 
     override var expressions by children(expressions)
@@ -145,16 +154,19 @@ class JKQualifiedExpressionImpl(
     override var selector: JKExpression by child(selector)
 }
 
-class JKExpressionStatementImpl(expression: JKExpression) : JKExpressionStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKExpressionStatementImpl(expression: JKExpression) : JKExpressionStatement(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitExpressionStatement(this, data)
 
     override val expression: JKExpression by child(expression)
 }
 
-class JKDeclarationStatementImpl(declaredStatements: List<JKDeclaration>) : JKDeclarationStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKDeclarationStatementImpl(declaredStatements: List<JKDeclaration>) : JKDeclarationStatement(),
+    PsiOwner by PsiOwnerImpl() {
     override val declaredStatements by children(declaredStatements)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitDeclarationStatement(this, data)
+
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
 }
 
 class JKArrayAccessExpressionImpl(
@@ -167,13 +179,15 @@ class JKArrayAccessExpressionImpl(
     override var indexExpression: JKExpression by child(indexExpression)
 }
 
-class JKParenthesizedExpressionImpl(expression: JKExpression) : JKParenthesizedExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKParenthesizedExpressionImpl(expression: JKExpression) : JKParenthesizedExpression, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitParenthesizedExpression(this, data)
 
     override var expression: JKExpression by child(expression)
 }
 
-class JKTypeCastExpressionImpl(override var expression: JKExpression, type: JKTypeElement) : JKTypeCastExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKTypeCastExpressionImpl(override var expression: JKExpression, type: JKTypeElement) : JKTypeCastExpression, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitTypeCastExpression(this, data)
 
     override var type by child(type)
@@ -189,7 +203,7 @@ data class JKClassTypeImpl(
     override val nullability: Nullability = Nullability.Default
 ) : JKClassType
 
-object JKNoTypeImpl: JKNoType {
+object JKNoTypeImpl : JKNoType {
     override val nullability: Nullability = Nullability.NotNull
 }
 
@@ -251,14 +265,16 @@ class JKLocalVariableImpl(
     type: JKTypeElement,
     name: JKNameIdentifier,
     initializer: JKExpression,
-    override var mutability: Mutability,
+    mutabilityElement: JKMutabilityModifierElement,
     annotationList: JKAnnotationList = JKAnnotationListImpl()
-) : JKLocalVariable, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKLocalVariable(), PsiOwner by PsiOwnerImpl() {
     override var initializer by child(initializer)
     override var name by child(name)
     override var type by child(type)
     override var annotationList by child(annotationList)
 
+    override var mutabilityElement by child(mutabilityElement)
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLocalVariable(this, data)
 }
 
@@ -266,19 +282,23 @@ class JKStubExpressionImpl : JKStubExpression, JKElementBase(), PsiOwner by PsiO
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitStubExpression(this, data)
 }
 
-object JKBodyStub : JKBlock, JKTreeElement {
+object JKBodyStubImpl : JKBodyStub() {
+    override var leftNonCodeElements: List<JKNonCodeElement> = emptyList()//TODO fix
+    override var rightNonCodeElements: List<JKNonCodeElement> = emptyList()//TODO fix
+
     override fun copy(): JKTreeElement = this
 
     override var statements: List<JKStatement>
         get() = emptyList()
         set(value) {}
 
-    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBlock(this, data)
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBodyStub(this, data)
 
     override fun <D> acceptChildren(visitor: JKVisitor<Unit, D>, data: D) {}
 
-    override val parent: JKElement?
+    override var parent: JKElement?
         get() = null
+        set (it) {}
 
     override fun detach(from: JKElement) {
     }
@@ -287,14 +307,15 @@ object JKBodyStub : JKBlock, JKTreeElement {
     }
 }
 
-class JKBlockStatementImpl(block: JKBlock) : JKBlockStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKBlockStatementImpl(block: JKBlock) : JKBlockStatement(), PsiOwner by PsiOwnerImpl() {
     override var block by child(block)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBlockStatement(this, data)
 }
 
-class JKBlockStatementWithoutBracketsImpl(block: JKBlock) : JKBlockStatementWithoutBrackets, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKBlockStatementWithoutBracketsImpl(statements: List<JKStatement>) : JKBlockStatementWithoutBrackets(),
-    override var block by child(block)
+    PsiOwner by PsiOwnerImpl() {
+    override var statements by children(statements)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBlockStatementWithoutBrackets(this, data)
 }
@@ -304,40 +325,44 @@ class JKThisExpressionImpl(qualifierLabel: JKLabel) : JKThisExpression, JKBranch
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitThisExpression(this, data)
 }
 
-class JKSuperExpressionImpl(qualifierLabel: JKLabel = JKLabelEmptyImpl()) : JKSuperExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKSuperExpressionImpl(qualifierLabel: JKLabel = JKLabelEmptyImpl()) : JKSuperExpression, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override var qualifierLabel: JKLabel by child(qualifierLabel)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitSuperExpression(this, data)
 }
 
-class JKWhileStatementImpl(condition: JKExpression, body: JKStatement) : JKWhileStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKWhileStatementImpl(condition: JKExpression, body: JKStatement) : JKWhileStatement(),
+    PsiOwner by PsiOwnerImpl() {
     override var condition by child(condition)
     override var body by child(body)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitWhileStatement(this, data)
 }
 
-class JKDoWhileStatementImpl(body: JKStatement, condition: JKExpression) : JKDoWhileStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKDoWhileStatementImpl(body: JKStatement, condition: JKExpression) : JKDoWhileStatement(),
+    PsiOwner by PsiOwnerImpl() {
     override var condition by child(condition)
     override var body by child(body)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitDoWhileStatement(this, data)
 }
 
-class JKBreakStatementImpl : JKBreakStatement, JKElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKBreakStatementImpl : JKBreakStatement(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBreakStatement(this, data)
 }
 
-class JKBreakWithLabelStatementImpl(override var label: JKNameIdentifier) : JKBreakWithLabelStatement, JKElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKBreakWithLabelStatementImpl(override var label: JKNameIdentifier) : JKBreakWithLabelStatement(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBreakWithLabelStatement(this, data)
 }
 
-class JKIfStatementImpl(condition: JKExpression, thenBranch: JKStatement) : JKIfStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKIfStatementImpl(condition: JKExpression, thenBranch: JKStatement) : JKIfStatement(),
+    PsiOwner by PsiOwnerImpl() {
     override var thenBranch by child(thenBranch)
     override var condition by child(condition)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitIfStatement(this, data)
 }
 
-class JKIfElseStatementImpl(condition: JKExpression, thenBranch: JKStatement, elseBranch: JKStatement) : JKIfElseStatement,
+class JKIfElseStatementImpl(condition: JKExpression, thenBranch: JKStatement, elseBranch: JKStatement) : JKIfElseStatement(), PsiOwner by PsiOwnerImpl() {
-    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var elseBranch by child(elseBranch)
     override var thenBranch by child(thenBranch)
     override var condition by child(condition)
@@ -354,7 +379,8 @@ class JKIfElseExpressionImpl(condition: JKExpression, thenBranch: JKExpression,
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitIfElseExpression(this, data)
 }
 
-class JKClassAccessExpressionImpl(override var identifier: JKClassSymbol) : JKClassAccessExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKClassAccessExpressionImpl(override var identifier: JKClassSymbol) : JKClassAccessExpression, JKElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitClassAccessExpression(this, data)
 }
 
@@ -392,9 +418,11 @@ class JKDelegationConstructorCallImpl(
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitDelegationConstructorCall(this, data)
 }
 
-class JKFieldAccessExpressionImpl(override var identifier: JKFieldSymbol) : JKFieldAccessExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKFieldAccessExpressionImpl(override var identifier: JKFieldSymbol) : JKFieldAccessExpression, JKElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitFieldAccessExpression(this, data)
 }
+
 val JKStatement.statements: List<JKStatement>
     get() =
         when (this) {
@@ -405,23 +433,25 @@ val JKStatement.statements: List<JKStatement>
 class JKLabelEmptyImpl : JKLabelEmpty, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLabelEmpty(this, data)
 }
+
 class JKLabelTextImpl(label: JKNameIdentifier) : JKLabelText, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override val label: JKNameIdentifier by child(label)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLabelText(this, data)
 }
 
-class JKContinueStatementImpl(label: JKLabel) : JKContinueStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKContinueStatementImpl(label: JKLabel) : JKContinueStatement(), PsiOwner by PsiOwnerImpl() {
     override var label: JKLabel by child(label)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitContinueStatement(this, data)
 }
 
-class JKLabeledStatementImpl(statement: JKStatement, labels: List<JKNameIdentifier>) : JKLabeledStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+class JKLabeledStatementImpl(statement: JKStatement, labels: List<JKNameIdentifier>) : JKLabeledStatement, JKBranchElementBase(),
+    PsiOwner by PsiOwnerImpl() {
     override var statement: JKStatement by child(statement)
     override val labels: List<JKNameIdentifier> by children(labels)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLabeledStatement(this, data)
 }
 
-class JKEmptyStatementImpl: JKEmptyStatement, JKElementBase() , PsiOwner by PsiOwnerImpl() {
+class JKEmptyStatementImpl : JKEmptyStatement(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitEmptyStatement(this, data)
 }
 
@@ -430,7 +460,7 @@ class PsiOwnerImpl(override var psi: PsiElement? = null) : PsiOwner
 val JKElement.psi: PsiElement?
     get() = (this as? PsiOwner)?.psi
 
-fun <Elem : PsiElement>JKElement.psi(): Elem? =
+fun <Elem : PsiElement> JKElement.psi(): Elem? =
     (this as? PsiOwner)?.psi as? Elem
 
 class JKTypeParameterListImpl(typeParameters: List<JKTypeParameter> = emptyList()) : JKTypeParameterList, JKBranchElementBase() {
@@ -460,7 +490,7 @@ class JKEnumConstantImpl(
     body: JKClassBody,
     type: JKTypeElement,
     annotationList: JKAnnotationList = JKAnnotationListImpl()
-) : JKEnumConstant, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKEnumConstant(), PsiOwner by PsiOwnerImpl() {
     override var name: JKNameIdentifier by child(name)
     override val arguments: JKArgumentList by child(arguments)
     override val body: JKClassBody by child(body)
@@ -475,7 +505,7 @@ fun JKTypeElement.present(): Boolean =
     type != JKNoTypeImpl
 
 class JKForInStatementImpl(declaration: JKDeclaration, iterationExpression: JKExpression, body: JKStatement) :
-    JKForInStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+    JKForInStatement(), PsiOwner by PsiOwnerImpl() {
     override var declaration: JKDeclaration by child(declaration)
     override var iterationExpression: JKExpression by child(iterationExpression)
     override var body: JKStatement by child(body)
@@ -485,12 +515,12 @@ class JKForInStatementImpl(declaration: JKDeclaration, iterationExpression: JKEx
 fun JKStatement.isEmpty(): Boolean =
     when (this) {
         is JKEmptyStatement -> true
-        is JKBlockStatement -> block is JKBodyStub
+        is JKBlockStatement -> block is JKBodyStubImpl
         is JKExpressionStatement -> expression is JKStubExpression
         else -> false
     }
 
-class JKPackageDeclarationImpl(packageName: JKNameIdentifier) : JKPackageDeclaration, JKBranchElementBase() {
+class JKPackageDeclarationImpl(packageName: JKNameIdentifier) : JKPackageDeclaration() {
     override var packageName: JKNameIdentifier by child(packageName)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitPackageDeclaration(this, data)
 }
@@ -506,6 +536,8 @@ class JKAnnotationImpl(
 ) : JKAnnotation, JKBranchElementBase() {
     override var arguments: List<JKAnnotationParameter> by children(arguments)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitAnnotation(this, data)
+
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
 }
 
 class JKTypeArgumentListImpl(typeArguments: List<JKTypeElement> = emptyList()) : JKTypeArgumentList, JKBranchElementBase(),
@@ -528,6 +560,8 @@ class JKImportStatementImpl(name: JKNameIdentifier) : JKImportStatement, JKBranc
     override val name: JKNameIdentifier by child(name)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitImportStatement(this, data)
+
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
 }
 
 class JKAnnotationParameterImpl(value: JKAnnotationMemberValue) : JKAnnotationParameter, JKBranchElementBase() {
@@ -569,3 +603,20 @@ class JKArgumentListImpl(arguments: List<JKArgument> = emptyList()) : JKArgument
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitArgumentList(this, data)
 }
+
+
+class JKMutabilityModifierElementImpl(override var mutability: Mutability) : JKMutabilityModifierElement() {
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitMutabilityModifierElement(this, data)
+}
+
+class JKModalityModifierElementImpl(override var modality: Modality) : JKModalityModifierElement() {
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitModalityModifierElement(this, data)
+}
+
+class JKVisibilityModifierElementImpl(override var visibility: Visibility) : JKVisibilityModifierElement() {
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitVisibilityModifierElement(this, data)
+}
+
+class JKExtraModifierElementImpl(override var extraModifier: ExtraModifier) : JKExtraModifierElement() {
+    override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitExtraModifierElement(this, data)
+}

nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/k.kt

@@ -30,11 +30,11 @@ class JKKtPropertyImpl(
     getter: JKKtGetterOrSetter,
     setter: JKKtGetterOrSetter,
     annotationList: JKAnnotationList,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality,
+    modalityElement: JKModalityModifierElement,
-    override var mutability: Mutability
+    mutabilityElement: JKMutabilityModifierElement
-) : JKBranchElementBase(), JKKtProperty, PsiOwner by PsiOwnerImpl() {
+) : JKKtProperty(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtProperty(this, data)
 
     override var annotationList: JKAnnotationList by child(annotationList)
@@ -43,6 +43,13 @@ class JKKtPropertyImpl(
     override var initializer: JKExpression by child(initializer)
     override var getter: JKKtGetterOrSetter by child(getter)
     override var setter: JKKtGetterOrSetter by child(setter)
+
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
+    override var mutabilityElement by child(mutabilityElement)
+
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
 }
 
 class JKKtFunctionImpl(
@@ -52,10 +59,10 @@ class JKKtFunctionImpl(
     block: JKBlock,
     typeParameterList: JKTypeParameterList,
     annotationList: JKAnnotationList,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality
+    modalityElement: JKModalityModifierElement
-) : JKBranchElementBase(), JKKtFunction, PsiOwner by PsiOwnerImpl() {
+) : JKKtFunction(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtFunction(this, data)
 
     override var returnType: JKTypeElement by child(returnType)
@@ -64,6 +71,11 @@ class JKKtFunctionImpl(
     override var block: JKBlock by child(block)
     override var typeParameterList: JKTypeParameterList by child(typeParameterList)
     override var annotationList: JKAnnotationList by child(annotationList)
+
+
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
 }
 
 sealed class JKKtQualifierImpl : JKQualifier, JKElementBase() {
@@ -147,11 +159,12 @@ class JKKtAssignmentStatementImpl(
     field: JKAssignableExpression,
     expression: JKExpression,
     override var operator: JKOperator
-) : JKKtAssignmentStatement, JKBranchElementBase() {
+) : JKKtAssignmentStatement() {
     override var field: JKAssignableExpression by child(field)
     override var expression by child(expression)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtAssignmentStatement(this, data)
 
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
 }
 
 object JKContextType : JKType {
@@ -165,10 +178,10 @@ class JKKtConstructorImpl(
     block: JKBlock,
     delegationCall: JKExpression,
     annotationList: JKAnnotationList,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality
+    modalityElement: JKModalityModifierElement
-) : JKBranchElementBase(), JKKtConstructor {
+) : JKKtConstructor() {
     override var returnType: JKTypeElement by child(JKTypeElementImpl(JKNoTypeImpl))
 
     override var name: JKNameIdentifier by child(name)
@@ -178,6 +191,10 @@ class JKKtConstructorImpl(
     override var typeParameterList: JKTypeParameterList by child(JKTypeParameterListImpl())
     override var annotationList: JKAnnotationList by child(annotationList)
 
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtConstructor(this, data)
 }
 
@@ -186,26 +203,30 @@ class JKKtPrimaryConstructorImpl(
     parameters: List<JKParameter>,
     delegationCall: JKExpression,
     annotationList: JKAnnotationList,
-    override var extraModifiers: List<ExtraModifier>,
+    extraModifierElements: List<JKExtraModifierElement>,
-    override var visibility: Visibility,
+    visibilityElement: JKVisibilityModifierElement,
-    override var modality: Modality
+    modalityElement: JKModalityModifierElement
-) : JKBranchElementBase(), JKKtPrimaryConstructor {
+) : JKKtPrimaryConstructor() {
     override var returnType: JKTypeElement by child(JKTypeElementImpl(JKNoTypeImpl))
 
     override var name: JKNameIdentifier by child(name)
     override var parameters: List<JKParameter> by children(parameters)
-    override var block: JKBlock by child(JKBodyStub)
+    override var block: JKBlock by child(JKBodyStubImpl)
     override var delegationCall: JKExpression by child(delegationCall)
     override var typeParameterList: JKTypeParameterList by child(JKTypeParameterListImpl())
     override var annotationList: JKAnnotationList by child(annotationList)
 
+    override var extraModifierElements by children(extraModifierElements)
+    override var visibilityElement by child(visibilityElement)
+    override var modalityElement by child(modalityElement)
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtPrimaryConstructor(this, data)
 }
 
 class JKKtWhenStatementImpl(
     expression: JKExpression,
     cases: List<JKKtWhenCase>
-) : JKKtWhenStatement, JKBranchElementBase() {
+) : JKKtWhenStatement() {
     override var expression: JKExpression by child(expression)
     override var cases: List<JKKtWhenCase> by children(cases)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtWhenStatement(this, data)
@@ -235,7 +256,7 @@ class JKKtIsExpressionImpl(expression: JKExpression, type: JKTypeElement) : JKKt
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtIsExpression(this, data)
 }
 
-class JKKtInitDeclarationImpl(block: JKBlock) : JKKtInitDeclaration, JKBranchElementBase() {
+class JKKtInitDeclarationImpl(block: JKBlock) : JKKtInitDeclaration() {
     override var block: JKBlock by child(block)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtInitDeclaration(this, data)
 }
@@ -244,7 +265,7 @@ class JKKtInitDeclarationImpl(block: JKBlock) : JKKtInitDeclaration, JKBranchEle
 class JKKtConvertedFromForLoopSyntheticWhileStatementImpl(
     variableDeclaration: JKStatement,
     whileStatement: JKWhileStatement
-) : JKKtConvertedFromForLoopSyntheticWhileStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
+) : JKKtConvertedFromForLoopSyntheticWhileStatement(), PsiOwner by PsiOwnerImpl() {
     override var variableDeclaration: JKStatement by child(variableDeclaration)
     override var whileStatement: JKWhileStatement by child(whileStatement)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R =
@@ -279,17 +300,21 @@ class JKKtTryCatchSectionImpl(
 class JKKtGetterOrSetterImpl(
     body: JKStatement,
     override val kind: JKKtGetterOrSetter.Kind,
-    override var visibility: Visibility
+    visibilityElement: JKVisibilityModifierElement
 ) : JKKtGetterOrSetter, JKBranchElementBase() {
     override var body: JKStatement by child(body)
+
+    override var visibilityElement by child(visibilityElement)
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtGetterOrSetter(this, data)
 }
 
 class JKKtEmptyGetterOrSetterImpl : JKKtEmptyGetterOrSetter, JKBranchElementBase() {
-    override var visibility: Visibility = Visibility.PUBLIC
     override var body: JKStatement by child(JKEmptyStatementImpl())
     override val kind: JKKtGetterOrSetter.Kind
         get() = error("Cannot get kind of JKKtEmptyGetterOrSetter")
 
+    override var visibilityElement by child(JKVisibilityModifierElementImpl(Visibility.PUBLIC))
+
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitKtEmptyGetterOrSetter(this, data)
 }

nj2k/src/org/jetbrains/kotlin/nj2k/tree/impl/tokens.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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.impl
+
+import org.jetbrains.kotlin.nj2k.tree.JKCommentElement
+import org.jetbrains.kotlin.nj2k.tree.JKNonCodeElement
+import org.jetbrains.kotlin.nj2k.tree.JKSpaceElement
+import org.jetbrains.kotlin.nj2k.tree.JKTokenElement
+
+
+class JKSpaceElementImpl(override val text: String) : JKSpaceElement
+
+class JKCommentElementImpl(override val text: String) : JKCommentElement
+
+class JKTokenElementImpl(override val text: String) : JKTokenElement {
+    override var leftNonCodeElements: List<JKNonCodeElement> = emptyList()
+    override var rightNonCodeElements: List<JKNonCodeElement> = emptyList()
+}

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

@@ -16,23 +16,28 @@
 
 package org.jetbrains.kotlin.nj2k.tree
 
-import org.jetbrains.kotlin.nj2k.tree.impl.JKClassSymbol
+import org.jetbrains.kotlin.nj2k.tree.impl.*
 
-interface JKField : JKVariable, JKVisibilityOwner, JKMutabilityOwner, JKModalityOwner, JKExtraModifiersOwner, JKAnnotationListOwner
+abstract class JKField : JKVariable(), JKVisibilityOwner, JKMutabilityOwner, JKModalityOwner, JKExtraModifiersOwner, JKAnnotationListOwner
 
-interface JKJavaField : JKField, JKBranchElement
+abstract class JKJavaField : JKField()
 
-interface JKJavaMethod : JKMethod, JKBranchElement {
+abstract class JKJavaMethod : JKMethod(), JKBranchElement {
-    var throwsList: List<JKTypeElement>
+    abstract var throwsList: List<JKTypeElement>
 }
 
 interface JKJavaMethodCallExpression : JKMethodCallExpression
 
-interface JKClassBody : JKTreeElement, JKBranchElement {
+abstract class JKClassBody : JKTreeElement, JKBranchElementBase() {
-    var declarations: List<JKDeclaration>
+    abstract var declarations: List<JKDeclaration>
+
+    val leftBrace = JKTokenElementImpl("{")
+    val rightBrace = JKTokenElementImpl("}")
 }
 
-interface JKEmptyClassBody : JKClassBody
+abstract class JKEmptyClassBody : JKClassBody() {
+    override var declarations: List<JKDeclaration> = emptyList()
+}
 
 interface JKJavaNewExpression : JKExpression, JKTypeArgumentListOwner, PsiOwner {
     val classSymbol: JKClassSymbol
@@ -61,20 +66,20 @@ interface JKJavaNewArray : JKExpression {
 
 interface JKJavaLiteralExpression : JKLiteralExpression
 
-interface JKReturnStatement : JKStatement {
+abstract class JKReturnStatement : JKStatement() {
-    val expression: JKExpression
+    abstract val expression: JKExpression
-    var label: JKLabel
+    abstract var label: JKLabel
 }
 
-interface JKJavaAssertStatement : JKStatement {
+abstract class JKJavaAssertStatement : JKStatement() {
-    val condition: JKExpression
+    abstract val condition: JKExpression
-    val description: JKExpression
+    abstract val description: JKExpression
 }
 
-interface JKJavaForLoopStatement : JKLoopStatement {
+abstract class JKJavaForLoopStatement : JKLoopStatement() {
-    var initializer: JKStatement
+    abstract var initializer: JKStatement
-    var condition: JKExpression
+    abstract var condition: JKExpression
-    var updaters: List<JKStatement>
+    abstract var updaters: List<JKStatement>
 }
 
 
@@ -91,15 +96,15 @@ interface JKJavaAssignmentExpression : JKExpression, JKBranchElement {
     var operator: JKOperator
 }
 
-interface JKJavaThrowStatement : JKStatement {
+abstract class JKJavaThrowStatement : JKStatement() {
-    var exception: JKExpression
+    abstract var exception: JKExpression
 }
 
-interface JKJavaTryStatement : JKStatement {
+abstract class JKJavaTryStatement : JKStatement() {
-    var resourceDeclarations: List<JKDeclaration>
+    abstract var resourceDeclarations: List<JKDeclaration>
-    var tryBlock: JKBlock
+    abstract var tryBlock: JKBlock
-    var finallyBlock: JKBlock
+    abstract var finallyBlock: JKBlock
-    var catchSections: List<JKJavaTryCatchSection>
+    abstract var catchSections: List<JKJavaTryCatchSection>
 }
 
 interface JKJavaTryCatchSection : JKTreeElement {
@@ -108,9 +113,9 @@ interface JKJavaTryCatchSection : JKTreeElement {
 }
 
 
-interface JKJavaSwitchStatement : JKStatement {
+abstract class JKJavaSwitchStatement : JKStatement() {
-    var expression: JKExpression
+    abstract var expression: JKExpression
-    var cases: List<JKJavaSwitchCase>
+    abstract var cases: List<JKJavaSwitchCase>
 }
 
 interface JKJavaSwitchCase : JKTreeElement {
@@ -127,17 +132,21 @@ interface JKJavaLabelSwitchCase : JKJavaSwitchCase {
     var label: JKExpression
 }
 
-interface JKJavaContinueStatement: JKStatement
+abstract class JKJavaContinueStatement : JKStatement()
 
-interface JKJavaSynchronizedStatement : JKStatement, JKBranchElement {
+abstract class JKJavaSynchronizedStatement : JKStatement(), JKBranchElement {
-    val lockExpression: JKExpression
+    abstract val lockExpression: JKExpression
-    val body: JKBlock
+    abstract val body: JKBlock
 }
 
-interface JKJavaAnnotationMethod : JKMethod, JKBranchElement {
+abstract class JKJavaAnnotationMethod : JKMethod(), JKBranchElement {
-    val defaultValue: JKAnnotationMemberValue
+    abstract val defaultValue: JKAnnotationMemberValue
+
+    override var extraModifierElements by children<JKExtraModifierElement>()
+    override var visibilityElement by child(JKVisibilityModifierElementImpl(Visibility.PUBLIC))
+    override var modalityElement by child(JKModalityModifierElementImpl(Modality.FINAL))
 }
 
-interface JKJavaStaticInitDeclaration : JKDeclaration, JKBranchElement {
+abstract class JKJavaStaticInitDeclaration : JKDeclaration(), JKBranchElement {
-    var block: JKBlock
+    abstract var block: JKBlock
 }

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

@@ -17,12 +17,10 @@
 package org.jetbrains.kotlin.nj2k.tree
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.nj2k.tree.impl.JKClassSymbol
+import org.jetbrains.kotlin.nj2k.tree.impl.*
-import org.jetbrains.kotlin.nj2k.tree.impl.JKFieldSymbol
-import org.jetbrains.kotlin.nj2k.tree.impl.JKMethodSymbol
 import org.jetbrains.kotlin.nj2k.tree.visitors.JKVisitor
 
-interface JKTreeElement : JKElement {
+interface JKTreeElement : JKElement, JKNonCodeElementsListOwner {
     fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R
 
     fun <R> accept(visitor: JKVisitor<R, Nothing?>): R = accept(visitor, null)
@@ -38,7 +36,9 @@ interface PsiOwner {
     var psi: PsiElement?
 }
 
-interface JKDeclaration : JKTreeElement
+abstract class JKDeclaration : JKTreeElement, JKBranchElementBase() {
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
+}
 
 interface JKImportStatement : JKTreeElement {
     val name: JKNameIdentifier
@@ -50,14 +50,15 @@ interface JKFile : JKTreeElement, JKBranchElement {
     var declarationList: List<JKDeclaration>
 }
 
-interface JKClass : JKDeclaration, JKVisibilityOwner, JKExtraModifiersOwner, JKModalityOwner, JKTypeParameterListOwner, JKAnnotationListOwner,
+abstract class JKClass : JKDeclaration(), JKVisibilityOwner, JKExtraModifiersOwner, JKModalityOwner, JKTypeParameterListOwner,
+    JKAnnotationListOwner,
     JKBranchElement {
-    val name: JKNameIdentifier
+    abstract val name: JKNameIdentifier
 
-    val inheritance: JKInheritanceInfo
+    abstract val inheritance: JKInheritanceInfo
 
-    var classBody: JKClassBody
+    abstract var classBody: JKClassBody
-    var classKind: ClassKind
+    abstract var classKind: ClassKind
 
     enum class ClassKind {
         ANNOTATION, CLASS, ENUM, INTERFACE, OBJECT, COMPANION
@@ -100,32 +101,71 @@ interface JKAnnotationListOwner : JKTreeElement {
     var annotationList: JKAnnotationList
 }
 
-interface JKMethod : JKDeclaration, JKVisibilityOwner, JKModalityOwner, JKExtraModifiersOwner, JKTypeParameterListOwner,
+abstract class JKMethod : JKDeclaration(), JKVisibilityOwner, JKModalityOwner, JKExtraModifiersOwner, JKTypeParameterListOwner,
     JKAnnotationListOwner {
-    val name: JKNameIdentifier
+    abstract val name: JKNameIdentifier
-    var parameters: List<JKParameter>
+    abstract var parameters: List<JKParameter>
-    var returnType: JKTypeElement
+    abstract var returnType: JKTypeElement
-    var block: JKBlock
+    abstract var block: JKBlock
+
+    val leftParen = JKTokenElementImpl("(")
+    val rightParen = JKTokenElementImpl(")")
 }
 
-interface JKVariable : JKDeclaration, JKAnnotationListOwner {
+abstract class JKVariable : JKDeclaration(), JKAnnotationListOwner {
-    var type: JKTypeElement
+    abstract var type: JKTypeElement
-    var name: JKNameIdentifier
+    abstract var name: JKNameIdentifier
-    var initializer: JKExpression
+    abstract var initializer: JKExpression
 }
 
-interface JKForLoopVariable : JKVariable
+abstract class JKForLoopVariable : JKVariable()
 
-interface JKLocalVariable : JKVariable, JKMutabilityOwner
+abstract class JKLocalVariable : JKVariable(), JKMutabilityOwner
 
-interface JKExtraModifiersOwner : JKModifiersListOwner {
+
-    var extraModifiers: List<ExtraModifier>
+abstract class JKModifierElement : JKTreeElement, JKBranchElementBase(), JKNonCodeElementsListOwner
+
+val JKModifierElement.modifier: Modifier
+    get() = when (this) {
+        is JKMutabilityModifierElement -> mutability
+        is JKModalityModifierElement -> modality
+        is JKVisibilityModifierElement -> visibility
+        is JKExtraModifierElement -> extraModifier
+        else -> error("")
+    }
+
+abstract class JKMutabilityModifierElement : JKModifierElement() {
+    abstract var mutability: Mutability
+}
+
+abstract class JKModalityModifierElement : JKModifierElement() {
+    abstract var modality: Modality
+}
+
+abstract class JKVisibilityModifierElement : JKModifierElement() {
+    abstract var visibility: Visibility
+}
+
+abstract class JKExtraModifierElement : JKModifierElement() {
+    abstract var extraModifier: ExtraModifier
 }
 
 interface Modifier {
     val text: String
 }
 
+interface JKExtraModifiersOwner : JKModifiersListOwner {
+    var extraModifierElements: List<JKExtraModifierElement>
+}
+
+fun JKExtraModifiersOwner.elementByModifier(modifier: ExtraModifier): JKExtraModifierElement? =
+    extraModifierElements.firstOrNull { it.extraModifier == modifier }
+
+fun JKExtraModifiersOwner.hasExtraModifier(modifier: ExtraModifier): Boolean =
+    extraModifierElements.any { it.extraModifier == modifier }
+
+//val JKExtraModifiersOwner.extraModifiers
+
 enum class ExtraModifier(override val text: String) : Modifier {
     ACTUAL("actual"),
     ANNOTATION("annotation"),
@@ -156,8 +196,8 @@ enum class ExtraModifier(override val text: String) : Modifier {
     VOLATILE("volatile")
 }
 
-interface JKVisibilityOwner : JKModifiersListOwner{
+interface JKVisibilityOwner : JKModifiersListOwner {
-    var visibility: Visibility
+    val visibilityElement: JKVisibilityModifierElement
 }
 
 enum class Visibility(override val text: String) : Modifier {
@@ -168,8 +208,14 @@ enum class Visibility(override val text: String) : Modifier {
     PRIVATE("private")
 }
 
+var JKVisibilityOwner.visibility: Visibility
+    get() = visibilityElement.visibility
+    set(value) {
+        visibilityElement.visibility = value
+    }
+
 interface JKModalityOwner : JKModifiersListOwner {
-    var modality: Modality
+    val modalityElement: JKModalityModifierElement
 }
 
 
@@ -180,8 +226,14 @@ enum class Modality(override val text: String) : Modifier {
     OVERRIDE("override")
 }
 
+var JKModalityOwner.modality: Modality
+    get() = modalityElement.modality
+    set(value) {
+        modalityElement.modality = value
+    }
+
 interface JKMutabilityOwner : JKModifiersListOwner {
-    var mutability: Mutability
+    val mutabilityElement: JKMutabilityModifierElement
 }
 
 enum class Mutability(override val text: String) : Modifier {
@@ -190,23 +242,38 @@ enum class Mutability(override val text: String) : Modifier {
     UNKNOWN("var")//TODO ???
 }
 
+var JKMutabilityOwner.mutability: Mutability
+    get() = mutabilityElement.mutability
+    set(value) {
+        mutabilityElement.mutability = value
+    }
+
 interface JKModifiersListOwner : JKTreeElement
 
-fun JKModifiersListOwner.modifiers(): List<Modifier> =
+fun JKModifiersListOwner.modifierElements(): List<JKModifierElement> =
-    listOfNotNull((this as? JKVisibilityOwner)?.visibility) +
+    listOfNotNull((this as? JKVisibilityOwner)?.visibilityElement) +
-            (this as? JKExtraModifiersOwner)?.extraModifiers.orEmpty() +
+            (this as? JKExtraModifiersOwner)?.extraModifierElements.orEmpty() +
-            listOfNotNull((this as? JKModalityOwner)?.modality) +
+            listOfNotNull((this as? JKModalityOwner)?.modalityElement) +
-            listOfNotNull((this as? JKMutabilityOwner)?.mutability)
+            listOfNotNull((this as? JKMutabilityOwner)?.mutabilityElement)
 
 
 interface JKTypeElement : JKTreeElement {
     val type: JKType
 }
 
-interface JKStatement : JKTreeElement
+abstract class JKStatement : JKTreeElement, JKBranchElementBase() {
+    override var rightNonCodeElements: List<JKNonCodeElement> = listOf(JKSpaceElementImpl("\n"))
+}
+
+abstract class JKBlock : JKTreeElement, JKBranchElementBase() {
+    abstract var statements: List<JKStatement>
+
+    val leftBrace = JKTokenElementImpl("{")
+    val rightBrace = JKTokenElementImpl("}")
+}
+
+abstract class JKBodyStub : JKBlock() {
 
-interface JKBlock : JKTreeElement {
-    var statements: List<JKStatement>
 }
 
 interface JKIdentifier : JKTreeElement
@@ -217,12 +284,12 @@ interface JKNameIdentifier : JKIdentifier {
 
 interface JKExpression : JKTreeElement, JKAnnotationMemberValue
 
-interface JKExpressionStatement : JKStatement, JKBranchElement {
+abstract class JKExpressionStatement : JKStatement() {
-    val expression: JKExpression
+    abstract val expression: JKExpression
 }
 
-interface JKDeclarationStatement : JKStatement {
+abstract class JKDeclarationStatement : JKStatement() {
-    val declaredStatements: List<JKDeclaration>
+    abstract val declaredStatements: List<JKDeclaration>
 }
 
 interface JKOperatorExpression : JKExpression {
@@ -310,8 +377,10 @@ interface JKLiteralExpression : JKExpression {
     }
 }
 
-interface JKParameter : JKVariable, JKModifiersListOwner {
+abstract class JKParameter : JKVariable(), JKModifiersListOwner {
-    var isVarArgs: Boolean
+    abstract var isVarArgs: Boolean
+
+    override var rightNonCodeElements: List<JKNonCodeElement> = emptyList()
 }
 
 interface JKStringLiteralExpression : JKLiteralExpression {
@@ -320,16 +389,16 @@ interface JKStringLiteralExpression : JKLiteralExpression {
 
 interface JKStubExpression : JKExpression
 
-interface JKLoopStatement : JKStatement {
+abstract class JKLoopStatement : JKStatement() {
-    var body: JKStatement
+    abstract var body: JKStatement
 }
 
-interface JKBlockStatement : JKStatement, JKBranchElement {
+abstract class JKBlockStatement : JKStatement() {
-    var block: JKBlock
+    abstract var block: JKBlock
 }
 
-interface JKBlockStatementWithoutBrackets : JKStatement, JKBranchElement {
+abstract class JKBlockStatementWithoutBrackets : JKStatement() {
-    var block: JKBlock
+    abstract var statements: List<JKStatement>
 }
 
 interface JKThisExpression : JKExpression {
@@ -340,27 +409,27 @@ interface JKSuperExpression : JKExpression {
     var qualifierLabel: JKLabel
 }
 
-interface JKWhileStatement : JKLoopStatement {
+abstract class JKWhileStatement : JKLoopStatement() {
-    var condition: JKExpression
+    abstract var condition: JKExpression
 }
 
-interface JKDoWhileStatement : JKLoopStatement {
+abstract class JKDoWhileStatement : JKLoopStatement() {
-    var condition: JKExpression
+    abstract var condition: JKExpression
 }
 
-interface JKBreakStatement : JKStatement
+abstract class JKBreakStatement : JKStatement()
 
-interface JKBreakWithLabelStatement : JKBreakStatement {
+abstract class JKBreakWithLabelStatement : JKBreakStatement() {
-    var label: JKNameIdentifier
+    abstract var label: JKNameIdentifier
 }
 
-interface JKIfStatement : JKStatement {
+abstract class JKIfStatement : JKStatement() {
-    var condition: JKExpression
+    abstract var condition: JKExpression
-    var thenBranch: JKStatement
+    abstract var thenBranch: JKStatement
 }
 
-interface JKIfElseStatement : JKIfStatement {
+abstract class JKIfElseStatement : JKIfStatement() {
-    var elseBranch: JKStatement
+    abstract var elseBranch: JKStatement
 }
 
 interface JKIfElseExpression : JKExpression {
@@ -391,16 +460,16 @@ interface JKLabelText : JKLabel {
     val label: JKNameIdentifier
 }
 
-interface JKContinueStatement : JKStatement {
+abstract class JKContinueStatement : JKStatement() {
-    var label: JKLabel
+    abstract var label: JKLabel
 }
 
-interface JKLabeledStatement : JKStatement {
+interface JKLabeledStatement : JKExpression {
     var statement: JKStatement
     val labels: List<JKNameIdentifier>
 }
 
-interface JKEmptyStatement : JKStatement
+abstract class JKEmptyStatement : JKStatement()
 
 interface JKTypeParameterList : JKTreeElement {
     var typeParameters: List<JKTypeParameter>
@@ -415,19 +484,19 @@ interface JKTypeParameterListOwner : JKTreeElement {
     var typeParameterList: JKTypeParameterList
 }
 
-interface JKEnumConstant : JKVariable {
+abstract class JKEnumConstant : JKVariable() {
-    val arguments: JKArgumentList
+    abstract val arguments: JKArgumentList
-    val body: JKClassBody
+    abstract val body: JKClassBody
 }
 
-interface JKForInStatement : JKStatement {
+abstract class JKForInStatement : JKStatement() {
-    var declaration: JKDeclaration
+    abstract var declaration: JKDeclaration
-    var iterationExpression: JKExpression
+    abstract var iterationExpression: JKExpression
-    var body: JKStatement
+    abstract var body: JKStatement
 }
 
-interface JKPackageDeclaration : JKDeclaration {
+abstract class JKPackageDeclaration : JKDeclaration() {
-    var packageName: JKNameIdentifier
+    abstract var packageName: JKNameIdentifier
 }
 
 interface JKClassLiteralExpression : JKExpression {

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

@@ -16,8 +16,6 @@
 
 package org.jetbrains.kotlin.nj2k.tree
 
-import org.jetbrains.kotlin.nj2k.tree.impl.JKMethodSymbol
-
 interface JKKtGetterOrSetter : JKTreeElement, JKVisibilityOwner, JKBranchElement {
     var body: JKStatement
     val kind: Kind
@@ -29,27 +27,24 @@ interface JKKtGetterOrSetter : JKTreeElement, JKVisibilityOwner, JKBranchElement
 
 interface JKKtEmptyGetterOrSetter : JKKtGetterOrSetter
 
-interface JKKtProperty : JKField, PsiOwner {
+abstract class JKKtProperty : JKField(), PsiOwner {
-    var getter: JKKtGetterOrSetter
+    abstract var getter: JKKtGetterOrSetter
-    var setter: JKKtGetterOrSetter
+    abstract var setter: JKKtGetterOrSetter
 }
 
 
-interface JKKtFunction : JKMethod, PsiOwner {
+abstract class JKKtFunction : JKMethod(), PsiOwner
+
+abstract class JKKtConstructor : JKMethod(), JKExtraModifiersOwner {
+    abstract var delegationCall: JKExpression
 }
 
-interface JKKtConstructor : JKDeclaration, JKExtraModifiersOwner, JKMethod, JKBranchElement {
+abstract class JKKtPrimaryConstructor : JKKtConstructor()
-    override var name: JKNameIdentifier
-    override var parameters: List<JKParameter>
-    var delegationCall: JKExpression
-}
 
-interface JKKtPrimaryConstructor : JKKtConstructor
+abstract class JKKtAssignmentStatement : JKStatement() {
-
+    abstract var field: JKAssignableExpression
-interface JKKtAssignmentStatement : JKStatement {
+    abstract var expression: JKExpression
-    var field: JKAssignableExpression
+    abstract var operator: JKOperator
-    var expression: JKExpression
-    var operator: JKOperator
 }
 
 interface JKKtCall : JKMethodCallExpression
@@ -62,9 +57,9 @@ interface JKKtAlsoCallExpression : JKKtMethodCallExpression {
 
 interface JKKtLiteralExpression : JKLiteralExpression
 
-interface JKKtWhenStatement : JKStatement {
+abstract class JKKtWhenStatement : JKStatement() {
-    var expression: JKExpression
+    abstract var expression: JKExpression
-    var cases: List<JKKtWhenCase>
+    abstract var cases: List<JKKtWhenCase>
 }
 
 interface JKKtWhenCase : JKTreeElement {
@@ -85,14 +80,14 @@ interface JKKtIsExpression : JKExpression, PsiOwner {
     var type: JKTypeElement
 }
 
-interface JKKtInitDeclaration : JKDeclaration {
+abstract class JKKtInitDeclaration : JKDeclaration() {
-    var block: JKBlock
+    abstract var block: JKBlock
 }
 
 
-interface JKKtConvertedFromForLoopSyntheticWhileStatement : JKStatement {
+abstract class JKKtConvertedFromForLoopSyntheticWhileStatement : JKStatement() {
-    var variableDeclaration: JKStatement
+    abstract var variableDeclaration: JKStatement
-    var whileStatement: JKWhileStatement
+    abstract var whileStatement: JKWhileStatement
 }
 
 

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. 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
+
+
+interface JKNonCodeElement {
+    val text: String
+}
+interface JKSpaceElement : JKNonCodeElement
+
+interface JKCommentElement : JKNonCodeElement
+
+fun JKCommentElement.isMultiline() =
+    text.startsWith("/*")
+
+fun JKCommentElement.isSingleline() =
+    text.startsWith("//")
+
+
+interface JKNonCodeElementsListOwner {
+    var leftNonCodeElements: List<JKNonCodeElement>
+    var rightNonCodeElements: List<JKNonCodeElement>
+}
+
+inline fun JKNonCodeElementsListOwner.takeNonCodeElementsFrom(other: JKNonCodeElementsListOwner) {
+    leftNonCodeElements = other.leftNonCodeElements
+    rightNonCodeElements = other.rightNonCodeElements
+}
+
+fun JKTreeElement.commentsFromInside(): List<JKCommentElement> {
+    val comments = mutableListOf<JKCommentElement>()
+    fun recurse(element: JKTreeElement): JKTreeElement {
+        comments += (element.leftNonCodeElements + element.rightNonCodeElements).filterIsInstance()
+        return applyRecursive(element, ::recurse)
+    }
+    applyRecursive(this, ::recurse)
+    return comments
+}
+
+inline fun <reified T : JKNonCodeElementsListOwner> T.withNonCodeElementsFrom(other: JKNonCodeElementsListOwner): T =
+    also { it.takeNonCodeElementsFrom(other) }
+
+
+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/visitors/GenerateJKTreeVisitorAndTransformer.kts

@@ -49,7 +49,7 @@
 import java.io.File
 
 
-val JK_ROOT = File("./j2k/newSrc/org/jetbrains/kotlin/j2k/tree")
+val JK_ROOT = File("./nj2k/src/org/jetbrains/kotlin/nj2k/tree")
 
 val JK_OUT_ROOT = File(JK_ROOT, "visitors")
 
@@ -58,17 +58,17 @@ val JK_JAVA_FILE = File(JK_ROOT, "j.kt")
 val JK_COMMON_FILE = File(JK_ROOT, "jk.kt")
 
 
-val interfaceRegex = "interface (JK[a-zA-Z]+)\\s?:?\\s?(JK[a-zA-Z]+)?".toRegex()
+val interfaceRegex = "(interface|abstract class)\\s+(JK[a-zA-Z]+)\\s+?:?\\s+(JK[a-zA-Z]+)?".toRegex()
 
 
 data class InterfaceData(val name: String, val extends: String?)
 
 fun File.interfaceNames() =
-        sequenceOf(this)
+    sequenceOf(this)
-                .map { it.readText() }
+        .map { it.readText() }
-                .flatMap { interfaceRegex.findAll(it) }
+        .flatMap { interfaceRegex.findAll(it) }
-                .map { match -> InterfaceData(match.groupValues[1], match.groupValues.getOrNull(2)) }
+        .map { match -> InterfaceData(match.groupValues[2], match.groupValues.getOrNull(3)) }
-                .toList()
+        .toList()
 
 
 fun String.safeVarName() = when (this) {
@@ -129,6 +129,41 @@ fun genVisitors(commonData: List<InterfaceData>, uncommonData: List<InterfaceDat
         appendln("}")
     })
 
+    File(JK_OUT_ROOT, "${visitorName}WithCommentsPrinting.kt").writeText(buildString {
+        appendln(pkg)
+        appendln()
+        appendln("import org.jetbrains.kotlin.nj2k.tree.*")
+        appendln()
+
+        appendln("interface ${visitorName}WithCommentsPrinting : ${visitorName}Void {")
+        appendln(
+            """
+        |    fun printLeftNonCodeElements(element: JKNonCodeElementsListOwner)
+        |    fun printRightNonCodeElements(element: JKNonCodeElementsListOwner)
+        |
+        """.trimMargin()
+        )
+
+        interfaceData.joinTo(this, separator = "\n\n") { (name, ext) ->
+            val nameWithoutPrefix = name.removePrefix("JK")
+            val argName = nameWithoutPrefix.decapitalize().safeVarName()
+            val arg = "$argName: $name"
+            val rawVisitSuffix = "Raw"
+            val generifyCall = if (name != "JKTreeElement") "= visit${ext!!.removePrefix("JK")}$rawVisitSuffix($argName)" else ""
+            """
+            |    override fun visit$nameWithoutPrefix($arg) {
+            |        printLeftNonCodeElements($argName)
+            |        visit$nameWithoutPrefix$rawVisitSuffix($argName)
+            |        printRightNonCodeElements($argName)
+            |    }
+            |
+            |    fun visit$nameWithoutPrefix$rawVisitSuffix($arg) $generifyCall
+            """.trimMargin()
+        }
+        appendln()
+        appendln("}")
+    })
+
     /*File(JK_OUT_ROOT, "$transformerName.kt").writeText(buildString {
         appendln(pkg)
         appendln()

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

@@ -18,6 +18,11 @@ interface JKVisitor<out R, in D> {
     fun visitVariable(variable: JKVariable, data: D): R = visitDeclaration(variable, data)
     fun visitForLoopVariable(forLoopVariable: JKForLoopVariable, data: D): R = visitVariable(forLoopVariable, data)
     fun visitLocalVariable(localVariable: JKLocalVariable, data: D): R = visitVariable(localVariable, data)
+    fun visitModifierElement(modifierElement: JKModifierElement, data: D): R = visitTreeElement(modifierElement, data)
+    fun visitMutabilityModifierElement(mutabilityModifierElement: JKMutabilityModifierElement, data: D): R = visitModifierElement(mutabilityModifierElement, data)
+    fun visitModalityModifierElement(modalityModifierElement: JKModalityModifierElement, data: D): R = visitModifierElement(modalityModifierElement, data)
+    fun visitVisibilityModifierElement(visibilityModifierElement: JKVisibilityModifierElement, data: D): R = visitModifierElement(visibilityModifierElement, data)
+    fun visitExtraModifierElement(extraModifierElement: JKExtraModifierElement, data: D): R = visitModifierElement(extraModifierElement, data)
     fun visitExtraModifiersOwner(extraModifiersOwner: JKExtraModifiersOwner, data: D): R = visitModifiersListOwner(extraModifiersOwner, data)
     fun visitVisibilityOwner(visibilityOwner: JKVisibilityOwner, data: D): R = visitModifiersListOwner(visibilityOwner, data)
     fun visitModalityOwner(modalityOwner: JKModalityOwner, data: D): R = visitModifiersListOwner(modalityOwner, data)
@@ -26,6 +31,7 @@ interface JKVisitor<out R, in D> {
     fun visitTypeElement(typeElement: JKTypeElement, data: D): R = visitTreeElement(typeElement, data)
     fun visitStatement(statement: JKStatement, data: D): R = visitTreeElement(statement, data)
     fun visitBlock(block: JKBlock, data: D): R = visitTreeElement(block, data)
+    fun visitBodyStub(bodyStub: JKBodyStub, data: D): R = visitBlock(bodyStub, data)
     fun visitIdentifier(identifier: JKIdentifier, data: D): R = visitTreeElement(identifier, data)
     fun visitNameIdentifier(nameIdentifier: JKNameIdentifier, data: D): R = visitIdentifier(nameIdentifier, data)
     fun visitExpression(expression: JKExpression, data: D): R = visitTreeElement(expression, data)
@@ -113,7 +119,7 @@ interface JKVisitor<out R, in D> {
     fun visitKtEmptyGetterOrSetter(ktEmptyGetterOrSetter: JKKtEmptyGetterOrSetter, data: D): R = visitKtGetterOrSetter(ktEmptyGetterOrSetter, data)
     fun visitKtProperty(ktProperty: JKKtProperty, data: D): R = visitField(ktProperty, data)
     fun visitKtFunction(ktFunction: JKKtFunction, data: D): R = visitMethod(ktFunction, data)
-    fun visitKtConstructor(ktConstructor: JKKtConstructor, data: D): R = visitDeclaration(ktConstructor, data)
+    fun visitKtConstructor(ktConstructor: JKKtConstructor, data: D): R = visitMethod(ktConstructor, data)
     fun visitKtPrimaryConstructor(ktPrimaryConstructor: JKKtPrimaryConstructor, data: D): R = visitKtConstructor(ktPrimaryConstructor, data)
     fun visitKtAssignmentStatement(ktAssignmentStatement: JKKtAssignmentStatement, data: D): R = visitStatement(ktAssignmentStatement, data)
     fun visitKtCall(ktCall: JKKtCall, data: D): R = visitMethodCallExpression(ktCall, data)

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

@@ -33,6 +33,16 @@ interface JKVisitorVoid : JKVisitor<Unit, Nothing?> {
     override fun visitForLoopVariable(forLoopVariable: JKForLoopVariable, data: Nothing?) = visitForLoopVariable(forLoopVariable)
     fun visitLocalVariable(localVariable: JKLocalVariable) = visitVariable(localVariable, null)
     override fun visitLocalVariable(localVariable: JKLocalVariable, data: Nothing?) = visitLocalVariable(localVariable)
+    fun visitModifierElement(modifierElement: JKModifierElement) = visitTreeElement(modifierElement, null)
+    override fun visitModifierElement(modifierElement: JKModifierElement, data: Nothing?) = visitModifierElement(modifierElement)
+    fun visitMutabilityModifierElement(mutabilityModifierElement: JKMutabilityModifierElement) = visitModifierElement(mutabilityModifierElement, null)
+    override fun visitMutabilityModifierElement(mutabilityModifierElement: JKMutabilityModifierElement, data: Nothing?) = visitMutabilityModifierElement(mutabilityModifierElement)
+    fun visitModalityModifierElement(modalityModifierElement: JKModalityModifierElement) = visitModifierElement(modalityModifierElement, null)
+    override fun visitModalityModifierElement(modalityModifierElement: JKModalityModifierElement, data: Nothing?) = visitModalityModifierElement(modalityModifierElement)
+    fun visitVisibilityModifierElement(visibilityModifierElement: JKVisibilityModifierElement) = visitModifierElement(visibilityModifierElement, null)
+    override fun visitVisibilityModifierElement(visibilityModifierElement: JKVisibilityModifierElement, data: Nothing?) = visitVisibilityModifierElement(visibilityModifierElement)
+    fun visitExtraModifierElement(extraModifierElement: JKExtraModifierElement) = visitModifierElement(extraModifierElement, null)
+    override fun visitExtraModifierElement(extraModifierElement: JKExtraModifierElement, data: Nothing?) = visitExtraModifierElement(extraModifierElement)
     fun visitExtraModifiersOwner(extraModifiersOwner: JKExtraModifiersOwner) = visitModifiersListOwner(extraModifiersOwner, null)
     override fun visitExtraModifiersOwner(extraModifiersOwner: JKExtraModifiersOwner, data: Nothing?) = visitExtraModifiersOwner(extraModifiersOwner)
     fun visitVisibilityOwner(visibilityOwner: JKVisibilityOwner) = visitModifiersListOwner(visibilityOwner, null)
@@ -49,6 +59,8 @@ interface JKVisitorVoid : JKVisitor<Unit, Nothing?> {
     override fun visitStatement(statement: JKStatement, data: Nothing?) = visitStatement(statement)
     fun visitBlock(block: JKBlock) = visitTreeElement(block, null)
     override fun visitBlock(block: JKBlock, data: Nothing?) = visitBlock(block)
+    fun visitBodyStub(bodyStub: JKBodyStub) = visitBlock(bodyStub, null)
+    override fun visitBodyStub(bodyStub: JKBodyStub, data: Nothing?) = visitBodyStub(bodyStub)
     fun visitIdentifier(identifier: JKIdentifier) = visitTreeElement(identifier, null)
     override fun visitIdentifier(identifier: JKIdentifier, data: Nothing?) = visitIdentifier(identifier)
     fun visitNameIdentifier(nameIdentifier: JKNameIdentifier) = visitIdentifier(nameIdentifier, null)
@@ -223,7 +235,7 @@ interface JKVisitorVoid : JKVisitor<Unit, Nothing?> {
     override fun visitKtProperty(ktProperty: JKKtProperty, data: Nothing?) = visitKtProperty(ktProperty)
     fun visitKtFunction(ktFunction: JKKtFunction) = visitMethod(ktFunction, null)
     override fun visitKtFunction(ktFunction: JKKtFunction, data: Nothing?) = visitKtFunction(ktFunction)
-    fun visitKtConstructor(ktConstructor: JKKtConstructor) = visitDeclaration(ktConstructor, null)
+    fun visitKtConstructor(ktConstructor: JKKtConstructor) = visitMethod(ktConstructor, null)
     override fun visitKtConstructor(ktConstructor: JKKtConstructor, data: Nothing?) = visitKtConstructor(ktConstructor)
     fun visitKtPrimaryConstructor(ktPrimaryConstructor: JKKtPrimaryConstructor) = visitKtConstructor(ktPrimaryConstructor, null)
     override fun visitKtPrimaryConstructor(ktPrimaryConstructor: JKKtPrimaryConstructor, data: Nothing?) = visitKtPrimaryConstructor(ktPrimaryConstructor)