New J2K: Replace backAnnotator with PsiOwner

j2k/newSrc/org/jetbrains/kotlin/j2k/ConversionContext.kt

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.j2k.tree.JKElement
 data class ConversionContext(
     val symbolProvider: JKSymbolProvider,
     val converter: NewJavaToKotlinConverter,
-    val backAnnotator: (JKElement) -> PsiElement?,
     val inConversionContext: (PsiElement) -> Boolean
 ) {
 

j2k/newSrc/org/jetbrains/kotlin/j2k/JavaToJKTreeBuilder.kt

@@ -41,8 +41,6 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
 
     private val modifierMapper = ModifierMapper()
 
-    val backAnnotation = mutableMapOf<JKElement, PsiElement>()//TODO preserve consistency
-
     private inner class ExpressionTreeMapper {
         fun PsiExpression?.toJK(): JKExpression {
             return when (this) {
@@ -74,7 +72,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                     throw RuntimeException("Not supported: ${this::class}")
                 }
             }.also {
-                if (this != null) backAnnotation[it] = this
+                if (this != null) (it as PsiOwner).psi = this
             }
         }
 
@@ -191,7 +189,9 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                     JKJavaNewEmptyArrayImpl(
                         dimensions.map { it?.toJK() ?: JKStubExpressionImpl() },
                         JKTypeElementImpl(generateSequence(type?.toJK(symbolProvider)) { it.safeAs<JKJavaArrayType>()?.type }.last())
-                    )
+                    ).also {
+                        it.psi = this
+                    }
                 }
             }
             val constructedClass = classOrAnonymousClassReference?.resolve()
@@ -227,7 +227,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
 
         fun PsiTypeElement.toJK(): JKTypeElement {
             return JKTypeElementImpl(type.toJK(symbolProvider)).also {
-                backAnnotation[it] = this
+                (it as PsiOwner).psi = this
             }
         }
     }
@@ -253,7 +253,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                 jkClassImpl.declarationList = children.mapNotNull {
                     ElementVisitor().apply { it.accept(this) }.resultElement as? JKDeclaration
                 }
-                backAnnotation[jkClassImpl] = this
+                jkClassImpl.psi = this
                 symbolProvider.provideUniverseSymbol(this, jkClassImpl)
             }
         }
@@ -266,7 +266,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                 with(expressionTreeMapper) { initializer.toJK() }
             ).also {
                 symbolProvider.provideUniverseSymbol(this, it)
-                backAnnotation[it] = this
+                it.psi = this
             }
         }
 
@@ -282,9 +282,8 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                 parameterList.parameters.map { it.toJK() },
                 body?.toJK() ?: JKBodyStub
             ).also {
-                backAnnotation[it] = this
+                it.psi = this
                 symbolProvider.provideUniverseSymbol(this, it)
-                backAnnotation[it] = this
             }
         }
 
@@ -299,7 +298,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                                    JKNameIdentifierImpl(name!!),
                                    with(modifierMapper) { modifierList.toJK() }).also {
                 symbolProvider.provideUniverseSymbol(this, it)
-                backAnnotation[it] = this
+                it.psi = this
             }
         }
 
@@ -317,7 +316,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                         with(expressionTreeMapper) { it.initializer.toJK() }
                     ).also { i ->
                         symbolProvider.provideUniverseSymbol(it, i)
-                        backAnnotation[i] = it
+                        i.psi = it
                     }
                 } else TODO()
             }
@@ -382,7 +381,7 @@ class JavaToJKTreeBuilder(var symbolProvider: JKSymbolProvider) {
                 }
                 else -> TODO("for ${this::class}")
             }.also {
-                if (this != null) backAnnotation[it] = this
+                if (this != null) (it as PsiOwner).psi = this
             }
         }
     }

j2k/newSrc/org/jetbrains/kotlin/j2k/NewJavaToKotlinConverter.kt

@@ -50,7 +50,6 @@ class NewJavaToKotlinConverter(
         val context = ConversionContext(
             symbolProvider,
             this,
-            { it: JKElement -> treeBuilder.backAnnotation[it] },
             { it: PsiElement -> it.containingFile in files }
         )
 

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/AssignmentAsExpressionToAlsoConversion.kt

@@ -17,7 +17,7 @@ class AssignmentAsExpressionToAlsoConversion(val context: ConversionContext) : R
         val alsoExpression = JKKtAlsoCallExpressionImpl(
             JKBlockStatementImpl(
                 JKBlockImpl(listOf(JKKtAssignmentStatementImpl(element.field, JKStubExpressionImpl(), element.operator)))
-            ), context.symbolProvider.provideByFqName("kotlin/also", context.backAnnotator.invoke(element.parentOfType<JKClass>()!!)!!)
+            ), context.symbolProvider.provideByFqName("kotlin/also", element.parentOfType<JKClass>()!!.psi!!)
         ).also {
             it.statement.cast<JKBlockStatement>().block.statements.first().cast<JKKtAssignmentStatement>().expression =
                     JKFieldAccessExpressionImpl(context.symbolProvider.provideUniverseSymbol(it.parameter))

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/BlockToRunConversion.kt

@@ -21,7 +21,7 @@ class BlockToRunConversion(private val context: ConversionContext) : RecursiveAp
         if (element.parent !is JKBlock) return recurse(element)
 
         val parentDeclaration = element.parentOfType<JKDeclaration>() ?: return recurse(element)
-        val psiContext = context.backAnnotator(parentDeclaration) ?: return recurse(element)
+        val psiContext = parentDeclaration.psi ?: return recurse(element)
 
         val unitType =
             context.symbolProvider.provideDirectSymbol(

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/ForConversion.kt

@@ -263,5 +263,5 @@ class ForConversion(private val context: ConversionContext) : RecursiveApplicabl
     }
 
     private inline fun <reified ElementType : PsiElement> JKElement.psi() =
-        context.backAnnotator(this) as? ElementType
+        this.psi as? ElementType
 }

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/ModalityConversion.kt

@@ -10,6 +10,7 @@ import com.intellij.psi.PsiMethod
 import org.jetbrains.kotlin.j2k.ConversionContext
 import org.jetbrains.kotlin.j2k.tree.*
 import org.jetbrains.kotlin.j2k.tree.impl.modality
+import org.jetbrains.kotlin.j2k.tree.impl.psi
 
 class ModalityConversion(private val context: ConversionContext) : RecursiveApplicableConversionBase() {
 
@@ -26,13 +27,13 @@ class ModalityConversion(private val context: ConversionContext) : RecursiveAppl
         if (klass.classKind != JKClass.ClassKind.CLASS) return
         if (klass.modifierList.modality != JKModalityModifier.Modality.OPEN) return
         if (context.converter.settings.openByDefault) return
-        if (!context.converter.converterServices.oldServices.referenceSearcher.hasInheritors(context.backAnnotator(klass) as PsiClass)) {
+        if (!context.converter.converterServices.oldServices.referenceSearcher.hasInheritors(klass.psi as PsiClass)) {
             klass.modifierList.modality = JKModalityModifier.Modality.FINAL
         }
     }
 
     private fun processMethod(method: JKJavaMethod) {
-        if ((context.backAnnotator(method)!! as PsiMethod).findSuperMethods().isNotEmpty()) {
+        if ((method.psi!! as PsiMethod).findSuperMethods().isNotEmpty()) {
             method.modifierList.modality = JKModalityModifier.Modality.OVERRIDE
         }
     }

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/PrintlnConversion.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.j2k.ConversionContext
 import org.jetbrains.kotlin.j2k.tree.*
 import org.jetbrains.kotlin.j2k.tree.impl.JKKtCallExpressionImpl
 import org.jetbrains.kotlin.j2k.tree.impl.JKMethodSymbol
+import org.jetbrains.kotlin.j2k.tree.impl.psi
 import org.jetbrains.kotlin.name.ClassId
 
 // TODO: Full special methods conversion
@@ -31,7 +32,7 @@ class PrintlnConversion(private val context: ConversionContext) : RecursiveAppli
         if (selector.identifier.name != "println") return element
 
         val contextElement = element.parentOfType<JKClass>() ?: return element
-        val targetElements = multiResolveFqName(ClassId.fromString("kotlin/io/println"), context.backAnnotator(contextElement)!!)
+        val targetElements = multiResolveFqName(ClassId.fromString("kotlin/io/println"), contextElement.psi!!)
         if (targetElements.isEmpty()) return element
         selector.invalidate()
 

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/SwitchStatementConversion.kt

@@ -92,7 +92,7 @@ class SwitchStatementConversion(private val context: ConversionContext) : Recurs
             is JKBreakStatement, is JKReturnStatement -> false
             is JKBlockStatement -> block.statements.fallsThrough()
             is JKIfStatement, is JKJavaSwitchStatement, is JKKtWhenStatement ->
-                context.backAnnotator(this)!!.canCompleteNormally()
+                this.psi!!.canCompleteNormally()
             else -> true
         }
 

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/TypeMappingConversion.kt

@@ -14,11 +14,10 @@ import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
 import org.jetbrains.kotlin.j2k.*
 import org.jetbrains.kotlin.j2k.ast.Nullability
 import org.jetbrains.kotlin.j2k.tree.*
-import org.jetbrains.kotlin.j2k.tree.impl.JKClassTypeImpl
+import org.jetbrains.kotlin.j2k.tree.impl.*
-import org.jetbrains.kotlin.j2k.tree.impl.JKJavaVoidType
-import org.jetbrains.kotlin.j2k.tree.impl.JKTypeElementImpl
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtClassOrObject
 
 class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicableConversionBase() {
 
@@ -59,7 +58,7 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
         is JKJavaVoidType -> JKClassTypeImpl(
             context.symbolProvider.provideByFqName(
                 ClassId.topLevel(KotlinBuiltIns.FQ_NAMES.unit.toSafe()),
-                context.backAnnotator.invoke(element.parentOfType<JKClass>()!!)!!
+                element.parentOfType<JKClass>()!!.psi!!
             ),
             nullability = Nullability.NotNull
         )
@@ -87,9 +86,9 @@ class TypeMappingConversion(val context: ConversionContext) : RecursiveApplicabl
 
     private fun calculateNullability(parent: JKElement?): Nullability {
         return when (parent) {
-            is JKJavaMethod -> typeFlavorCalculator.methodNullability(context.backAnnotator(parent) as PsiMethod)
+            is JKJavaMethod -> typeFlavorCalculator.methodNullability(parent.psi as PsiMethod)
-            is JKJavaField -> typeFlavorCalculator.variableNullability(context.backAnnotator(parent) as PsiVariable)
+            is JKJavaField -> typeFlavorCalculator.variableNullability(parent.psi as PsiVariable)
-            is JKLocalVariable -> typeFlavorCalculator.variableNullability(context.backAnnotator(parent) as PsiVariable)
+            is JKLocalVariable -> typeFlavorCalculator.variableNullability(parent.psi as PsiVariable)
             else -> Nullability.Default
         }
     }

j2k/newSrc/org/jetbrains/kotlin/j2k/conversions/resolveUtil.kt

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiPolyVariantReference
 import org.jetbrains.kotlin.j2k.ConversionContext
 import org.jetbrains.kotlin.j2k.tree.JKTreeElement
+import org.jetbrains.kotlin.j2k.tree.impl.psi
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtImportDirective
@@ -19,7 +20,7 @@ import org.jetbrains.kotlin.resolve.ImportPath
 
 
 internal fun resolveFqName(classId: ClassId, contextElement: JKTreeElement, context: ConversionContext): PsiElement? {
-    val element = context.backAnnotator(contextElement) ?: return null
+    val element = contextElement.psi ?: return null
     return resolveFqName(classId, element)
 }
 

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/impl/j.kt

@@ -27,7 +27,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.resolve.jvm.JvmPrimitiveType
 
 class JKJavaFieldImpl(modifierList: JKModifierList, type: JKTypeElement, name: JKNameIdentifier, initializer: JKExpression) : JKJavaField,
-    JKBranchElementBase() {
+    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaField(this, data)
 
     override var initializer: JKExpression by child(initializer)
@@ -38,7 +38,7 @@ class JKJavaFieldImpl(modifierList: JKModifierList, type: JKTypeElement, name: J
 
 class JKJavaMethodImpl(
     modifierList: JKModifierList, returnType: JKTypeElement, name: JKNameIdentifier, parameters: List<JKParameter>, block: JKBlock
-) : JKJavaMethod, JKBranchElementBase() {
+) : JKJavaMethod, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaMethod(this, data)
 
     override var modifierList: JKModifierList by child(modifierList)
@@ -52,7 +52,7 @@ class JKJavaMethodImpl(
 class JKJavaLiteralExpressionImpl(
     override val literal: String,
     override val type: JKLiteralExpression.LiteralType
-) : JKJavaLiteralExpression, JKElementBase() {
+) : JKJavaLiteralExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaLiteralExpression(this, data)
 
     init {
@@ -60,7 +60,7 @@ class JKJavaLiteralExpressionImpl(
     }
 }
 
-class JKJavaModifierImpl(override val type: JKJavaModifier.JavaModifierType) : JKJavaModifier, JKElementBase() {
+class JKJavaModifierImpl(override val type: JKJavaModifier.JavaModifierType) : JKJavaModifier, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaModifier(this, data)
 }
 
@@ -173,17 +173,18 @@ class JKJavaMethodCallExpressionImpl(
     override var identifier: JKMethodSymbol,
     arguments: JKExpressionList,
     typeArguments: List<JKTypeElement> = emptyList()
-) : JKJavaMethodCallExpression, JKBranchElementBase() {
+) : JKJavaMethodCallExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaMethodCallExpression(this, data)
 
     override val arguments: JKExpressionList by child(arguments)
     override var typeArguments: List<JKTypeElement> by children(typeArguments)
 }
 
+
 class JKJavaNewExpressionImpl(
     override val constructorSymbol: JKMethodSymbol,
     arguments: JKExpressionList
-) : JKJavaNewExpression, JKBranchElementBase() {
+) : JKJavaNewExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaNewExpression(this, data)
 
     override var arguments by child(arguments)
@@ -191,9 +192,9 @@ class JKJavaNewExpressionImpl(
 
 class JKJavaDefaultNewExpressionImpl(
     override val classSymbol: JKClassSymbol
-) : JKJavaDefaultNewExpression, JKElementBase()
+) : JKJavaDefaultNewExpression, JKElementBase(), PsiOwner by PsiOwnerImpl()
 
-class JKJavaNewEmptyArrayImpl(initializer: List<JKExpression>, type: JKTypeElement) : JKJavaNewEmptyArray, JKBranchElementBase() {
+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)
@@ -231,13 +232,13 @@ object JKJavaVoidType : JKType {
 class JKJavaArrayTypeImpl(override val type: JKType, override var nullability: Nullability = Nullability.Default) : JKJavaArrayType {
 }
 
-class JKReturnStatementImpl(expression: JKExpression) : JKBranchElementBase(), JKReturnStatement {
+class JKReturnStatementImpl(expression: JKExpression) : JKBranchElementBase(), 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)
 }
 
-class JKJavaAssertStatementImpl(condition: JKExpression, description: JKExpression) : JKJavaAssertStatement, JKBranchElementBase() {
+class JKJavaAssertStatementImpl(condition: JKExpression, description: JKExpression) : JKJavaAssertStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override val description by child(description)
     override val condition by child(condition)
 
@@ -245,7 +246,7 @@ class JKJavaAssertStatementImpl(condition: JKExpression, description: JKExpressi
 }
 
 class JKJavaForLoopStatementImpl(initializer: JKStatement, condition: JKExpression, updater: JKStatement, body: JKStatement) :
-    JKJavaForLoopStatement, JKBranchElementBase() {
+    JKJavaForLoopStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var body by child(body)
     override var updater by child(updater)
     override var condition by child(condition)
@@ -254,7 +255,7 @@ class JKJavaForLoopStatementImpl(initializer: JKStatement, condition: JKExpressi
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaForLoopStatement(this, data)
 }
 
-class JKJavaInstanceOfExpressionImpl(expression: JKExpression, type: JKTypeElement) : JKJavaInstanceOfExpression, JKBranchElementBase() {
+class JKJavaInstanceOfExpressionImpl(expression: JKExpression, type: JKTypeElement) : JKJavaInstanceOfExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var type by child(type)
     override var expression by child(expression)
 
@@ -262,7 +263,7 @@ class JKJavaInstanceOfExpressionImpl(expression: JKExpression, type: JKTypeEleme
 }
 
 class JKJavaPolyadicExpressionImpl(operands: List<JKExpression>, override var tokens: List<JKOperator>) : JKJavaPolyadicExpression,
-    JKBranchElementBase() {
+    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var operands by children(operands)
 
     override fun getTokenBeforeOperand(operand: JKExpression): JKOperator? {
@@ -277,7 +278,7 @@ class JKJavaAssignmentExpressionImpl(
     override var field: JKAssignableExpression,
     expression: JKExpression,
     override var operator: JKOperator
-) : JKBranchElementBase(), JKJavaAssignmentExpression {
+) : JKBranchElementBase(), JKJavaAssignmentExpression, PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaAssignmentExpression(this, data)
 
     override var expression: JKExpression by child(expression)
@@ -286,13 +287,13 @@ class JKJavaAssignmentExpressionImpl(
 class JKJavaSwitchStatementImpl(
     expression: JKExpression,
     cases: List<JKJavaSwitchCase>
-) : JKJavaSwitchStatement, JKBranchElementBase() {
+) : JKJavaSwitchStatement, JKBranchElementBase(), 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() {
+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)
 }
@@ -300,7 +301,7 @@ class JKJavaDefaultSwitchCaseImpl(statements: List<JKStatement>) : JKJavaDefault
 class JKJavaLabelSwitchCaseImpl(
     label: JKExpression,
     statements: List<JKStatement>
-) : JKJavaLabelSwitchCase, JKBranchElementBase() {
+) : JKJavaLabelSwitchCase, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var statements: List<JKStatement> by children(statements)
     override var label: JKExpression by child(label)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitJavaLabelSwitchCase(this, data)

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/impl/jk.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.j2k.tree.impl
 
 import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.j2k.ConversionContext
 import org.jetbrains.kotlin.j2k.JKSymbolProvider
@@ -32,7 +33,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.KtClass
 import org.jetbrains.kotlin.psi.KtNamedFunction
 
-class JKFileImpl : JKFile, JKBranchElementBase() {
+class JKFileImpl : JKFile, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitFile(this, data)
 
     override var declarationList by children<JKDeclaration>()
@@ -43,7 +44,7 @@ class JKClassImpl(
     name: JKNameIdentifier,
     inheritance: JKInheritanceInfo,
     override var classKind: JKClass.ClassKind
-) : JKClass, JKBranchElementBase() {
+) : JKClass, JKBranchElementBase(), 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)
@@ -52,13 +53,13 @@ class JKClassImpl(
     override val inheritance by child(inheritance)
 }
 
-class JKNameIdentifierImpl(override val value: String) : JKNameIdentifier, JKElementBase() {
+class JKNameIdentifierImpl(override val value: String) : JKNameIdentifier, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitNameIdentifier(this, data)
 }
 
 class JKModifierListImpl(
     modifiers: List<JKModifier> = emptyList()
-) : JKModifierList, JKBranchElementBase() {
+) : JKModifierList, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     constructor(vararg modifiers: JKModifier) : this(modifiers.asList())
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitModifierList(this, data)
@@ -92,7 +93,7 @@ class JKParameterImpl(
     modifierList: JKModifierList,
     initializer: JKExpression = JKStubExpressionImpl()
 
-) : JKParameter, JKBranchElementBase() {
+) : JKParameter, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitParameter(this, data)
 
     override var modifierList by child(modifierList)
@@ -101,7 +102,7 @@ class JKParameterImpl(
     override var type by child(type)
 }
 
-class JKBlockImpl(statements: List<JKStatement> = emptyList()) : JKBlock, JKBranchElementBase() {
+class JKBlockImpl(statements: List<JKStatement> = emptyList()) : JKBlock, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBlock(this, data)
 
     override var statements by children(statements)
@@ -113,7 +114,7 @@ class JKBinaryExpressionImpl(
     right: JKExpression,
     override var operator: JKOperator
 ) : JKBinaryExpression,
-    JKBranchElementBase() {
+    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBinaryExpression(this, data)
     override var right by child(right)
     override var left by child(left)
@@ -163,21 +164,20 @@ 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() {
     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() {
+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() {
+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)
@@ -187,20 +187,20 @@ class JKQualifiedExpressionImpl(
     receiver: JKExpression,
     override var operator: JKQualifier,
     selector: JKExpression
-) : JKQualifiedExpression, JKBranchElementBase() {
+) : JKQualifiedExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitQualifiedExpression(this, data)
 
     override var receiver: JKExpression by child(receiver)
     override var selector: JKExpression by child(selector)
 }
 
-class JKExpressionStatementImpl(expression: JKExpression) : JKExpressionStatement, JKBranchElementBase() {
+class JKExpressionStatementImpl(expression: JKExpression) : JKExpressionStatement, JKBranchElementBase(), 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() {
+class JKDeclarationStatementImpl(declaredStatements: List<JKDeclaration>) : JKDeclarationStatement, JKBranchElementBase(), 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)
@@ -209,26 +209,26 @@ class JKDeclarationStatementImpl(declaredStatements: List<JKDeclaration>) : JKDe
 class JKArrayAccessExpressionImpl(
     expression: JKExpression,
     indexExpression: JKExpression
-) : JKArrayAccessExpression, JKBranchElementBase() {
+) : JKArrayAccessExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitArrayAccessExpression(this, data)
 
     override var expression: JKExpression by child(expression)
     override var indexExpression: JKExpression by child(indexExpression)
 }
 
-class JKParenthesizedExpressionImpl(expression: JKExpression) : JKParenthesizedExpression, JKBranchElementBase() {
+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() {
+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)
 }
 
-class JKTypeElementImpl(override var type: JKType) : JKTypeElement, JKElementBase() {
+class JKTypeElementImpl(override var type: JKType) : JKTypeElement, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitTypeElement(this, data)
 }
 
@@ -264,7 +264,7 @@ fun JKType.fqName(): String =
 fun JKType.equalsByName(other: JKType) = fqName() == other.fqName()
 
 
-class JKNullLiteral : JKLiteralExpression, JKElementBase() {
+class JKNullLiteral : JKLiteralExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override val literal: String
         get() = "null"
     override val type: LiteralType
@@ -273,7 +273,7 @@ class JKNullLiteral : JKLiteralExpression, JKElementBase() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLiteralExpression(this, data)
 }
 
-class JKBooleanLiteral(val value: Boolean) : JKLiteralExpression, JKElementBase() {
+class JKBooleanLiteral(val value: Boolean) : JKLiteralExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override val literal: String
         get() = value.toString()
     override val type: LiteralType
@@ -305,7 +305,7 @@ fun JKLiteralExpression.LiteralType.toJkType(symbolProvider: JKSymbolProvider):
 }
 
 class JKLocalVariableImpl(modifierList: JKModifierList, type: JKTypeElement, name: JKNameIdentifier, initializer: JKExpression) :
-    JKLocalVariable, JKBranchElementBase() {
+    JKLocalVariable, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var modifierList by child(modifierList)
     override var initializer by child(initializer)
     override var name by child(name)
@@ -314,7 +314,7 @@ class JKLocalVariableImpl(modifierList: JKModifierList, type: JKTypeElement, nam
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLocalVariable(this, data)
 }
 
-class JKStubExpressionImpl : JKStubExpression, JKElementBase() {
+class JKStubExpressionImpl : JKStubExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitStubExpression(this, data)
 }
 
@@ -339,42 +339,41 @@ object JKBodyStub : JKBlock, JKTreeElement {
     }
 }
 
-class JKBlockStatementImpl(block: JKBlock) : JKBlockStatement, JKBranchElementBase() {
+class JKBlockStatementImpl(block: JKBlock) : JKBlockStatement, JKBranchElementBase(), 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 JKThisExpressionImpl : JKThisExpression, JKElementBase() {
+class JKThisExpressionImpl : JKThisExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitThisExpression(this, data)
 }
 
-class JKSuperExpressionImpl : JKSuperExpression, JKElementBase() {
+class JKSuperExpressionImpl : JKSuperExpression, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitSuperExpression(this, data)
 }
 
-class JKWhileStatementImpl(condition: JKExpression, body: JKStatement) : JKWhileStatement, JKBranchElementBase() {
+class JKWhileStatementImpl(condition: JKExpression, body: JKStatement) : JKWhileStatement, JKBranchElementBase(), 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() {
+class JKDoWhileStatementImpl(body: JKStatement, condition: JKExpression) : JKDoWhileStatement, JKBranchElementBase(), 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, JKElementBase() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitBreakStatement(this, data)
 }
 
-class JKBreakWithLabelStatementImpl(override var label: JKNameIdentifier) : JKBreakWithLabelStatement, JKElementBase() {
+class JKBreakWithLabelStatementImpl(override var label: JKNameIdentifier) : JKBreakWithLabelStatement, JKElementBase(), 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() {
+class JKIfStatementImpl(condition: JKExpression, thenBranch: JKStatement) : JKIfStatement, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var thenBranch by child(thenBranch)
     override var condition by child(condition)
 
@@ -382,7 +381,7 @@ class JKIfStatementImpl(condition: JKExpression, thenBranch: JKStatement) : JKIf
 }
 
 class JKIfElseStatementImpl(condition: JKExpression, thenBranch: JKStatement, elseBranch: JKStatement) : JKIfElseStatement,
-    JKBranchElementBase() {
+    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var elseBranch by child(elseBranch)
     override var thenBranch by child(thenBranch)
     override var condition by child(condition)
@@ -391,7 +390,7 @@ class JKIfElseStatementImpl(condition: JKExpression, thenBranch: JKStatement, el
 }
 
 class JKIfElseExpressionImpl(condition: JKExpression, thenBranch: JKExpression, elseBranch: JKExpression) : JKIfElseExpression,
-    JKBranchElementBase() {
+    JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var elseBranch by child(elseBranch)
     override var thenBranch by child(thenBranch)
     override var condition by child(condition)
@@ -399,19 +398,19 @@ 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() {
+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)
 }
 
-class JKModalityModifierImpl(override val modality: JKModalityModifier.Modality) : JKModalityModifier, JKElementBase() {
+class JKModalityModifierImpl(override val modality: JKModalityModifier.Modality) : JKModalityModifier, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitModalityModifier(this, data)
 }
 
-class JKAccessModifierImpl(override val visibility: JKAccessModifier.Visibility) : JKAccessModifier, JKElementBase() {
+class JKAccessModifierImpl(override val visibility: JKAccessModifier.Visibility) : JKAccessModifier, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitAccessModifier(this, data)
 }
 
-class JKMutabilityModifierImpl(override val mutability: Mutability) : JKMutabilityModifier, JKElementBase() {
+class JKMutabilityModifierImpl(override val mutability: Mutability) : JKMutabilityModifier, JKElementBase(), PsiOwner by PsiOwnerImpl() {
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitMutabilityModifier(this, data)
 }
 
@@ -419,14 +418,14 @@ class JKLambdaExpressionImpl(
     parameters: List<JKParameter> = listOf(
         JKParameterImpl(JKTypeElementImpl(JKJavaVoidType), JKNameIdentifierImpl("it"), JKModifierListImpl())
     ), statement: JKStatement, returnType: JKTypeElement = JKTypeElementImpl(JKContextType)
-) : JKLambdaExpression, JKBranchElementBase() {
+) : JKLambdaExpression, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override var statement by child(statement)
     override val returnType by child(returnType)
     override var parameters by children(parameters)
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitLambdaExpression(this, data)
 }
 
-class JKInheritanceInfoImpl(implements: List<JKTypeElement>) : JKInheritanceInfo, JKBranchElementBase() {
+class JKInheritanceInfoImpl(implements: List<JKTypeElement>) : JKInheritanceInfo, JKBranchElementBase(), PsiOwner by PsiOwnerImpl() {
     override val inherit: List<JKTypeElement> by children(implements)
 
     override fun <R, D> accept(visitor: JKVisitor<R, D>, data: D): R = visitor.visitInheritanceInfo(this, data)
@@ -436,7 +435,7 @@ class JKDelegationConstructorCallImpl(
     override val identifier: JKMethodSymbol,
     expression: JKExpression,
     arguments: JKExpressionList
-) : JKBranchElementBase(), JKDelegationConstructorCall {
+) : JKBranchElementBase(), JKDelegationConstructorCall, PsiOwner by PsiOwnerImpl() {
     override var typeArguments by children(emptyList<JKTypeElement>())
     override val expression: JKExpression by child(expression)
     override val arguments: JKExpressionList by child(arguments)
@@ -444,7 +443,7 @@ 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() {
+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>
@@ -454,21 +453,26 @@ val JKStatement.statements: List<JKStatement>
             else -> listOf(this)
         }
 
-class JKLabelEmptyImpl : JKLabelEmpty, JKElementBase() {
+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() {
+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() {
+class JKContinueStatementImpl(label: JKLabel) : JKContinueStatement, JKBranchElementBase(), 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() {
+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 PsiOwnerImpl(override var psi: PsiElement? = null) : PsiOwner
+
+val JKElement.psi: PsiElement?
+    get() = (this as? PsiOwner)?.psi

j2k/newSrc/org/jetbrains/kotlin/j2k/tree/jk.kt

@@ -35,6 +35,10 @@ interface JKTreeElement : JKElement {
     fun copy(): JKTreeElement
 }
 
+interface PsiOwner {
+    var psi: PsiElement?
+}
+
 interface JKDeclaration : JKTreeElement
 
 interface JKFile : JKTreeElement, JKBranchElement {