Add quick-fix "Replace with safe call & elvis" #KT-17815 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/ReplaceCallFix.kt

@@ -32,7 +32,8 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 
 abstract class ReplaceCallFix(
         expression: KtQualifiedExpression,
-        private val operation: String
+        private val operation: String,
+        private val notNullNeeded: Boolean = false
 ) : KotlinQuickFixAction<KtQualifiedExpression>(expression) {
 
     override fun getFamilyName() = text
@@ -44,25 +45,39 @@ abstract class ReplaceCallFix(
 
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
         val element = element ?: return
-        val newExpression = KtPsiFactory(element).createExpressionByPattern("$0$operation$1",
+        val elvis = elvisOrEmpty(notNullNeeded)
+        val newExpression = KtPsiFactory(element).createExpressionByPattern("$0$operation$1$elvis",
                                                                             element.receiverExpression, element.selectorExpression!!)
-        element.replace(newExpression)
+        val replacement = element.replace(newExpression)
+        if (notNullNeeded) {
+            replacement.moveCaretToEnd(editor, project)
+        }
     }
 }
 
-class ReplaceImplicitReceiverCallFix(expression: KtExpression) : KotlinQuickFixAction<KtExpression>(expression) {
+class ReplaceImplicitReceiverCallFix(
+        expression: KtExpression,
+        private val notNullNeeded: Boolean
+) : KotlinQuickFixAction<KtExpression>(expression) {
     override fun getFamilyName() = text
 
     override fun getText() = "Replace with safe (this?.) call"
 
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
         val element = element ?: return
-        val newExpression = KtPsiFactory(element).createExpressionByPattern("this?.$0", element)
-        element.replace(newExpression)
+        val elvis = elvisOrEmpty(notNullNeeded)
+        val newExpression = KtPsiFactory(element).createExpressionByPattern("this?.$0$elvis", element)
+        val replacement = element.replace(newExpression)
+        if (notNullNeeded) {
+            replacement.moveCaretToEnd(editor, project)
+        }
     }
 }
 
-class ReplaceWithSafeCallFix(expression: KtDotQualifiedExpression): ReplaceCallFix(expression, "?.") {
+class ReplaceWithSafeCallFix(
+        expression: KtDotQualifiedExpression,
+        notNullNeeded: Boolean
+) : ReplaceCallFix(expression, "?.", notNullNeeded) {
 
     override fun getText() = "Replace with safe (?.) call"
 
@@ -71,12 +86,13 @@ class ReplaceWithSafeCallFix(expression: KtDotQualifiedExpression): ReplaceCallF
             val psiElement = diagnostic.psiElement
             val qualifiedExpression = psiElement.parent as? KtDotQualifiedExpression
             if (qualifiedExpression != null) {
-                return ReplaceWithSafeCallFix(qualifiedExpression)
-            } else {
+                return ReplaceWithSafeCallFix(qualifiedExpression, qualifiedExpression.shouldHaveNotNullType())
+            }
+            else {
                 psiElement as? KtNameReferenceExpression ?: return null
                 if (psiElement.getResolvedCall(psiElement.analyze())?.getImplicitReceiverValue() != null) {
                     val expressionToReplace: KtExpression = psiElement.parent as? KtCallExpression ?: psiElement
-                    return ReplaceImplicitReceiverCallFix(expressionToReplace)
+                    return ReplaceImplicitReceiverCallFix(expressionToReplace, expressionToReplace.shouldHaveNotNullType())
                 }
                 return null
             }
@@ -84,7 +100,10 @@ class ReplaceWithSafeCallFix(expression: KtDotQualifiedExpression): ReplaceCallF
     }
 }
 
-class ReplaceWithSafeCallForScopeFunctionFix(expression: KtDotQualifiedExpression) : ReplaceCallFix(expression, "?.") {
+class ReplaceWithSafeCallForScopeFunctionFix(
+        expression: KtDotQualifiedExpression,
+        notNullNeeded: Boolean
+) : ReplaceCallFix(expression, "?.", notNullNeeded) {
 
     override fun getText() = "Replace scope function with safe (?.) call"
 
@@ -118,7 +137,8 @@ class ReplaceWithSafeCallForScopeFunctionFix(expression: KtDotQualifiedExpressio
                 }
             }
 
-            return ReplaceWithSafeCallForScopeFunctionFix(scopeDotQualifiedExpression)
+            return ReplaceWithSafeCallForScopeFunctionFix(
+                    scopeDotQualifiedExpression, scopeDotQualifiedExpression.shouldHaveNotNullType())
         }
 
         private fun KtCallExpression.scopeFunctionKind(context: BindingContext): ScopeFunctionKind? {
@@ -133,7 +153,7 @@ class ReplaceWithSafeCallForScopeFunctionFix(expression: KtDotQualifiedExpressio
     }
 }
 
-class ReplaceWithDotCallFix(expression: KtSafeQualifiedExpression): ReplaceCallFix(expression, "."), CleanupFix {
+class ReplaceWithDotCallFix(expression: KtSafeQualifiedExpression) : ReplaceCallFix(expression, "."), CleanupFix {
     override fun getText() = "Replace with dot call"
 
     companion object : KotlinSingleIntentionActionFactory() {
@@ -143,3 +163,4 @@ class ReplaceWithDotCallFix(expression: KtSafeQualifiedExpression): ReplaceCallF
         }
     }
 }
+

idea/src/org/jetbrains/kotlin/idea/quickfix/ReplaceCallFixUtils.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.quickfix
+
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.psi.KtBinaryExpression
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
+
+fun elvisOrEmpty(notNullNeeded: Boolean): String = if (notNullNeeded) "?:" else ""
+
+fun KtExpression.shouldHaveNotNullType(): Boolean {
+    val parent = parent
+    val type = when (parent) {
+                   is KtBinaryExpression -> parent.left?.let { it.getType(it.analyze()) }
+                   is KtProperty -> parent.typeReference?.let { it.analyze()[BindingContext.TYPE, it] }
+                   else -> null
+               } ?: return false
+    return !type.isMarkedNullable
+}
+
+fun PsiElement.moveCaretToEnd(editor: Editor?, project: Project) {
+    editor?.run {
+        PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(document)
+        val endOffset = if (text.endsWith(")")) endOffset - 1 else endOffset
+        document.insertString(endOffset, " ")
+        caretModel.moveToOffset(endOffset + 1)
+    }
+}

idea/src/org/jetbrains/kotlin/idea/quickfix/ReplaceInfixOrOperatorCallFix.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.quickfix
 import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.OperatorToFunctionIntention
@@ -29,7 +30,10 @@ import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 
-class ReplaceInfixOrOperatorCallFix(element: KtExpression) : KotlinQuickFixAction<KtExpression>(element) {
+class ReplaceInfixOrOperatorCallFix(
+        element: KtExpression,
+        private val notNullNeeded: Boolean
+) : KotlinQuickFixAction<KtExpression>(element) {
 
     override fun getText() = "Replace with safe (?.) call"
 
@@ -38,6 +42,8 @@ class ReplaceInfixOrOperatorCallFix(element: KtExpression) : KotlinQuickFixActio
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
         val element = element ?: return
         val psiFactory = KtPsiFactory(file)
+        val elvis = elvisOrEmpty(notNullNeeded)
+        var replacement: PsiElement? = null
         when (element) {
             is KtArrayAccessExpression -> {
                 val assignment = element.getAssignmentByLHS()
@@ -51,31 +57,34 @@ class ReplaceInfixOrOperatorCallFix(element: KtExpression) : KotlinQuickFixActio
                 }
                 else {
                     val newExpression = psiFactory.createExpressionByPattern(
-                            "$0?.get($1)", arrayExpression, element.indexExpressions.joinToString(", ") { it.text })
-                    element.replace(newExpression)
+                            "$0?.get($1)$elvis", arrayExpression, element.indexExpressions.joinToString(", ") { it.text })
+                    replacement = element.replace(newExpression)
                 }
             }
             is KtCallExpression -> {
                 val newExpression = psiFactory.createExpressionByPattern(
-                        "$0?.invoke($1)", element.calleeExpression ?: return, element.valueArguments.joinToString(", ") { it.text })
-                element.replace(newExpression)
+                        "$0?.invoke($1)$elvis", element.calleeExpression ?: return, element.valueArguments.joinToString(", ") { it.text })
+                replacement = element.replace(newExpression)
             }
             is KtBinaryExpression -> {
                 if (element.operationToken == KtTokens.IDENTIFIER) {
                     val newExpression = psiFactory.createExpressionByPattern(
-                            "$0?.$1($2)", element.left ?: return, element.operationReference, element.right ?: return)
-                    element.replace(newExpression)
+                            "$0?.$1($2)$elvis", element.left ?: return, element.operationReference, element.right ?: return)
+                    replacement = element.replace(newExpression)
                 }
                 else {
                     val nameExpression = OperatorToFunctionIntention.convert(element).second
                     val callExpression = nameExpression.parent as KtCallExpression
                     val qualifiedExpression = callExpression.parent as KtDotQualifiedExpression
                     val safeExpression = psiFactory.createExpressionByPattern(
-                            "$0?.$1", qualifiedExpression.receiverExpression, callExpression)
-                    qualifiedExpression.replace(safeExpression)
+                            "$0?.$1$elvis", qualifiedExpression.receiverExpression, callExpression)
+                    replacement = qualifiedExpression.replace(safeExpression)
                 }
             }
         }
+        if (notNullNeeded) {
+            replacement?.moveCaretToEnd(editor, project)
+        }
     }
 
     override fun startInWriteAction() = true
@@ -85,22 +94,21 @@ class ReplaceInfixOrOperatorCallFix(element: KtExpression) : KotlinQuickFixActio
             val expression = diagnostic.psiElement
             if (expression is KtArrayAccessExpression) {
                 if (expression.arrayExpression == null) return null
-                return ReplaceInfixOrOperatorCallFix(expression)
+                return ReplaceInfixOrOperatorCallFix(expression, expression.shouldHaveNotNullType())
             }
             val parent = expression.parent
             return when (parent) {
                 is KtBinaryExpression -> {
                     if (parent.left == null || parent.right == null) null
-                    else {
-                        if (parent.operationToken in OperatorConventions.COMPARISON_OPERATIONS) null
-                        else ReplaceInfixOrOperatorCallFix(parent)
-                    }
+                    else if (parent.operationToken == KtTokens.EQ) null
+                    else if (parent.operationToken in OperatorConventions.COMPARISON_OPERATIONS) null
+                    else ReplaceInfixOrOperatorCallFix(parent, parent.shouldHaveNotNullType())
                 }
                 is KtCallExpression -> {
                     if (parent.calleeExpression == null) null
                     else if (parent.parent is KtQualifiedExpression) null
                     else if (parent.getResolvedCall(parent.analyze())?.getImplicitReceiverValue() != null) null
-                    else ReplaceInfixOrOperatorCallFix(parent)
+                    else ReplaceInfixOrOperatorCallFix(parent, parent.shouldHaveNotNullType())
                 }
                 else -> null
             }

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentArray.kt

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun foo(array: Array<String>?) {
+    var s = ""
+    s = array[0]<caret>
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentArray.kt.after

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun foo(array: Array<String>?) {
+    var s = ""
+    s = array?.get(0) ?: <caret>
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentBinaryOperator.kt

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun foo(bar: Int?) {
+    var i: Int = 1
+    i = bar +<caret> 1
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentBinaryOperator.kt.after

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun foo(bar: Int?) {
+    var i: Int = 1
+    i = bar?.plus(1) ?: <caret>
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentCallExpression.kt

@@ -0,0 +1,8 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun bar() {
+    val fff: (() -> Int)? = { 1 }
+    var i: Int = 1
+    i = fff<caret>()
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentCallExpression.kt.after

@@ -0,0 +1,8 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun bar() {
+    val fff: (() -> Int)? = { 1 }
+    var i: Int = 1
+    i = fff?.invoke() ?: <caret>
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentList.kt

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun foo(list: List<String>?) {
+    var s = ""
+    s = list[0]<caret>
+}

idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentList.kt.after

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+
+fun foo(list: List<String>?) {
+    var s = ""
+    s = list?.get(0) ?: <caret>
+}

idea/testData/quickfix/replaceWithSafeCall/assignment.kt

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+var i = 0
+
+fun foo(s: String?) {
+    i = s<caret>.length
+}

idea/testData/quickfix/replaceWithSafeCall/assignment.kt.after

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+var i = 0
+
+fun foo(s: String?) {
+    i = s?.length ?: <caret>
+}

idea/testData/quickfix/replaceWithSafeCall/assignmentFromImplicitParameter.kt

@@ -0,0 +1,9 @@
+// "Replace with safe (this?.) call" "true"
+// WITH_RUNTIME
+var i = 0
+
+fun foo(a: String?) {
+    a.run {
+        i = <caret>length
+    }
+}

idea/testData/quickfix/replaceWithSafeCall/assignmentFromImplicitParameter.kt.after

@@ -0,0 +1,9 @@
+// "Replace with safe (this?.) call" "true"
+// WITH_RUNTIME
+var i = 0
+
+fun foo(a: String?) {
+    a.run {
+        i = this?.length ?: <caret>
+    }
+}

idea/testData/quickfix/replaceWithSafeCall/assignmentToNullable.kt

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+var i: Int? = 0
+
+fun foo(s: String?) {
+    i = s<caret>.length
+}

idea/testData/quickfix/replaceWithSafeCall/assignmentToNullable.kt.after

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+var i: Int? = 0
+
+fun foo(s: String?) {
+    i = s?.length
+}

idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt

@@ -0,0 +1,5 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T(s: String?) {
+    var i: Int = s<caret>.length
+}

idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt.after

@@ -0,0 +1,5 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+class T(s: String?) {
+    var i: Int = s?.length ?: <caret>
+}

idea/testData/quickfix/replaceWithSafeCall/expression.kt

@@ -0,0 +1,5 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(s: String?) {
+    1 + s<caret>.length
+}

idea/testData/quickfix/replaceWithSafeCall/expression.kt.after

@@ -0,0 +1,5 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(s: String?) {
+    1 + (s?.length ?: <caret>)
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/assignment.kt

@@ -0,0 +1,9 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+var i = 0
+
+fun foo(a: String?) {
+    i = a.run {
+        length<caret>
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/assignment.kt.after

@@ -0,0 +1,9 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+var i = 0
+
+fun foo(a: String?) {
+    i = a?.run {
+        length
+    } ?: <caret>
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -8624,6 +8624,30 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("assignmentArray.kt")
+        public void testAssignmentArray() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentArray.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("assignmentBinaryOperator.kt")
+        public void testAssignmentBinaryOperator() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentBinaryOperator.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("assignmentCallExpression.kt")
+        public void testAssignmentCallExpression() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentCallExpression.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("assignmentList.kt")
+        public void testAssignmentList() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceInfixOrOperatorCall/assignmentList.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("binaryOperator.kt")
         public void testBinaryOperator() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceInfixOrOperatorCall/binaryOperator.kt");
@@ -8753,6 +8777,36 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("assignment.kt")
+        public void testAssignment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCall/assignment.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("assignmentFromImplicitParameter.kt")
+        public void testAssignmentFromImplicitParameter() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCall/assignmentFromImplicitParameter.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("assignmentToNullable.kt")
+        public void testAssignmentToNullable() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCall/assignmentToNullable.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("assignmentToProperty.kt")
+        public void testAssignmentToProperty() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("expression.kt")
+        public void testExpression() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCall/expression.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("extFunction.kt")
         public void testExtFunction() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCall/extFunction.kt");
@@ -8846,6 +8900,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("assignment.kt")
+        public void testAssignment() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/assignment.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("let.kt")
         public void testLet() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/let.kt");