KT-20010 'Replace safe access expression with 'if' expression' IDEA Kotlin plugin intention may failed (#1288)

* 'Replace safe access expression with 'if' expression' IDEA Kotlin plugin intention may failed #KT-20010 Fixed

* #KT-20010 Fixed

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

@@ -20,11 +20,14 @@ 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.core.replaced
 import org.jetbrains.kotlin.idea.intentions.SelfTargetingRangeIntention
 import org.jetbrains.kotlin.idea.intentions.branchedTransformations.convertToIfNotNullExpression
 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.*
+import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsStatement
 
 class SafeAccessToIfThenIntention : SelfTargetingRangeIntention<KtSafeQualifiedExpression>(KtSafeQualifiedExpression::class.java, "Replace safe access expression with 'if' expression"), LowPriorityAction {
@@ -43,12 +46,25 @@ class SafeAccessToIfThenIntention : SelfTargetingRangeIntention<KtSafeQualifiedE
         val dotQualified = psiFactory.createExpressionByPattern("$0.$1", receiver, selector)
 
         val elseClause = if (element.isUsedAsStatement(element.analyze())) null else psiFactory.createExpression("null")
-        val ifExpression = element.convertToIfNotNullExpression(receiver, dotQualified, elseClause)
+        var ifExpression = element.convertToIfNotNullExpression(receiver, dotQualified, elseClause)
+
+        var isAssignment = false
+        val binaryExpression = (ifExpression.parent as? KtParenthesizedExpression)?.parent as? KtBinaryExpression
+        val right = binaryExpression?.right
+        if (right != null && binaryExpression.operationToken == KtTokens.EQ) {
+            val replaced = binaryExpression.replaced(psiFactory.createExpressionByPattern("$0 = $1", ifExpression.text, right))
+            ifExpression = replaced.findDescendantOfType()!!
+            isAssignment = true
+        }
 
         if (!receiverIsStable) {
-            val valueToExtract = (ifExpression.then as KtDotQualifiedExpression).receiverExpression
-            ifExpression.introduceValueForCondition(valueToExtract, editor)
+            val valueToExtract = if (isAssignment)
+                ((ifExpression.then as? KtBinaryExpression)?.left as? KtDotQualifiedExpression)?.receiverExpression
+            else
+                (ifExpression.then as? KtDotQualifiedExpression)?.receiverExpression
+
+            if (valueToExtract != null) ifExpression.introduceValueForCondition(valueToExtract, editor)
         }
+
     }
 }
-

idea/testData/intentions/branched/safeAccessToIfThen/assignment.kt

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    foo?<caret>.bar = Bar()
+}

idea/testData/intentions/branched/safeAccessToIfThen/assignment.kt.after

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    if (foo != null) foo.bar = Bar()
+}

idea/testData/intentions/branched/safeAccessToIfThen/assignment2.kt

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    foo?<caret>.bar?.baz = 2
+}

idea/testData/intentions/branched/safeAccessToIfThen/assignment2.kt.after

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    (if (foo != null) foo.bar else null)?.baz = 2
+}

idea/testData/intentions/branched/safeAccessToIfThen/assignment3.kt

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    foo?.bar?<caret>.baz = 2
+}

idea/testData/intentions/branched/safeAccessToIfThen/assignment3.kt.after

@@ -0,0 +1,10 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    val bar = foo?.bar
+    if (bar != null) bar.baz = 2
+}

idea/testData/intentions/branched/safeAccessToIfThen/equality.kt

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    foo?<caret>.bar == null
+}

idea/testData/intentions/branched/safeAccessToIfThen/equality.kt.after

@@ -0,0 +1,9 @@
+class Foo {
+    var bar: Bar? = null
+}
+class Bar {
+    var baz = 1
+}
+fun test(foo: Foo?) {
+    (if (foo != null) foo.bar else null) == null
+}

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

@@ -2344,6 +2344,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/branched/safeAccessToIfThen"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
             }
 
+            @TestMetadata("assignment.kt")
+            public void testAssignment() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/safeAccessToIfThen/assignment.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("assignment2.kt")
+            public void testAssignment2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/safeAccessToIfThen/assignment2.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("assignment3.kt")
+            public void testAssignment3() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/safeAccessToIfThen/assignment3.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("binaryExpressionLhs.kt")
             public void testBinaryExpressionLhs() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/safeAccessToIfThen/binaryExpressionLhs.kt");
@@ -2368,6 +2386,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("equality.kt")
+            public void testEquality() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/safeAccessToIfThen/equality.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("localValAsReceiver.kt")
             public void testLocalValAsReceiver() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/branched/safeAccessToIfThen/localValAsReceiver.kt");