Return correct parent for parts of converted elvis expression

2017-09-15 18:48:07 +02:00
parent e55f4c6e78
commit a7efc4f5da
2 changed files with 42 additions and 21 deletions
@@ -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
 }

@@ -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 {
-        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 KotlinUElvisExpression(elvisExpression, left, right, givenParent)
+}

-            return CommonSupertypes
-                    .commonSupertype(listOf(leftType, rightType))
-                    .toPsiType(this, elvisExpression, boxed = false)
-        }
+class KotlinUElvisExpression(
+    private val elvisExpression: KtBinaryExpression,
+    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)
+    }
+}