From 2117b30b76d403bb4e7435b72495275b528e9a1a Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Mon, 10 Oct 2016 16:24:03 +0300
Subject: [PATCH] If expression to elvis now handles call chains #KT-14084
 Fixed

---
 .../ReplaceSingleLineLetIntention.kt          | 23 -------
 .../jetbrains/kotlin/idea/intentions/Utils.kt | 25 +++++++
 .../intentions/IfThenToElvisIntention.kt      | 65 ++++++++++++++-----
 .../ifThenToElvis/isCheckWithSelector.kt      |  5 ++
 .../isCheckWithSelector.kt.after              |  5 ++
 .../ifThenToElvis/nullCheckWithSelector.kt    |  5 ++
 .../nullCheckWithSelector.kt.after            |  5 ++
 .../nullCheckWithSelectorCall.kt              |  9 +++
 .../nullCheckWithSelectorCall.kt.after        |  5 ++
 .../nullCheckWithSelectorCallChain.kt         |  9 +++
 .../nullCheckWithSelectorCallChain.kt.after   |  5 ++
 .../intentions/IntentionTestGenerated.java    | 24 +++++++
 12 files changed, 146 insertions(+), 39 deletions(-)
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt.after
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt.after
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt.after
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt
 create mode 100644 idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ReplaceSingleLineLetIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ReplaceSingleLineLetIntention.kt
index 35a8713b779..9234f07e092 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ReplaceSingleLineLetIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ReplaceSingleLineLetIntention.kt
@@ -58,26 +58,6 @@ class ReplaceSingleLineLetIntention : SelfTargetingOffsetIndependentIntention<Kt
         return this
     }
 
-    private fun KtDotQualifiedExpression.replaceFirstReceiver(
-            factory: KtPsiFactory,
-            newReceiver: KtExpression,
-            safeAccess: Boolean = false
-    ): KtExpression {
-        val receiver = receiverExpression
-        when (receiver) {
-            is KtDotQualifiedExpression -> {
-                receiver.replaceFirstReceiver(factory, newReceiver, safeAccess)
-            }
-            else -> {
-                if (safeAccess) {
-                    operationTokenNode.psi.replace(factory.createSafeCallNode().psi)
-                }
-                receiver.replace(newReceiver)
-            }
-        }
-        return this
-    }
-
     override fun isApplicableTo(element: KtCallExpression): Boolean {
         if (!isLetMethod(element)) return false
         val lambdaExpression = element.lambdaArguments.firstOrNull()?.getLambdaExpression() ?: return false
@@ -98,9 +78,6 @@ class ReplaceSingleLineLetIntention : SelfTargetingOffsetIndependentIntention<Kt
         return if (parameters.size == 1) parameters[0].text else "it"
     }
 
-    private fun KtDotQualifiedExpression.getLeftMostReceiverExpression(): KtExpression =
-            (receiverExpression as? KtDotQualifiedExpression)?.getLeftMostReceiverExpression() ?: receiverExpression
-
     private fun KtDotQualifiedExpression.receiverUsedAsArgument(receiverName: String): Boolean {
         if ((selectorExpression as? KtCallExpression)?.valueArguments?.firstOrNull { it.text == receiverName } != null) return true
         return (receiverExpression as? KtDotQualifiedExpression)?.receiverUsedAsArgument(receiverName) ?: false
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt
index ba823176f7c..e5586ce8ed7 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/Utils.kt
@@ -260,3 +260,28 @@ fun KtElement?.isSizeOrLength(): Boolean {
         else -> false
     }
 }
+
+
+fun KtDotQualifiedExpression.getLeftMostReceiverExpression(): KtExpression =
+        (receiverExpression as? KtDotQualifiedExpression)?.getLeftMostReceiverExpression() ?: receiverExpression
+
+fun KtDotQualifiedExpression.replaceFirstReceiver(
+        factory: KtPsiFactory,
+        newReceiver: KtExpression,
+        safeAccess: Boolean = false
+): KtExpression {
+    val receiver = receiverExpression
+    when (receiver) {
+        is KtDotQualifiedExpression -> {
+            receiver.replaceFirstReceiver(factory, newReceiver, safeAccess)
+        }
+        else -> {
+            if (safeAccess) {
+                operationTokenNode.psi.replace(factory.createSafeCallNode().psi)
+            }
+            receiver.replace(newReceiver)
+        }
+    }
+    return this
+}
+
diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfThenToElvisIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfThenToElvisIntention.kt
index afca1236b4f..39193792b4a 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfThenToElvisIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/intentions/IfThenToElvisIntention.kt
@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.inspections.IntentionBasedInspection
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingOffsetIndependentIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.*
+import org.jetbrains.kotlin.idea.intentions.getLeftMostReceiverExpression
+import org.jetbrains.kotlin.idea.intentions.replaceFirstReceiver
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 
@@ -33,31 +35,33 @@ class IfThenToElvisIntention : SelfTargetingOffsetIndependentIntention<KtIfExpre
 ) {
 
     private fun KtExpression.clausesReplaceableByElvis(firstClause: KtExpression, secondClause: KtExpression) =
-            firstClause.isNotNullExpression() && secondClause.evaluatesTo(this) &&
+            firstClause.isNotNullExpression() &&
+            (secondClause.evaluatesTo(this) || secondClause.hasFirstReceiverOf(this)) &&
             !(firstClause is KtThrowExpression && firstClause.throwsNullPointerExceptionWithNoArguments())
 
+    private fun KtExpression.checkedExpression() = when (this) {
+        is KtBinaryExpression -> expressionComparedToNull()
+        is KtIsExpression -> leftHandSide
+        else -> null
+    }
+
     override fun isApplicableTo(element: KtIfExpression): Boolean {
         val condition = element.condition as? KtOperationExpression ?: return false
         val thenClause = element.then ?: return false
         val elseClause = element.`else` ?: return false
 
-        val expression = when (condition) {
-            is KtBinaryExpression -> condition.expressionComparedToNull() ?: return false
-            is KtIsExpression -> condition.leftHandSide
-            else -> return false
-
-        }
-        if (!expression.isStableVariable()) return false
+        val checkedExpression = condition.checkedExpression() ?: return false
+        if (!checkedExpression.isStableVariable()) return false
 
         return when (condition) {
             is KtBinaryExpression -> when (condition.operationToken) {
-                KtTokens.EQEQ -> expression.clausesReplaceableByElvis(thenClause, elseClause)
-                KtTokens.EXCLEQ -> expression.clausesReplaceableByElvis(elseClause, thenClause)
+                KtTokens.EQEQ -> checkedExpression.clausesReplaceableByElvis(thenClause, elseClause)
+                KtTokens.EXCLEQ -> checkedExpression.clausesReplaceableByElvis(elseClause, thenClause)
                 else -> false
             }
             is KtIsExpression -> when (condition.isNegated) {
-                true -> expression.clausesReplaceableByElvis(thenClause, elseClause)
-                false -> expression.clausesReplaceableByElvis(elseClause, thenClause)
+                true -> checkedExpression.clausesReplaceableByElvis(thenClause, elseClause)
+                false -> checkedExpression.clausesReplaceableByElvis(elseClause, thenClause)
             }
             else -> false
         }
@@ -68,6 +72,20 @@ class IfThenToElvisIntention : SelfTargetingOffsetIndependentIntention<KtIfExpre
         return innerExpression !is KtBlockExpression && innerExpression.node.elementType != KtNodeTypes.NULL
     }
 
+    private fun KtExpression.hasFirstReceiverOf(receiver: KtExpression): Boolean {
+        val actualReceiver = (unwrapBlockOrParenthesis() as? KtDotQualifiedExpression)?.getLeftMostReceiverExpression() ?: return false
+        return actualReceiver.evaluatesTo(receiver)
+    }
+
+    private fun KtExpression.insertSafeCalls(factory: KtPsiFactory): KtExpression {
+        if (this !is KtQualifiedExpression) return this
+        if (this is KtDotQualifiedExpression) {
+            operationTokenNode.psi.replace(factory.createSafeCallNode().psi)
+        }
+        receiverExpression.insertSafeCalls(factory)
+        return this
+    }
+
     override fun applyTo(element: KtIfExpression, editor: Editor?) {
         val condition = element.condition as KtOperationExpression
 
@@ -90,11 +108,26 @@ class IfThenToElvisIntention : SelfTargetingOffsetIndependentIntention<KtIfExpre
         }
 
         val factory = KtPsiFactory(element)
-        val newExpr = factory.createExpressionByPattern("$0 ?: $1", left, right) as KtBinaryExpression
-        if (condition is KtIsExpression) {
-            newExpr.left!!.replace(factory.createExpressionByPattern("$0 as? $1", left, condition.typeReference!!))
+        val newReceiver = (condition as? KtIsExpression)?.let {
+            factory.createExpressionByPattern("$0 as? $1",
+                                              (left as? KtDotQualifiedExpression)?.getLeftMostReceiverExpression() ?: left,
+                                              it.typeReference!!)
         }
-        val elvis = KtPsiUtil.deparenthesize(element.replaced(newExpr)) as KtBinaryExpression
+        val checkedExpression = condition.checkedExpression()!!
+        val replacedLeft = if (left.evaluatesTo(checkedExpression)) {
+            if (condition is KtIsExpression) newReceiver!! else left
+        }
+        else {
+            if (condition is KtIsExpression) {
+                (left as KtDotQualifiedExpression).replaceFirstReceiver(
+                        factory, newReceiver!!, safeAccess = true)
+            }
+            else {
+                left.insertSafeCalls(factory)
+            }
+        }
+        val newExpr = element.replaced(factory.createExpressionByPattern("$0 ?: $1", replacedLeft, right))
+        val elvis = KtPsiUtil.deparenthesize(newExpr) as KtBinaryExpression
 
         if (editor != null) {
             elvis.inlineLeftSideIfApplicableWithPrompt(editor)
diff --git a/idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt b/idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt
new file mode 100644
index 00000000000..91ffecc6c3c
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt
@@ -0,0 +1,5 @@
+class My(val x: Int)
+
+fun foo(arg: Any?): Int {
+    return if (<caret>arg is My) arg.x else 42
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt.after b/idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt.after
new file mode 100644
index 00000000000..83fe36fe19f
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt.after
@@ -0,0 +1,5 @@
+class My(val x: Int)
+
+fun foo(arg: Any?): Int {
+    return (arg as? My)?.x ?: 42
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt
new file mode 100644
index 00000000000..f8aa7da8130
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt
@@ -0,0 +1,5 @@
+class My(val x: Int)
+
+fun foo(arg: My?): Int {
+    return if (<caret>arg != null) arg.x else 42
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt.after b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt.after
new file mode 100644
index 00000000000..acbd77dc6a7
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt.after
@@ -0,0 +1,5 @@
+class My(val x: Int)
+
+fun foo(arg: My?): Int {
+    return arg?.x ?: 42
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt
new file mode 100644
index 00000000000..77c1d528b02
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+
+val nullableString: String? = "abc"
+
+val foo = if (<caret>nullableString != null) {
+    nullableString.toUpperCase()
+} else {
+    ""
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt.after b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt.after
new file mode 100644
index 00000000000..d7efccae838
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt.after
@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+val nullableString: String? = "abc"
+
+val foo = nullableString?.toUpperCase() ?: ""
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt
new file mode 100644
index 00000000000..da880bec592
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt
@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+
+val nullableString: String? = "abc"
+
+val foo = if (<caret>nullableString != null) {
+    nullableString.toUpperCase().toLowerCase()
+} else {
+    ""
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt.after b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt.after
new file mode 100644
index 00000000000..8e8c151a34c
--- /dev/null
+++ b/idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt.after
@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+val nullableString: String? = "abc"
+
+val foo = nullableString?.toUpperCase()?.toLowerCase() ?: ""
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 2dac8606a02..bc7d811766a 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -1302,6 +1302,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("isCheckWithSelector.kt")
+            public void testIsCheckWithSelector() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifThenToElvis/isCheckWithSelector.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("lhsEqualsNull.kt")
             public void testLhsEqualsNull() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifThenToElvis/lhsEqualsNull.kt");
@@ -1374,6 +1380,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("nullCheckWithSelector.kt")
+            public void testNullCheckWithSelector() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelector.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("nullCheckWithSelectorCall.kt")
+            public void testNullCheckWithSelectorCall() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCall.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("nullCheckWithSelectorCallChain.kt")
+            public void testNullCheckWithSelectorCallChain() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifThenToElvis/nullCheckWithSelectorCallChain.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("otherBlockHasMoreThanOneStatement.kt")
             public void testOtherBlockHasMoreThanOneStatement() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/ifThenToElvis/otherBlockHasMoreThanOneStatement.kt");