Surround with null check : applicable also for TYPE_MISMATCH

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

@@ -194,6 +194,7 @@ class QuickFixRegistrar : QuickFixContributor {
         UNSAFE_IMPLICIT_INVOKE_CALL.registerFactory(SurroundWithNullCheckFix)
         UNSAFE_INFIX_CALL.registerFactory(SurroundWithNullCheckFix)
         ITERATOR_ON_NULLABLE.registerFactory(SurroundWithNullCheckFix.IteratorOnNullableFactory)
+        TYPE_MISMATCH.registerFactory(SurroundWithNullCheckFix.TypeMismatchFactory)
 
         UNSAFE_CALL.registerFactory(WrapWithSafeLetCallFix.UnsafeFactory)
         UNSAFE_IMPLICIT_INVOKE_CALL.registerFactory(WrapWithSafeLetCallFix.UnsafeFactory)

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

@@ -20,13 +20,16 @@ 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.diagnostics.Errors
 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.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypesAndPredicate
 import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
+import org.jetbrains.kotlin.types.typeUtil.isNullabilityMismatch
 
 class SurroundWithNullCheckFix(
         expression: KtExpression,
@@ -82,6 +85,23 @@ class SurroundWithNullCheckFix(
             return SurroundWithNullCheckFix(forExpression, nullableExpression)
         }
     }
+
+    object TypeMismatchFactory : KotlinSingleIntentionActionFactory() {
+
+        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 call = argument.getParentOfType<KtCallExpression>(true) ?: return null
+            if (call.parent !is KtBlockExpression) return null
+
+            if (!isNullabilityMismatch(expected = typeMismatch.a, actual = typeMismatch.b)) return null
+
+            if (!nullableExpression.isPredictable()) return null
+
+            return SurroundWithNullCheckFix(call, nullableExpression)
+        }
+    }
 }
 
 private fun KtExpression.isPredictable(): Boolean {

idea/testData/quickfix/surroundWithNullCheck/argumentNullable.kt

@@ -0,0 +1,7 @@
+// "Surround with null check" "true"
+
+fun foo(x: String?) {
+    bar(<caret>x)
+}
+
+fun bar(s: String) = s.hashCode()

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

@@ -0,0 +1,9 @@
+// "Surround with null check" "true"
+
+fun foo(x: String?) {
+    if (x != null) {
+        bar(x)
+    }
+}
+
+fun bar(s: String) = s.hashCode()

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

@@ -7438,6 +7438,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/surroundWithNullCheck"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }
 
+        @TestMetadata("argumentNullable.kt")
+        public void testArgumentNullable() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/argumentNullable.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("chainedUnsafeCall.kt")
         public void testChainedUnsafeCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/chainedUnsafeCall.kt");