KT-11104 : quick fix "surround with null check"

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/util/callUtil.kt

@@ -31,6 +31,7 @@ import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.CallTransformer
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
 import org.jetbrains.kotlin.resolve.calls.model.*
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.sure
 
@@ -200,6 +201,16 @@ fun KtExpression.getFunctionResolvedCallWithAssert(context: BindingContext): Res
     return resolvedCall as ResolvedCall<out FunctionDescriptor>
 }
 
+fun KtExpression.getType(context: BindingContext): KotlinType? {
+    val type = context.getType(this)
+    if (type != null) return type
+    val resolvedCall = this.getResolvedCall(context)
+    if (resolvedCall is VariableAsFunctionResolvedCall) {
+        return resolvedCall.variableCall.resultingDescriptor.type
+    }
+    return null
+}
+
 fun Call.isSafeCall(): Boolean {
     if (this is CallTransformer.CallForImplicitInvoke) {
         //implicit safe 'invoke'

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

@@ -190,6 +190,10 @@ class QuickFixRegistrar : QuickFixContributor {
         UNNECESSARY_SAFE_CALL.registerFactory(ReplaceWithDotCallFix)
         UNSAFE_CALL.registerFactory(ReplaceWithSafeCallFix)
 
+        UNSAFE_CALL.registerFactory(SurroundWithNullCheckFix)
+        UNSAFE_IMPLICIT_INVOKE_CALL.registerFactory(SurroundWithNullCheckFix)
+        UNSAFE_INFIX_CALL.registerFactory(SurroundWithNullCheckFix)
+
         UNSAFE_CALL.registerFactory(AddExclExclCallFix)
         UNNECESSARY_NOT_NULL_ASSERTION.registerFactory(RemoveExclExclCallFix)
         UNSAFE_INFIX_CALL.registerFactory(ReplaceInfixCallFix)

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

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2010-2016 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.codeInsight.intention.IntentionAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.diagnostics.Diagnostic
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.util.getResolutionScope
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypesAndPredicate
+import org.jetbrains.kotlin.resolve.calls.callUtil.getType
+import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
+
+class SurroundWithNullCheckFix(
+        expression: KtExpression,
+        val nullableExpression: KtExpression
+) : KotlinQuickFixAction<KtExpression>(expression) {
+
+    override fun getFamilyName() = text
+
+    override fun getText() = "Surround with null check"
+
+    override fun invoke(project: Project, editor: Editor?, file: KtFile) {
+        val factory = KtPsiFactory(element)
+        val surrounded = factory.createExpressionByPattern("if ($0 != null) { $1 }", nullableExpression, element)
+        element.replace(surrounded)
+    }
+
+    companion object : KotlinSingleIntentionActionFactory() {
+
+        override fun createAction(diagnostic: Diagnostic): IntentionAction? {
+            val element = diagnostic.psiElement
+            val expression = element.getParentOfTypesAndPredicate(false, KtExpression::class.java) {
+                it is KtDeclaration ||
+                it.parent is KtBlockExpression ||
+                it.parent?.parent is KtIfExpression ||
+                it.parent?.parent is KtWhenExpression ||
+                it.parent?.parent is KtLoopExpression
+            } ?: return null
+            if (expression is KtDeclaration) return null
+
+            val parent = element.parent
+            val nullableExpression = when (parent) {
+                is KtDotQualifiedExpression -> parent.receiverExpression
+                is KtBinaryExpression -> parent.left
+                is KtCallExpression -> parent.calleeExpression
+                else -> return null
+            } as? KtReferenceExpression ?: return null
+
+            val context = expression.analyze()
+            val nullableType = nullableExpression.getType(context) ?: return null
+            val containingDescriptor = nullableExpression.getResolutionScope(context, expression.getResolutionFacade()).ownerDescriptor
+            val dataFlowValue = DataFlowValueFactory.createDataFlowValue(nullableExpression, nullableType, context, containingDescriptor)
+            if (!dataFlowValue.isPredictable) return null
+
+            return SurroundWithNullCheckFix(expression, nullableExpression)
+        }
+    }
+
+}

idea/testData/quickfix/surroundWithNullCheck/chainedUnsafeCall.kt

@@ -0,0 +1,9 @@
+// "Surround with null check" "false"
+// 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 Int?
+
+fun foo(arg: Int?) {
+    arg?.hashCode()<caret>.toLong()
+}

idea/testData/quickfix/surroundWithNullCheck/expressionUnsafeCall.kt

@@ -0,0 +1,8 @@
+// "Surround with null check" "false"
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Convert to block body
+// ACTION: Introduce local variable
+// ACTION: Replace with safe (?.) call
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type Int?
+
+fun foo(arg: Int?) = arg<caret>.hashCode()

idea/testData/quickfix/surroundWithNullCheck/infixUnsafeCall.kt

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+infix fun Int.op(arg: Int) = this
+
+fun foo(arg: Int?) {
+    arg <caret>op 42
+}

idea/testData/quickfix/surroundWithNullCheck/infixUnsafeCall.kt.after

@@ -0,0 +1,9 @@
+// "Surround with null check" "true"
+
+infix fun Int.op(arg: Int) = this
+
+fun foo(arg: Int?) {
+    if (arg != null) {
+        arg op 42
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/invokeFuncUnsafe.kt

@@ -0,0 +1,5 @@
+// "Surround with null check" "true"
+
+fun foo(exec: (() -> Unit)?) {
+    <caret>exec()
+}

idea/testData/quickfix/surroundWithNullCheck/invokeFuncUnsafe.kt.after

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(exec: (() -> Unit)?) {
+    if (exec != null) {
+        exec()
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/invokeUnsafe.kt

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+operator fun Int.invoke() = this
+
+fun foo(arg: Int?) {
+    <caret>arg()
+}

idea/testData/quickfix/surroundWithNullCheck/invokeUnsafe.kt.after

@@ -0,0 +1,9 @@
+// "Surround with null check" "true"
+
+operator fun Int.invoke() = this
+
+fun foo(arg: Int?) {
+    if (arg != null) {
+        arg()
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/simpleUnsafeCall.kt

@@ -0,0 +1,5 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?) {
+    arg<caret>.hashCode()
+}

idea/testData/quickfix/surroundWithNullCheck/simpleUnsafeCall.kt.after

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?) {
+    if (arg != null) {
+        arg.hashCode()
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInBinary.kt

@@ -0,0 +1,5 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?) {
+    42 + arg<caret>.hashCode() - 13
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInBinary.kt.after

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?) {
+    if (arg != null) {
+        42 + arg.hashCode() - 13
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInDeclaration.kt

@@ -0,0 +1,8 @@
+// "Surround with null check" "false"
+// 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 String?
+
+fun foo(s: String?) {
+    val x = s<caret>.hashCode()
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInGetter.kt

@@ -0,0 +1,11 @@
+// "Surround with null check" "false"
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Convert to block body
+// ACTION: Introduce local variable
+// ACTION: Replace with safe (?.) call
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type String?
+
+class My(val x: String?) {
+    val y: Int
+        get() = x<caret>.hashCode()
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideAnonymous.kt

@@ -0,0 +1,12 @@
+// "Surround with null check" "false"
+// WITH_RUNTIME
+// ACTION: Add 'block =' to argument
+// ACTION: Add non-null asserted (!!) call
+// ACTION: Convert to block body
+// ACTION: Introduce local variable
+// ACTION: Replace with safe (?.) call
+// ERROR: Only safe (?.) or non-null asserted (!!.) calls are allowed on a nullable receiver of type Int?
+
+fun foo(arg: Int?) {
+    run(fun() = arg<caret>.hashCode())
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideIf.kt

@@ -0,0 +1,5 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?, flag: Boolean) {
+    if (flag) arg<caret>.hashCode()
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideIf.kt.after

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?, flag: Boolean) {
+    if (flag) if (arg != null) {
+        arg.hashCode()
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideWhen.kt

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?, flag: Boolean) {
+    when (flag) {
+        true -> arg<caret>.hashCode()
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideWhen.kt.after

@@ -0,0 +1,9 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?, flag: Boolean) {
+    when (flag) {
+        true -> if (arg != null) {
+            arg.hashCode()
+        }
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideWhile.kt

@@ -0,0 +1,5 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?, flag: Boolean) {
+    while (flag) arg<caret>.hashCode()
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideWhile.kt.after

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(arg: Int?, flag: Boolean) {
+    while (flag) if (arg != null) {
+        arg.hashCode()
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unstableValue.kt

@@ -0,0 +1,12 @@
+// "Surround with null check" "false"
+// 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 Int?
+
+class My(var x: Int?) {
+
+    fun foo() {
+        x<caret>.hashCode()
+    }
+}

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

@@ -7430,6 +7430,99 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
         }
     }
 
+    @TestMetadata("idea/testData/quickfix/surroundWithNullCheck")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class SurroundWithNullCheck extends AbstractQuickFixTest {
+        public void testAllFilesPresentInSurroundWithNullCheck() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/surroundWithNullCheck"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
+        }
+
+        @TestMetadata("chainedUnsafeCall.kt")
+        public void testChainedUnsafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/chainedUnsafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("expressionUnsafeCall.kt")
+        public void testExpressionUnsafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/expressionUnsafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("infixUnsafeCall.kt")
+        public void testInfixUnsafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/infixUnsafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("invokeFuncUnsafe.kt")
+        public void testInvokeFuncUnsafe() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/invokeFuncUnsafe.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("invokeUnsafe.kt")
+        public void testInvokeUnsafe() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/invokeUnsafe.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("simpleUnsafeCall.kt")
+        public void testSimpleUnsafeCall() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/simpleUnsafeCall.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInBinary.kt")
+        public void testUnsafeCallInBinary() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInBinary.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInDeclaration.kt")
+        public void testUnsafeCallInDeclaration() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInDeclaration.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInGetter.kt")
+        public void testUnsafeCallInGetter() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInGetter.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInsideAnonymous.kt")
+        public void testUnsafeCallInsideAnonymous() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideAnonymous.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInsideIf.kt")
+        public void testUnsafeCallInsideIf() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideIf.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInsideWhen.kt")
+        public void testUnsafeCallInsideWhen() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideWhen.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unsafeCallInsideWhile.kt")
+        public void testUnsafeCallInsideWhile() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideWhile.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("unstableValue.kt")
+        public void testUnstableValue() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/unstableValue.kt");
+            doTest(fileName);
+        }
+    }
+
     @TestMetadata("idea/testData/quickfix/toString")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)