Surround with null check: suggested for an assignment

#KT-30215 Fixed

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

@@ -104,17 +104,22 @@ class SurroundWithNullCheckFix(
         override fun createAction(diagnostic: Diagnostic): IntentionAction? {
             val typeMismatch = Errors.TYPE_MISMATCH.cast(diagnostic)
             val nullableExpression = typeMismatch.psiElement as? KtReferenceExpression ?: return null
-            val argument = nullableExpression.parent as? KtValueArgument ?: return null
+            val parent = nullableExpression.parent
-            val call = argument.getParentOfType<KtCallExpression>(true) ?: return null
+            val root = when (parent) {
-
+                is KtValueArgument -> {
-            val rootCall = call.getLastParentOfTypeInRow<KtQualifiedExpression>() ?: call
+                    val call = parent.getParentOfType<KtCallExpression>(true) ?: return null
-            if (rootCall.parent !is KtBlockExpression) return null
+                    call.getLastParentOfTypeInRow<KtQualifiedExpression>() ?: call
-
+                }
+                is KtBinaryExpression -> {
+                    if (parent.right != nullableExpression) return null
+                    parent
+                }
+                else -> return null
+            }
+            if (root.parent !is KtBlockExpression) return null
             if (!isNullabilityMismatch(expected = typeMismatch.a, actual = typeMismatch.b)) return null
-
             if (!nullableExpression.isStableSimpleExpression()) return null
-
+            return SurroundWithNullCheckFix(root, nullableExpression)
-            return SurroundWithNullCheckFix(rootCall, nullableExpression)
         }
     }
 }

idea/testData/quickfix/surroundWithNullCheck/assignment.kt

@@ -0,0 +1,6 @@
+// "Surround with null check" "true"
+
+fun foo(s: String?) {
+    var ss: String = ""
+    ss = <caret>s
+}

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

@@ -0,0 +1,8 @@
+// "Surround with null check" "true"
+
+fun foo(s: String?) {
+    var ss: String = ""
+    if (s != null) {
+        ss = s
+    }
+}

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

@@ -11526,6 +11526,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/surroundWithNullCheck/argumentNullable.kt");
         }
 
+        @TestMetadata("assignment.kt")
+        public void testAssignment() throws Exception {
+            runTest("idea/testData/quickfix/surroundWithNullCheck/assignment.kt");
+        }
+
         @TestMetadata("chainedUnsafeCall.kt")
         public void testChainedUnsafeCall() throws Exception {
             runTest("idea/testData/quickfix/surroundWithNullCheck/chainedUnsafeCall.kt");