Elvis -> if-then: handle case with safe cast separately #KT-17816 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/IfThenUtils.kt

@@ -106,7 +106,12 @@ fun KtExpression.convertToIfStatement(condition: KtExpression, thenClause: KtExp
 
 fun KtIfExpression.introduceValueForCondition(occurrenceInThenClause: KtExpression, editor: Editor?) {
     val project = this.project
-    val occurrenceInConditional = (this.condition as KtBinaryExpression).left!!
+    val condition = condition
+    val occurrenceInConditional = when (condition) {
+        is KtBinaryExpression -> condition.left
+        is KtIsExpression -> condition.leftHandSide
+        else -> throw AssertionError("Only binary / is expressions are supported here: ${condition?.text}")
+    }!!
     KotlinIntroduceVariableHandler.doRefactoring(project,
                                                  editor,
                                                  occurrenceInConditional,

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/ElvisToIfThenIntention.kt

@@ -19,13 +19,18 @@ package org.jetbrains.kotlin.idea.intentions.branchedTransformations.intentions
 import com.intellij.codeInsight.intention.LowPriorityAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.convertToIfNotNullExpression
+import org.jetbrains.kotlin.idea.intentions.branchedTransformations.convertToIfStatement
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.introduceValueForCondition
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.isStable
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.kotlin.psi.KtBinaryExpression
-import org.jetbrains.kotlin.psi.KtPsiUtil
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.types.TypeUtils
 
 class ElvisToIfThenIntention : SelfTargetingRangeIntention<KtBinaryExpression>(KtBinaryExpression::class.java, "Replace elvis expression with 'if' expression"), LowPriorityAction {
     override fun applicabilityRange(element: KtBinaryExpression): TextRange? {
@@ -35,13 +40,57 @@ class ElvisToIfThenIntention : SelfTargetingRangeIntention<KtBinaryExpression>(K
             null
     }
 
+    private fun KtExpression.findSafeCastReceiver(context: BindingContext): KtBinaryExpressionWithTypeRHS? {
+        var current = this
+        while (current is KtQualifiedExpression) {
+            val resolvedCall = current.selectorExpression.getResolvedCall(context) ?: return null
+            val type = resolvedCall.resultingDescriptor.returnType
+            if (type != null && TypeUtils.isNullableType(type)) {
+                return null
+            }
+            current = current.receiverExpression
+        }
+        current = KtPsiUtil.safeDeparenthesize(current)
+        return (current as? KtBinaryExpressionWithTypeRHS)?.takeIf {
+            it.operationReference.getReferencedNameElementType() === KtTokens.AS_SAFE &&
+            it.right != null
+        }
+    }
+
+    private fun KtExpression.buildExpressionWithReplacedReceiver(
+            factory: KtPsiFactory,
+            newReceiver: KtExpression,
+            topLevel: Boolean = true
+    ): KtExpression {
+        if (this !is KtQualifiedExpression) {
+            return newReceiver
+        }
+        return factory.buildExpression(reformat = topLevel) {
+            appendExpression(receiverExpression.buildExpressionWithReplacedReceiver(factory, newReceiver, topLevel = false))
+            appendFixedText(".")
+            appendExpression(selectorExpression)
+        }
+    }
+
     override fun applyTo(element: KtBinaryExpression, editor: Editor?) {
+        val context = element.analyze(BodyResolveMode.PARTIAL)
         val left = KtPsiUtil.safeDeparenthesize(element.left!!)
         val right = KtPsiUtil.safeDeparenthesize(element.right!!)
 
-        val leftIsStable = left.isStable()
-
-        val ifStatement = element.convertToIfNotNullExpression(left, left, right)
+        val leftSafeCastReceiver = left.findSafeCastReceiver(context)
+        val (leftIsStable, ifStatement) = if (leftSafeCastReceiver != null) {
+            val newReceiver = leftSafeCastReceiver.left
+            val typeReference = leftSafeCastReceiver.right!!
+            val factory = KtPsiFactory(element)
+            newReceiver.isStable(context) to element.convertToIfStatement(
+                    factory.createExpressionByPattern("$0 is $1", newReceiver, typeReference),
+                    left.buildExpressionWithReplacedReceiver(factory, newReceiver),
+                    right
+            )
+        }
+        else {
+            left.isStable(context) to element.convertToIfNotNullExpression(left, left, right)
+        }
 
         if (!leftIsStable) {
             ifStatement.introduceValueForCondition(ifStatement.then!!, editor)

idea/testData/intentions/branched/elvisToIfThen/fakeSafeCast.kt

@@ -0,0 +1,4 @@
+class My(val x: Int?)
+fun foo(arg: Any) {
+    val y = (arg as? My)?.x <caret>?: 42
+}

idea/testData/intentions/branched/elvisToIfThen/fakeSafeCast.kt.after

@@ -0,0 +1,4 @@
+class My(val x: Int?)
+fun foo(arg: Any) {
+    val y = if ((arg as? My)?.x != null) (arg as? My)?.x else 42
+}

idea/testData/intentions/branched/elvisToIfThen/longSafeCast.kt

@@ -0,0 +1,4 @@
+class My(val x: Int)
+fun foo(arg: Any) {
+    val y = (arg as? My)?.x?.hashCode() <caret>?: 42
+}

idea/testData/intentions/branched/elvisToIfThen/longSafeCast.kt.after

@@ -0,0 +1,4 @@
+class My(val x: Int)
+fun foo(arg: Any) {
+    val y = if (arg is My) arg.x.hashCode() else 42
+}

idea/testData/intentions/branched/elvisToIfThen/safeCast.kt

@@ -0,0 +1,4 @@
+class My(val x: Int)
+fun foo(arg: Any) {
+    val y = (arg as? My)?.x <caret>?: 42
+}

idea/testData/intentions/branched/elvisToIfThen/safeCast.kt.after

@@ -0,0 +1,4 @@
+class My(val x: Int)
+fun foo(arg: Any) {
+    val y = if (arg is My) arg.x else 42
+}

idea/testData/intentions/branched/elvisToIfThen/safeCastUnstable.kt

@@ -0,0 +1,5 @@
+class My(var z: Any, val x: Int) {
+    fun foo() {
+        val y = (z as? My)?.x <caret>?: 42
+    }
+}

idea/testData/intentions/branched/elvisToIfThen/safeCastUnstable.kt.after

@@ -0,0 +1,6 @@
+class My(var z: Any, val x: Int) {
+    fun foo() {
+        val z1 = z
+        val y = if (z1 is My) z1 else 42
+    }
+}

idea/testData/intentions/branched/elvisToIfThen/simpleSafeCast.kt

@@ -0,0 +1,4 @@
+class My(val x: Int)
+fun foo(arg: Any) {
+    val y = arg as? My <caret>?: My(42)
+}

idea/testData/intentions/branched/elvisToIfThen/simpleSafeCast.kt.after

@@ -0,0 +1,4 @@
+class My(val x: Int)
+fun foo(arg: Any) {
+    val y = if (arg is My) arg else My(42)
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -1225,6 +1225,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("fakeSafeCast.kt")
+            public void testFakeSafeCast() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/fakeSafeCast.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("localValLhs.kt")
             public void testLocalValLhs() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/localValLhs.kt");
@@ -1237,6 +1243,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("longSafeCast.kt")
+            public void testLongSafeCast() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/longSafeCast.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("safeCast.kt")
+            public void testSafeCast() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/safeCast.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("safeCastUnstable.kt")
+            public void testSafeCastUnstable() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/safeCastUnstable.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("simpleNameExpression.kt")
             public void testSimpleNameExpression() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/simpleNameExpression.kt");
@@ -1249,6 +1273,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("simpleSafeCast.kt")
+            public void testSimpleSafeCast() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/simpleSafeCast.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("topLevelVal.kt")
             public void testTopLevelVal() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/elvisToIfThen/topLevelVal.kt");