Add better quickfix for scope functions #KT-13676 Fixed

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

@@ -229,6 +229,7 @@ class QuickFixRegistrar : QuickFixContributor {
         UNSAFE_INFIX_CALL.registerFactory(ReplaceInfixOrOperatorCallFix)
         UNSAFE_CALL.registerFactory(ReplaceInfixOrOperatorCallFix) // [] only
         UNSAFE_IMPLICIT_INVOKE_CALL.registerFactory(ReplaceInfixOrOperatorCallFix)
+        UNSAFE_CALL.registerFactory(ReplaceWithSafeCallForScopeFunctionFix)
 
         AMBIGUOUS_ANONYMOUS_TYPE_INFERRED.registerActions(SpecifyTypeExplicitlyFix())
         PROPERTY_WITH_NO_TYPE_NO_INITIALIZER.registerActions(SpecifyTypeExplicitlyFix())

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

@@ -24,8 +24,11 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 
 abstract class ReplaceCallFix(
         expression: KtQualifiedExpression,
@@ -81,6 +84,55 @@ class ReplaceWithSafeCallFix(expression: KtDotQualifiedExpression): ReplaceCallF
     }
 }
 
+class ReplaceWithSafeCallForScopeFunctionFix(expression: KtDotQualifiedExpression) : ReplaceCallFix(expression, "?.") {
+
+    override fun getText() = "Replace scope function with safe (?.) call"
+
+    companion object : KotlinSingleIntentionActionFactory() {
+        override fun createAction(diagnostic: Diagnostic): KotlinQuickFixAction<KtExpression>? {
+            val element = diagnostic.psiElement
+            val scopeFunctionLiteral = element.getStrictParentOfType<KtFunctionLiteral>() ?: return null
+            val scopeCallExpression = scopeFunctionLiteral.getStrictParentOfType<KtCallExpression>() ?: return null
+            val scopeDotQualifiedExpression = scopeCallExpression.getStrictParentOfType<KtDotQualifiedExpression>() ?: return null
+
+            val context = scopeCallExpression.analyze()
+            val scopeFunctionLiteralDescriptor = context[BindingContext.FUNCTION, scopeFunctionLiteral] ?: return null
+            val scopeFunctionKind = scopeCallExpression.scopeFunctionKind(context) ?: return null
+
+            val internalReceiver = (element.parent as? KtDotQualifiedExpression)?.receiverExpression
+            val internalReceiverDescriptor = internalReceiver.getResolvedCall(context)?.candidateDescriptor
+            val internalResolvedCall = (element.getParentOfType<KtElement>(strict = false))?.getResolvedCall(context)
+                                       ?: return null
+
+            when (scopeFunctionKind) {
+                ScopeFunctionKind.WITH_PARAMETER -> {
+                    if (internalReceiverDescriptor != scopeFunctionLiteralDescriptor.valueParameters.singleOrNull()) {
+                        return null
+                    }
+                }
+                ScopeFunctionKind.WITH_RECEIVER -> {
+                    if (internalReceiverDescriptor != scopeFunctionLiteralDescriptor.extensionReceiverParameter &&
+                        internalResolvedCall.getImplicitReceiverValue() == null) {
+                        return null
+                    }
+                }
+            }
+
+            return ReplaceWithSafeCallForScopeFunctionFix(scopeDotQualifiedExpression)
+        }
+
+        private fun KtCallExpression.scopeFunctionKind(context: BindingContext): ScopeFunctionKind? {
+            val methodName = getResolvedCall(context)?.resultingDescriptor?.fqNameUnsafe?.asString()
+            return ScopeFunctionKind.values().firstOrNull { kind -> kind.names.contains(methodName) }
+        }
+
+        private enum class ScopeFunctionKind(vararg val names: String) {
+            WITH_PARAMETER("kotlin.let", "kotlin.also"),
+            WITH_RECEIVER("kotlin.apply", "kotlin.run")
+        }
+    }
+}
+
 class ReplaceWithDotCallFix(expression: KtSafeQualifiedExpression): ReplaceCallFix(expression, "."), CleanupFix {
     override fun getText() = "Replace with dot call"
 

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/also.kt

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a.also {
+        it.<caret>length
+    }
+}

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.also {
+        it.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithFake.kt

@@ -0,0 +1,13 @@
+// "Replace scope function with safe (?.) call" "false"
+// WITH_RUNTIME
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Introduce local variable
+// ACTION: Replace with safe (?.) call
+// ACTION: Surround with null check
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type String?
+
+fun foo(a: String?, b: String?) {
+    a.apply { ->
+        b<caret>.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThis.kt

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a.apply {
+        this<caret>.length
+    }
+}

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.apply {
+        this.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThisAndArrow.kt

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a.apply { ->
+        this<caret>.length
+    }
+}

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.apply { ->
+        this.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThisLabeled.kt

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a.apply {
+        this@apply<caret>.length
+    }
+}

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.apply {
+        this@apply.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThisLabeledFake.kt

@@ -0,0 +1,14 @@
+// "Replace scope function with safe (?.) call" "false"
+// WITH_RUNTIME
+// ACTION: Introduce local variable
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Replace with safe (?.) call
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type My?
+
+class My(val prop: Int)
+
+fun My?.foo(a: String?) {
+    a.apply {
+        this@foo<caret>.prop
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithoutThis.kt

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

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.apply {
+        length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithoutThisMethodCall.kt

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a.apply {
+        <caret>toLowerCase()
+    }
+}

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.apply {
+        toLowerCase()
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/let.kt

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a.let {
+        it<caret>.length
+    }
+}

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.let {
+        it.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithParam.kt

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

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.let { s ->
+        s.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithWrongImplicitThis.kt

@@ -0,0 +1,13 @@
+// "Replace scope function with safe (?.) call" "false"
+// WITH_RUNTIME
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Introduce local variable
+// ACTION: Move lambda argument into parentheses
+// ACTION: Replace with safe (this?.) call
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type String?
+
+fun String?.foo(a: String?) {
+    a.let { s ->
+        <caret>length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithWrongParam.kt

@@ -0,0 +1,13 @@
+// "Replace scope function with safe (?.) call" "false"
+// WITH_RUNTIME
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Introduce local variable
+// ACTION: Replace with safe (?.) call
+// ACTION: Surround with null check
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type String?
+
+fun foo(a: String?, b: String?) {
+    a.let { s ->
+        b<caret>.length
+    }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/notInsideScope.kt

@@ -0,0 +1,10 @@
+// "Replace scope function with safe (?.) call" "false"
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Introduce local variable
+// ACTION: Replace with safe (?.) call
+// ACTION: Surround with null check
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type String?
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a<caret>.length
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/run.kt

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

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

@@ -0,0 +1,7 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    a?.run {
+        length
+    }
+}

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

@@ -8763,6 +8763,99 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
         }
     }
 
+    @TestMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ReplaceWithSafeCallForScopeFunction extends AbstractQuickFixTest {
+        public void testAllFilesPresentInReplaceWithSafeCallForScopeFunction() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/replaceWithSafeCallForScopeFunction"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
+        }
+
+        @TestMetadata("also.kt")
+        public void testAlso() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/also.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithFake.kt")
+        public void testApplyWithFake() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithFake.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithThis.kt")
+        public void testApplyWithThis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThis.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithThisAndArrow.kt")
+        public void testApplyWithThisAndArrow() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThisAndArrow.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithThisLabeled.kt")
+        public void testApplyWithThisLabeled() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThisLabeled.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithThisLabeledFake.kt")
+        public void testApplyWithThisLabeledFake() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithThisLabeledFake.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithoutThis.kt")
+        public void testApplyWithoutThis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithoutThis.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("applyWithoutThisMethodCall.kt")
+        public void testApplyWithoutThisMethodCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/applyWithoutThisMethodCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("let.kt")
+        public void testLet() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/let.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("letWithParam.kt")
+        public void testLetWithParam() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithParam.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("letWithWrongImplicitThis.kt")
+        public void testLetWithWrongImplicitThis() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithWrongImplicitThis.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("letWithWrongParam.kt")
+        public void testLetWithWrongParam() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithWrongParam.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("notInsideScope.kt")
+        public void testNotInsideScope() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/notInsideScope.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("run.kt")
+        public void testRun() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/run.kt");
+            doTest(fileName);
+        }
+    }
+
     @TestMetadata("idea/testData/quickfix/simplifyComparison")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)