Surround with null check : apply for ITERATOR_ON_NULLABLE

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

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

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

@@ -69,6 +69,19 @@ class SurroundWithNullCheckFix(
             return SurroundWithNullCheckFix(expression, nullableExpression)
         }
     }
+
+    object IteratorOnNullableFactory : KotlinSingleIntentionActionFactory() {
+
+        override fun createAction(diagnostic: Diagnostic): IntentionAction? {
+            val nullableExpression = diagnostic.psiElement as? KtReferenceExpression ?: return null
+            val forExpression = nullableExpression.parent?.parent as? KtForExpression ?: return null
+            if (forExpression.parent !is KtBlockExpression) return null
+
+            if (!nullableExpression.isPredictable()) return null
+
+            return SurroundWithNullCheckFix(forExpression, nullableExpression)
+        }
+    }
 }
 
 private fun KtExpression.isPredictable(): Boolean {

idea/testData/quickfix/expressions/fixNullableWithExclExclAbsentWithBadIterator.kt

@@ -1,5 +1,6 @@
 // "Add non-null asserted (!!) call" "false"
 // ACTION: Replace with a 'forEach' function call
+// ACTION: Surround with null check
 // ERROR: Not nullable value required to call an 'iterator()' method on for-loop range
 
 class Some {

idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt

@@ -0,0 +1,6 @@
+// "Surround with null check" "true"
+// WITH_RUNTIME
+
+fun foo(list: List<String>?) {
+    for (element in <caret>list) {}
+}

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

@@ -0,0 +1,8 @@
+// "Surround with null check" "true"
+// WITH_RUNTIME
+
+fun foo(list: List<String>?) {
+    if (list != null) {
+        for (element in list) {}
+    }
+}

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

@@ -7468,6 +7468,12 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             doTest(fileName);
         }
 
+        @TestMetadata("iteratorUnsafe.kt")
+        public void testIteratorUnsafe() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("simpleUnsafeCall.kt")
         public void testSimpleUnsafeCall() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/quickfix/surroundWithNullCheck/simpleUnsafeCall.kt");