Surround with null check for unsafe things: find first parent statement to surround (which is not used as expression) #KT-13958 Fixed

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

@@ -29,6 +29,8 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getLastParentOfTypeInRow
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypesAndPredicate
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.bindingContextUtil.isUsedAsExpression
 import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.types.typeUtil.isNullabilityMismatch
@@ -52,26 +54,26 @@ class SurroundWithNullCheckFix(
 
         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 expressionParent = element.getParentOfType<KtExpression>(strict = element is KtOperationReferenceExpression) ?: return null
+            val context = expressionParent.analyze()
 
             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 nullableExpression =
+                    when (parent) {
+                        is KtDotQualifiedExpression -> parent.receiverExpression
+                        is KtBinaryExpression -> parent.left
+                        is KtCallExpression -> parent.calleeExpression
+                        else -> return null
+                    } as? KtReferenceExpression ?: return null
 
-            if (!nullableExpression.isStable()) return null
+            if (!nullableExpression.isStable(context)) return null
 
-            return SurroundWithNullCheckFix(expression, nullableExpression)
+            val expressionTarget = expressionParent.getParentOfTypesAndPredicate(strict = false, parentClasses = KtExpression::class.java) {
+                !it.isUsedAsExpression(context)
+            } ?: return null
+            if (expressionTarget is KtDeclaration) return null
+
+            return SurroundWithNullCheckFix(expressionTarget, nullableExpression)
         }
     }
 
@@ -108,8 +110,7 @@ class SurroundWithNullCheckFix(
     }
 }
 
-private fun KtExpression.isStable(): Boolean {
-    val context = this.analyze()
+private fun KtExpression.isStable(context: BindingContext = this.analyze()): Boolean {
     val nullableType = this.getType(context) ?: return false
     val containingDescriptor = this.getResolutionScope(context, this.getResolutionFacade()).ownerDescriptor
     return DataFlowValueFactory.createDataFlowValue(this, nullableType, context, containingDescriptor).isStable

idea/testData/quickfix/surroundWithNullCheck/nullableInCondition.kt

@@ -0,0 +1,8 @@
+// "Surround with null check" "true"
+// WITH_RUNTIME
+
+fun foz(arg: String?) {
+    if (arg<caret>.isNotEmpty()) {
+
+    }
+}

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

@@ -0,0 +1,10 @@
+// "Surround with null check" "true"
+// WITH_RUNTIME
+
+fun foz(arg: String?) {
+    if (arg != null) {
+        if (arg.isNotEmpty()) {
+
+        }
+    }
+}

idea/testData/quickfix/surroundWithNullCheck/unsafeCallInsideAnonymous.kt

@@ -1,11 +1,5 @@
-// "Surround with null check" "false"
+// "Surround with null check" "true"
 // 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/unsafeCallInsideAnonymous.kt.after

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

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

@@ -8405,6 +8405,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("nullableInCondition.kt")
+        public void testNullableInCondition() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/nullableInCondition.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("objectQualifier.kt")
         public void testObjectQualifier() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/objectQualifier.kt");