UAST: support for JvmDeclarationUElement

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt

@@ -28,7 +28,7 @@ import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
-abstract class KotlinAbstractUElement(private val givenParent: UElement?) : UElement {
+abstract class KotlinAbstractUElement(private val givenParent: UElement?) : UElement, JvmDeclarationUElement {
 
     override val uastParent: UElement? by lz {
         givenParent ?: convertParent()
@@ -152,7 +152,11 @@ private fun findAnnotationClassFromConstructorParameter(parameter: KtParameter):
 }
 
 abstract class KotlinAbstractUExpression(givenParent: UElement?)
-    : KotlinAbstractUElement(givenParent), UExpression {
+    : KotlinAbstractUElement(givenParent), UExpression, JvmDeclarationUElement {
+
+    override val javaPsi = null
+    override val sourcePsi
+        get() = psi
 
     override val annotations: List<UAnnotation>
         get() {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinUastLanguagePlugin.kt

@@ -64,7 +64,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
     override fun convertElement(element: PsiElement, parent: UElement?, requiredType: Class<out UElement>?): UElement? {
         return convertDeclarationOrElement(element, parent, requiredType)
     }
-    
+
     override fun convertElementWithParent(element: PsiElement, requiredType: Class<out UElement>?): UElement? {
         if (element is PsiFile) return convertDeclaration(element, null, requiredType)
         if (element is KtLightClassForFacade) return convertDeclaration(element, null, requiredType)
@@ -100,7 +100,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
         val resolvedCall = element.getResolvedCall(element.analyze()) ?: return null
         val resultingDescriptor = resolvedCall.resultingDescriptor
         if (resultingDescriptor !is FunctionDescriptor || resultingDescriptor.name.asString() != methodName) return null
-        
+
         val parent = element.parent
         val parentUElement = convertElementWithParent(parent, null) ?: return null
 
@@ -117,7 +117,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
         if (element !is KtCallExpression) return null
         val resolvedCall = element.getResolvedCall(element.analyze()) ?: return null
         val resultingDescriptor = resolvedCall.resultingDescriptor
-        if (resultingDescriptor !is ConstructorDescriptor 
+        if (resultingDescriptor !is ConstructorDescriptor
                 || resultingDescriptor.returnType.constructor.declarationDescriptor?.name?.asString() != fqName) {
             return null
         }
@@ -134,9 +134,13 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
     internal fun convertDeclaration(element: PsiElement,
                                    givenParent: UElement?,
                                    requiredType: Class<out UElement>?): UElement? {
-        fun <P : PsiElement> build(ctor: (P, UElement?) -> UElement): () -> UElement? {
+        fun <P : PsiElement> build(ctor: (P, UElement?) -> UElement): () -> UElement? = { ctor(element as P, givenParent) }
-            return { ctor(element as P, givenParent) }
+
-        }
+        fun <P : PsiElement, K : KtElement> buildKt(ktElement: K, ctor: (P, K, UElement?) -> UElement): () -> UElement? =
+                { ctor(element as P, ktElement, givenParent) }
+
+        fun <P : PsiElement, K : KtElement> buildKtOpt(ktElement: K?, ctor: (P, K?, UElement?) -> UElement): () -> UElement? =
+                { ctor(element as P, ktElement, givenParent) }
 
         val original = element.originalElement
         return with(requiredType) {
@@ -148,10 +152,11 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
                     }
                     else -> el<UClass> { KotlinUClass.create(original, givenParent) }
                 }
-                is KtLightFieldImpl.KtLightEnumConstant -> el<UEnumConstant>(build(::KotlinUEnumConstant))
+                is KtLightFieldImpl.KtLightEnumConstant -> el<UEnumConstant>(buildKtOpt(original.kotlinOrigin, ::KotlinUEnumConstant))
-                is KtLightField -> el<UField>(build(::KotlinUField))
+                is KtLightField -> el<UField>(buildKtOpt(original.kotlinOrigin, ::KotlinUField))
-                is KtLightParameter, is UastKotlinPsiParameter -> el<UParameter>(build(::KotlinUParameter))
+                is KtLightParameter -> el<UParameter>(buildKtOpt(original.kotlinOrigin, ::KotlinUParameter))
-                is UastKotlinPsiVariable -> el<UVariable>(build(::KotlinUVariable))
+                is UastKotlinPsiParameter -> el<UParameter>(buildKt(original.ktParameter, ::KotlinUParameter))
+                is UastKotlinPsiVariable -> el<UVariable>(buildKt(original.ktElement, ::KotlinUVariable))
 
                 is KtEnumEntry -> el<UEnumConstant> {
                     convertEnumEntry(original, givenParent)
@@ -198,7 +203,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
                     val ownerFunction = original.ownerFunction as? KtFunction ?: return null
                     val lightMethod = LightClassUtil.getLightClassMethod(ownerFunction) ?: return null
                     val lightParameter = lightMethod.parameterList.parameters.find { it.name == original.name } ?: return null
-                    KotlinUParameter(lightParameter, givenParent)
+                    KotlinUParameter(lightParameter, original, givenParent)
                 }
 
                 is KtFile -> el<UFile> { KotlinUFile(original, this@KotlinUastLanguagePlugin) }
@@ -213,7 +218,7 @@ class KotlinUastLanguagePlugin : UastLanguagePlugin {
     private fun convertEnumEntry(original: KtEnumEntry, givenParent: UElement?): UElement? {
         return LightClassUtil.getLightClassBackingField(original)?.let { psiField ->
             if (psiField is KtLightFieldImpl.KtLightEnumConstant) {
-                KotlinUEnumConstant(psiField, givenParent)
+                KotlinUEnumConstant(psiField, psiField.kotlinOrigin, givenParent)
             }
             else {
                 null
@@ -247,7 +252,7 @@ private fun convertNonLocalProperty(property: KtProperty,
     val methods = LightClassUtil.getLightClassPropertyMethods(property)
     return methods.backingField?.let { backingField ->
         with(requiredType) {
-            el<UField> { KotlinUField(backingField, givenParent) }
+            el<UField> { KotlinUField(backingField, (backingField as? KtLightElement<*,*>)?.kotlinOrigin,  givenParent) }
         }
     } ?: methods.getter?.let { getter ->
         KotlinUastLanguagePlugin().convertDeclaration(getter, givenParent, requiredType)
@@ -277,7 +282,7 @@ internal object KotlinConverter {
                 val declarationsExpression = KotlinUDeclarationsExpression(givenParent)
                 declarationsExpression.apply {
                     declarations = element.parameters.mapIndexed { i, p ->
-                        KotlinUParameter(UastKotlinPsiParameter.create(p, element, declarationsExpression, i), this)
+                        KotlinUParameter(UastKotlinPsiParameter.create(p, element, declarationsExpression, i), p, this)
                     }
                 }
             }
@@ -287,7 +292,7 @@ internal object KotlinConverter {
                 if (element is KtProperty && !element.isLocal) {
                     el<UField> {
                         LightClassUtil.getLightClassBackingField(element)?.let {
-                            KotlinUField(it, givenParent)
+                            KotlinUField(it, element, givenParent)
                         }
                     }
                 }
@@ -375,13 +380,13 @@ internal object KotlinConverter {
             is KtDestructuringDeclaration -> expr<UDeclarationsExpression> {
                 val declarationsExpression = KotlinUDestructuringDeclarationExpression(givenParent, expression)
                 declarationsExpression.apply {
-                    val tempAssignment = KotlinULocalVariable(UastKotlinPsiVariable.create(expression, declarationsExpression), declarationsExpression)
+                    val tempAssignment = KotlinULocalVariable(UastKotlinPsiVariable.create(expression, declarationsExpression), expression, declarationsExpression)
                     val destructuringAssignments = expression.entries.mapIndexed { i, entry ->
                         val psiFactory = KtPsiFactory(expression.project)
                         val initializer = psiFactory.createAnalyzableExpression("${tempAssignment.name}.component${i + 1}()",
                                                                                 expression.containingFile)
                         initializer.destructuringDeclarationInitializer = true
-                        KotlinULocalVariable(UastKotlinPsiVariable.create(entry, tempAssignment.psi, declarationsExpression, initializer), declarationsExpression)
+                        KotlinULocalVariable(UastKotlinPsiVariable.create(entry, tempAssignment.psi, declarationsExpression, initializer), entry, declarationsExpression)
                     }
                     declarations = listOf(tempAssignment) + destructuringAssignments
                 }
@@ -505,7 +510,7 @@ private fun convertVariablesDeclaration(
     val declarationsExpression = KotlinUDeclarationsExpression(null, parent, psi)
     val parentPsiElement = parent?.psi
     val variable = KotlinUAnnotatedLocalVariable(
-            UastKotlinPsiVariable.create(psi, parentPsiElement, declarationsExpression), declarationsExpression) { annotationParent ->
+            UastKotlinPsiVariable.create(psi, parentPsiElement, declarationsExpression), psi, declarationsExpression) { annotationParent ->
         psi.annotationEntries.map { KotlinUAnnotation(it, annotationParent) }
     }
     return declarationsExpression.apply { declarations = listOf(variable) }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUAnnotation.kt

@@ -1,6 +1,7 @@
 package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiClass
+import org.jetbrains.kotlin.asJava.toLightAnnotation
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
 import org.jetbrains.kotlin.psi.KtAnnotationEntry
 import org.jetbrains.kotlin.psi.KtParameter
@@ -17,6 +18,11 @@ class KotlinUAnnotation(
         override val psi: KtAnnotationEntry,
         givenParent: UElement?
 ) : KotlinAbstractUElement(givenParent), UAnnotation {
+
+    override val javaPsi = psi.toLightAnnotation()
+
+    override val sourcePsi = psi
+
     private val resolvedAnnotation: AnnotationDescriptor? by lz { psi.analyze()[BindingContext.ANNOTATION, psi] }
 
     private val resolvedCall: ResolvedCall<*>? by lz { psi.getResolvedCall(psi.analyze()) }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUClass.kt

@@ -31,7 +31,7 @@ import org.jetbrains.uast.java.AbstractJavaUClass
 import org.jetbrains.uast.kotlin.declarations.KotlinUMethod
 import org.jetbrains.uast.kotlin.declarations.UastLightIdentifier
 
-abstract class AbstractKotlinUClass(private val givenParent: UElement?) : AbstractJavaUClass() {
+abstract class AbstractKotlinUClass(private val givenParent: UElement?) : AbstractJavaUClass(), JvmDeclarationUElement {
     override val uastParent: UElement? by lz { givenParent ?: convertParent() }
 
     //TODO: should be merged with KotlinAbstractUElement.convertParent() after detaching from AbstractJavaUClass
@@ -42,7 +42,6 @@ abstract class AbstractKotlinUClass(private val givenParent: UElement?) : Abstra
                     else -> it.toUElement()
                 }
             } ?: (psi.parent ?: psi.containingFile).toUElement()
-
 }
 
 open class KotlinUClass private constructor(
@@ -52,6 +51,10 @@ open class KotlinUClass private constructor(
 
     val ktClass = psi.kotlinOrigin
 
+    override val javaPsi: KtLightClass = psi
+
+    override val sourcePsi: KtClassOrObject? = ktClass
+
     override val psi = unwrap<UClass, PsiClass>(psi)
 
     override fun getOriginalElement(): PsiElement? = super.getOriginalElement()
@@ -114,6 +117,7 @@ open class KotlinUClass private constructor(
             else -> KotlinUClass(psi, containingElement)
         }
     }
+
 }
 
 open class KotlinConstructorUMethod(
@@ -148,6 +152,10 @@ open class KotlinConstructorUMethod(
         }
     }
 
+    override val javaPsi = psi
+
+    override val sourcePsi = psi.kotlinOrigin
+
     open protected fun getBodyExpressions(): List<KtExpression> {
         if (isPrimary) return getInitializers()
         val bodyExpression = (psi.kotlinOrigin as? KtFunction)?.bodyExpression ?: return emptyList()
@@ -177,6 +185,10 @@ class KotlinUAnonymousClass(
 
     override val psi: PsiAnonymousClass = unwrap<UAnonymousClass, PsiAnonymousClass>(psi)
 
+    override val javaPsi: PsiAnonymousClass = psi
+
+    override val sourcePsi: KtClassOrObject? = (psi as? KtLightClass)?.kotlinOrigin
+
     override fun getOriginalElement(): PsiElement? = super<AbstractKotlinUClass>.getOriginalElement()
 
     override fun getSuperClass(): UClass? = super<AbstractKotlinUClass>.getSuperClass()
@@ -206,6 +218,10 @@ class KotlinScriptUClass(
     override val uastAnchor: UElement
         get() = UIdentifier(nameIdentifier, this)
 
+    override val javaPsi: PsiClass = psi
+
+    override val sourcePsi: KtClassOrObject? = psi.kotlinOrigin
+
     override val psi = unwrap<UClass, KtLightClassForScript>(psi)
 
     override fun getSuperClass(): UClass? = super.getSuperClass()
@@ -239,5 +255,7 @@ class KotlinScriptUClass(
             val initializers = script.declarations.filterIsInstance<KtScriptInitializer>()
             KotlinUBlockExpression.create(initializers, this)
         }
+        override val javaPsi = psi
+        override val sourcePsi = psi.kotlinOrigin
     }
 }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUFile.kt

@@ -18,6 +18,7 @@ package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiRecursiveElementWalkingVisitor
 import org.jetbrains.kotlin.asJava.findFacadeClass
 import org.jetbrains.kotlin.asJava.toLightClass
@@ -26,13 +27,17 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.uast.*
 import java.util.*
 
-class KotlinUFile(override val psi: KtFile, override val languagePlugin: UastLanguagePlugin) : UFile {
+class KotlinUFile(override val psi: KtFile, override val languagePlugin: UastLanguagePlugin) : UFile, JvmDeclarationUElement {
     override val packageName: String
         get() = psi.packageFqName.asString()
 
     override val annotations: List<UAnnotation>
         get() = psi.annotationEntries.map { KotlinUAnnotation(it, this) }
 
+    override val javaPsi: PsiElement? = null
+
+    override val sourcePsi: PsiElement? = psi
+
     override val allCommentsInFile by lz {
         val comments = ArrayList<UComment>(0)
         psi.accept(object : PsiRecursiveElementWalkingVisitor() {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUImportStatement.kt

@@ -30,6 +30,11 @@ class KotlinUImportStatement(
         override val psi: KtImportDirective,
         givenParent: UElement?
 ) : KotlinAbstractUElement(givenParent), UImportStatement {
+
+    override val javaPsi = null
+
+    override val sourcePsi = psi
+
     override val isOnDemand: Boolean
         get() = psi.isAllUnder
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt

@@ -38,6 +38,10 @@ open class KotlinUMethod(
 ) : KotlinAbstractUElement(givenParent), UAnnotationMethod, JavaUElementWithComments, PsiMethod by psi {
     override val psi: KtLightMethod = unwrap<UMethod, KtLightMethod>(psi)
 
+    override val javaPsi = psi
+
+    override val sourcePsi = psi.kotlinOrigin
+
     override val uastDefaultValue by lz {
         val annotationParameter = psi.kotlinOrigin as? KtParameter ?: return@lz null
         val defaultValue = annotationParameter.defaultValue ?: return@lz null
@@ -57,7 +61,7 @@ open class KotlinUMethod(
     }
 
     override val uastParameters by lz {
-        psi.parameterList.parameters.map { KotlinUParameter(it, this) }
+        psi.parameterList.parameters.map { KotlinUParameter(it, (it as? KtLightElement<*, *>)?.kotlinOrigin, this) }
     }
 
     override val uastAnchor: UElement

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUVariable.kt

@@ -19,10 +19,7 @@ package org.jetbrains.uast.kotlin
 import com.intellij.psi.*
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
-import org.jetbrains.kotlin.psi.KtNamedDeclaration
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtParameter
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.KtVariableDeclaration
 import org.jetbrains.uast.*
 import org.jetbrains.uast.internal.acceptList
 import org.jetbrains.uast.java.JavaAbstractUExpression
@@ -86,9 +83,13 @@ abstract class AbstractKotlinUVariable(givenParent: UElement?)
 
 class KotlinUVariable(
         psi: PsiVariable,
+        override val sourcePsi: KtElement,
         givenParent: UElement?
 ) : AbstractKotlinUVariable(givenParent), UVariable, PsiVariable by psi {
-    override val psi = unwrap<UVariable, PsiVariable>(psi)
+
+    override val javaPsi = unwrap<UVariable, PsiVariable>(psi)
+
+    override val psi = javaPsi
 
     override val annotations by lz { psi.annotations.map { JavaUAnnotation(it, this) } }
 
@@ -113,10 +114,13 @@ class KotlinUVariable(
 
 open class KotlinUParameter(
         psi: PsiParameter,
+        override val sourcePsi: KtElement?,
         givenParent: UElement?
 ) : AbstractKotlinUVariable(givenParent), UParameter, PsiParameter by psi {
 
-    override val psi = unwrap<UParameter, PsiParameter>(psi)
+    override val javaPsi = unwrap<UParameter, PsiParameter>(psi)
+
+    override val psi = javaPsi
 
     override fun getInitializer(): PsiExpression? {
         return super<AbstractKotlinUVariable>.getInitializer()
@@ -137,10 +141,13 @@ open class KotlinUParameter(
 
 open class KotlinUField(
         psi: PsiField,
+        override val sourcePsi: KtElement?,
         givenParent: UElement?
 ) : AbstractKotlinUVariable(givenParent), UField, PsiField by psi {
 
-    override val psi = unwrap<UField, PsiField>(psi)
+    override val javaPsi  = unwrap<UField, PsiField>(psi)
+
+    override val psi = javaPsi
 
     override fun getInitializer(): PsiExpression? {
         return super<AbstractKotlinUVariable>.getInitializer()
@@ -173,10 +180,13 @@ open class KotlinUField(
 
 open class KotlinULocalVariable(
         psi: PsiLocalVariable,
+        override val sourcePsi: KtElement,
         givenParent: UElement?
 ) : AbstractKotlinUVariable(givenParent), ULocalVariable, PsiLocalVariable by psi {
 
-    override val psi = unwrap<ULocalVariable, PsiLocalVariable>(psi)
+    override val javaPsi = unwrap<ULocalVariable, PsiLocalVariable>(psi)
+
+    override val psi = javaPsi
 
     override fun getInitializer(): PsiExpression? {
         return super<AbstractKotlinUVariable>.getInitializer()
@@ -205,24 +215,29 @@ open class KotlinULocalVariable(
 
 open class KotlinUAnnotatedLocalVariable(
         psi: PsiLocalVariable,
+        sourcePsi: KtElement,
         uastParent: UElement?,
         computeAnnotations: (parent: UElement) -> List<UAnnotation>
-) : KotlinULocalVariable(psi, uastParent) {
+) : KotlinULocalVariable(psi, sourcePsi, uastParent) {
 
     override val annotations: List<UAnnotation> by lz { computeAnnotations(this) }
 }
 
 open class KotlinUEnumConstant(
         psi: PsiEnumConstant,
+        override val sourcePsi: KtElement?,
         givenParent: UElement?
 ) : AbstractKotlinUVariable(givenParent), UEnumConstant, PsiEnumConstant by psi {
+
     override val initializingClass: UClass? by lz {
         (psi.initializingClass as? KtLightClass)?.let { initializingClass ->
             KotlinUClass.create(initializingClass, this)
         }
     }
 
-    override val psi = unwrap<UEnumConstant, PsiEnumConstant>(psi)
+    override val javaPsi = unwrap<UEnumConstant, PsiEnumConstant>(psi)
+
+    override val psi = javaPsi
 
     override fun getContainingFile(): PsiFile {
         return super.getContainingFile()
@@ -245,7 +260,7 @@ open class KotlinUEnumConstant(
         get() = null
 
     override val classReference: UReferenceExpression?
-        get() = KotlinEnumConstantClassReference(psi, this)
+        get() = KotlinEnumConstantClassReference(psi, sourcePsi, this)
 
     override val typeArgumentCount: Int
         get() = 0
@@ -272,8 +287,12 @@ open class KotlinUEnumConstant(
 
     private class KotlinEnumConstantClassReference(
             override val psi: PsiEnumConstant,
+            override val sourcePsi: KtElement?,
             private val givenParent: UElement?
     ) : JavaAbstractUExpression(), USimpleNameReferenceExpression {
+        override val javaPsi: PsiElement?
+            get() = psi
+
         override val uastParent: UElement? by lz {
             givenParent ?: KotlinUastLanguagePlugin().convertElementWithParent(psi.parent ?: psi.containingFile, null)
         }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/ElvisExpression.kt

@@ -14,25 +14,29 @@ import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
 
 private fun createVariableReferenceExpression(variable: UVariable, containingElement: UElement?) =
-        object : USimpleNameReferenceExpression {
+        object : USimpleNameReferenceExpression, JvmDeclarationUElement {
             override val psi: PsiElement? = null
             override fun resolve(): PsiElement? = variable
             override val uastParent: UElement? = containingElement
             override val resolvedName: String? = variable.name
             override val annotations: List<UAnnotation> = emptyList()
             override val identifier: String = variable.name.orAnonymous()
+            override val javaPsi: PsiElement? = null
+            override val sourcePsi: PsiElement? = null
         }
 
 private fun createNullLiteralExpression(containingElement: UElement?) =
-        object : ULiteralExpression {
+        object : ULiteralExpression, JvmDeclarationUElement {
             override val psi: PsiElement? = null
             override val uastParent: UElement? = containingElement
             override val value: Any? = null
             override val annotations: List<UAnnotation> = emptyList()
+            override val javaPsi: PsiElement? = null
+            override val sourcePsi: PsiElement? = null
         }
 
 private fun createNotEqWithNullExpression(variable: UVariable, containingElement: UElement?) =
-        object : UBinaryExpression {
+        object : UBinaryExpression, JvmDeclarationUElement {
             override val psi: PsiElement? = null
             override val uastParent: UElement? = containingElement
             override val leftOperand: UExpression by lz { createVariableReferenceExpression(variable, this) }
@@ -41,6 +45,8 @@ private fun createNotEqWithNullExpression(variable: UVariable, containingElement
             override val operatorIdentifier: UIdentifier? = UIdentifier(null, this)
             override fun resolveOperator(): PsiMethod? = null
             override val annotations: List<UAnnotation> = emptyList()
+            override val javaPsi: PsiElement? = null
+            override val sourcePsi: PsiElement? = null
         }
 
 private fun createElvisExpressions(
@@ -50,12 +56,14 @@ private fun createElvisExpressions(
         psiParent: PsiElement): List<UExpression> {
 
     val declaration = KotlinUDeclarationsExpression(containingElement)
-    val tempVariable = KotlinULocalVariable(UastKotlinPsiVariable.create(left, declaration, psiParent), declaration)
+    val tempVariable = KotlinULocalVariable(UastKotlinPsiVariable.create(left, declaration, psiParent), left, declaration)
     declaration.declarations = listOf(tempVariable)
 
-    val ifExpression = object : UIfExpression {
+    val ifExpression = object : UIfExpression, JvmDeclarationUElement {
         override val psi: PsiElement? = null
         override val uastParent: UElement? = containingElement
+        override val javaPsi: PsiElement? = null
+        override val sourcePsi: PsiElement? = null
         override val condition: UExpression by lz { createNotEqWithNullExpression(tempVariable, this) }
         override val thenExpression: UExpression? by lz { createVariableReferenceExpression(tempVariable, this) }
         override val elseExpression: UExpression? by lz { KotlinConverter.convertExpression(right, this ) }
@@ -82,7 +90,9 @@ class KotlinUElvisExpression(
     givenParent: UElement?
 ) : KotlinAbstractUElement(givenParent), UExpressionList, KotlinEvaluatableUElement {
 
-    override val psi: PsiElement? = elvisExpression
+    override val javaPsi: PsiElement? = null
+    override val sourcePsi: PsiElement? = elvisExpression
+    override val psi: PsiElement? = sourcePsi
     override val kind = KotlinSpecialExpressionKinds.ELVIS
     override val annotations: List<UAnnotation> = emptyList()
     override val expressions: List<UExpression> by lz {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUBlockExpression.kt

@@ -30,8 +30,10 @@ class KotlinUBlockExpression(
     class KotlinLazyUBlockExpression(
             override val uastParent: UElement?,
             expressionProducer: (expressionParent: UElement) -> List<UExpression>
-    ) : UBlockExpression {
+    ) : UBlockExpression, JvmDeclarationUElement {
         override val psi: PsiElement? = null
+        override val javaPsi: PsiElement? = null
+        override val sourcePsi: PsiElement? = null
         override val annotations: List<UAnnotation> = emptyList()
         override val expressions by lz { expressionProducer(this) }
     }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUCatchClause.kt

@@ -27,11 +27,15 @@ class KotlinUCatchClause(
         override val psi: KtCatchClause,
         givenParent: UElement?
 ) : KotlinAbstractUElement(givenParent), UCatchClause {
+
+    override val javaPsi = null
+    override val sourcePsi = psi
+
     override val body by lz { KotlinConverter.convertOrEmpty(psi.catchBody, this) }
     
     override val parameters by lz {
         val parameter = psi.catchParameter ?: return@lz emptyList<UParameter>()
-        listOf(KotlinUParameter(UastKotlinPsiParameter.create(parameter, psi, this, 0), this))
+        listOf(KotlinUParameter(UastKotlinPsiParameter.create(parameter, psi, this, 0), psi, this))
     }
 
     override val typeReferences by lz {

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUDeclarationsExpression.kt

@@ -25,6 +25,9 @@ open class KotlinUDeclarationsExpression(
         val psiAnchor: PsiElement? = null
 ) : KotlinAbstractUExpression(givenParent), UDeclarationsExpression {
 
+    override val sourcePsi: PsiElement?
+        get() = psiAnchor
+
     override val uastParent: UElement?
         get() = if (psiAnchor != null) doConvertParent(this, psiAnchor.parent) else super.uastParent
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUForEachExpression.kt

@@ -34,7 +34,7 @@ class KotlinUForEachExpression(
     override val variable by lz {
         val parameter = psi.loopParameter?.let { UastKotlinPsiParameter.create(it, psi, this, 0) } 
                 ?: UastPsiParameterNotResolved(psi, KotlinLanguage.INSTANCE)
-        KotlinUParameter(parameter, this)
+        KotlinUParameter(parameter, psi, this)
     }
 
     override val forIdentifier: UIdentifier

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinULambdaExpression.kt

@@ -35,7 +35,7 @@ class KotlinULambdaExpression(
     
     override val valueParameters by lz {
         psi.valueParameters.mapIndexed { i, p ->
-            KotlinUParameter(UastKotlinPsiParameter.create(p, psi, this, i), this)
+            KotlinUParameter(UastKotlinPsiParameter.create(p, psi, this, i), psi, this)
         }
     }
     

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUNamedExpression.kt

@@ -23,19 +23,23 @@ import org.jetbrains.uast.*
 
 class KotlinUNamedExpression private constructor(
         override val name: String?,
+        override val sourcePsi: PsiElement?,
         givenParent: UElement?,
         expressionProducer: (UElement) -> UExpression
 ) : KotlinAbstractUElement(givenParent), UNamedExpression {
+
     override val expression: UExpression by lz { expressionProducer(this) }
 
     override val annotations: List<UAnnotation> = emptyList()
 
     override val psi: PsiElement? = null
 
+    override val javaPsi: PsiElement? = null
+
     companion object {
         internal fun create(name: String?, valueArgument: ValueArgument, uastParent: UElement?): UNamedExpression {
             val expression = valueArgument.getArgumentExpression()
-            return KotlinUNamedExpression(name, uastParent) { expressionParent ->
+            return KotlinUNamedExpression(name, valueArgument.asElement(), uastParent) { expressionParent ->
                 expression?.let { expressionParent.getLanguagePlugin().convert<UExpression>(it, expressionParent) } ?: UastEmptyExpression
             }
         }
@@ -44,7 +48,7 @@ class KotlinUNamedExpression private constructor(
                 name: String?,
                 valueArguments: List<ValueArgument>,
                 uastParent: UElement?): UNamedExpression {
-            return KotlinUNamedExpression(name, uastParent) { expressionParent ->
+            return KotlinUNamedExpression(name, null, uastParent) { expressionParent ->
                 KotlinUVarargExpression(valueArguments, expressionParent)
             }
         }

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUObjectLiteralExpression.kt

@@ -75,6 +75,10 @@ class KotlinUObjectLiteralExpression(
             override val psi: KtSuperTypeCallEntry,
             givenParent: UElement?
     ) : KotlinAbstractUElement(givenParent), USimpleNameReferenceExpression {
+
+        override val javaPsi = null
+        override val sourcePsi = psi
+
         override fun resolve() = (psi.resolveCallToDeclaration(this) as? PsiMethod)?.containingClass
 
         override val annotations: List<UAnnotation>

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/KotlinUSwitchExpression.kt

@@ -69,7 +69,9 @@ class KotlinUSwitchEntry(
                 is KtBlockExpression -> exprPsi.statements.map { KotlinConverter.convertOrEmpty(it, this) }
                 else -> listOf(KotlinConverter.convertOrEmpty(exprPsi, this))
             }
-            expressions = userExpressions + object : UBreakExpression {
+            expressions = userExpressions + object : UBreakExpression, JvmDeclarationUElement {
+                override val javaPsi: PsiElement? = null
+                override val sourcePsi: PsiElement? = null
                 override val psi: PsiElement?
                     get() = null
                 override val label: String?

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.uast.kotlin.expressions
 
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiType
 import com.intellij.psi.PsiVariable
 import org.jetbrains.kotlin.psi.KtFunction
@@ -14,6 +15,10 @@ internal class KotlinLocalFunctionUVariable(
         override val psi: PsiVariable,
         givenParent: UElement?
 ) : KotlinAbstractUElement(givenParent), UVariable, PsiVariable by psi {
+
+    override val javaPsi = psi
+    override val sourcePsi: PsiElement? = (psi as? UastKotlinPsiVariable?)?.ktElement ?: psi
+
     override val uastInitializer: UExpression? by lz {
         createLocalFunctionLambdaExpression(function, this)
     }
@@ -33,7 +38,7 @@ private class KotlinLocalFunctionULambdaExpression(
 
     override val valueParameters by lz {
         psi.valueParameters.mapIndexed { i, p ->
-            KotlinUParameter(UastKotlinPsiParameter.create(p, psi, this, i), this)
+            KotlinUParameter(UastKotlinPsiParameter.create(p, psi, this, i), p, this)
         }
     }
 

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/internal/KotlinUElementWithComments.kt

@@ -17,10 +17,11 @@
 package org.jetbrains.uast.kotlin.internal
 
 import com.intellij.psi.PsiComment
+import org.jetbrains.uast.JvmDeclarationUElement
 import org.jetbrains.uast.UComment
 import org.jetbrains.uast.UElement
 
-interface KotlinUElementWithComments : UElement {
+interface KotlinUElementWithComments : UElement, JvmDeclarationUElement {
     override val comments: List<UComment>
         get() {
             val psi = psi ?: return emptyList()

plugins/uast-kotlin/tests/AbstractKotlinRenderLogTest.kt

@@ -3,11 +3,11 @@ package org.jetbrains.uast.test.kotlin
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiRecursiveElementVisitor
+import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.addToStdlib.assertedCast
-import org.jetbrains.uast.UDeclaration
+import org.jetbrains.uast.*
-import org.jetbrains.uast.UElement
+import org.jetbrains.uast.java.JavaUAnnotation
-import org.jetbrains.uast.UFile
 import org.jetbrains.uast.kotlin.KOTLIN_CACHED_UELEMENT_KEY
 import org.jetbrains.uast.kotlin.KotlinUastLanguagePlugin
 import org.jetbrains.uast.test.common.RenderLogTestBase
@@ -32,6 +32,7 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
         }
 
         file.checkContainingFileForAllElements()
+        file.checkJvmDeclarationsImplementations()
     }
 
     private fun checkParentConsistency(file: UFile) {
@@ -95,6 +96,30 @@ abstract class AbstractKotlinRenderLogTest : AbstractKotlinUastTest(), RenderLog
             }
         })
     }
+
+    private fun UFile.checkJvmDeclarationsImplementations() {
+        accept(object : UastVisitor {
+            override fun visitElement(node: UElement): Boolean {
+
+                val jvmDeclaration = node as? JvmDeclarationUElement
+                                     ?: throw AssertionError("${node.javaClass} should implement 'JvmDeclarationUElement'")
+
+                if (jvmDeclaration is JavaUAnnotation) return false // but actually it's strange to meet JavaUAnnotation here
+                if (jvmDeclaration is UIdentifier) return false // probably should be fixed in platform to fully support in in Kotlin
+
+                jvmDeclaration.sourcePsi?.let {
+                    assertTrue("sourcePsi should be physical but ${it.javaClass} found for [${it.text}] " +
+                               "for ${jvmDeclaration.javaClass}->${jvmDeclaration.uastParent?.javaClass}", it is KtElement)
+                }
+                jvmDeclaration.javaPsi?.let {
+                    assertTrue("javaPsi should be light but ${it.javaClass} found for [${it.text}] " +
+                               "for ${jvmDeclaration.javaClass}->${jvmDeclaration.uastParent?.javaClass}", it !is KtElement)
+                }
+
+                return false
+            }
+        })
+    }
 }
 
 private fun PsiFile.clearUastCaches() {