Return correct parent for parts of converted elvis expression

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
 
 abstract class KotlinAbstractUElement(private val givenParent: UElement?) : UElement {
 
@@ -96,6 +97,13 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
         return result.tempVarAssignment
     }
 
+    if (result is KotlinUElvisExpression && parent is KtBinaryExpression) {
+        when (element.psi) {
+            parent.left -> return result.lhsDeclaration
+            parent.right -> return result.rhsIfExpression
+        }
+    }
+
     return result
 }
 

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

@@ -72,26 +72,39 @@ fun createElvisExpression(elvisExpression: KtBinaryExpression, givenParent: UEle
     val left = elvisExpression.left ?: return UastEmptyExpression
     val right = elvisExpression.right ?: return UastEmptyExpression
 
-    return object : KotlinAbstractUElement(givenParent), UExpressionList, KotlinEvaluatableUElement {
+    return KotlinUElvisExpression(elvisExpression, left, right, givenParent)
-        override val psi: PsiElement? = elvisExpression
+}
-        override val kind = KotlinSpecialExpressionKinds.ELVIS
-        override val annotations: List<UAnnotation> = emptyList()
-        override val expressions: List<UExpression> by lz {
-            createElvisExpressions(left, right, this, elvisExpression.parent)
-        }
-        override fun asRenderString(): String {
-            return kind.name + " " +
-                   expressions.joinToString(separator = "\n", prefix = "{\n", postfix = "\n}") {
-                       it.asRenderString().withMargin
-                   }
-        }
-        override fun getExpressionType(): PsiType? {
-            val leftType = left.analyze()[BindingContext.EXPRESSION_TYPE_INFO, left]?.type ?: return null
-            val rightType = right.analyze()[BindingContext.EXPRESSION_TYPE_INFO, right]?.type ?: return null
 
-            return CommonSupertypes
+class KotlinUElvisExpression(
-                    .commonSupertype(listOf(leftType, rightType))
+    private val elvisExpression: KtBinaryExpression,
-                    .toPsiType(this, elvisExpression, boxed = false)
+    private val left: KtExpression,
-        }
+    private val right: KtExpression,
+    givenParent: UElement?
+) : KotlinAbstractUElement(givenParent), UExpressionList, KotlinEvaluatableUElement {
+
+    override val psi: PsiElement? = elvisExpression
+    override val kind = KotlinSpecialExpressionKinds.ELVIS
+    override val annotations: List<UAnnotation> = emptyList()
+    override val expressions: List<UExpression> by lz {
+        createElvisExpressions(left, right, this, elvisExpression.parent)
     }
-}
+
+    val lhsDeclaration get() = (expressions[0] as UDeclarationsExpression).declarations.single()
+    val rhsIfExpression get() = expressions[1] as UIfExpression
+
+    override fun asRenderString(): String {
+        return kind.name + " " +
+               expressions.joinToString(separator = "\n", prefix = "{\n", postfix = "\n}") {
+                   it.asRenderString().withMargin
+               }
+    }
+
+    override fun getExpressionType(): PsiType? {
+        val leftType = left.analyze()[BindingContext.EXPRESSION_TYPE_INFO, left]?.type ?: return null
+        val rightType = right.analyze()[BindingContext.EXPRESSION_TYPE_INFO, right]?.type ?: return null
+
+        return CommonSupertypes
+            .commonSupertype(listOf(leftType, rightType))
+            .toPsiType(this, elvisExpression, boxed = false)
+    }
+}