From 954952ae71b0d6afdb752640294109636967ad6e Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Mon, 13 Jun 2016 14:52:14 +0300
Subject: [PATCH] Surround with null check : apply for ITERATOR_ON_NULLABLE

---
 .../kotlin/idea/quickfix/QuickFixRegistrar.kt       |  1 +
 .../idea/quickfix/SurroundWithNullCheckFix.kt       | 13 +++++++++++++
 .../fixNullableWithExclExclAbsentWithBadIterator.kt |  1 +
 .../surroundWithNullCheck/iteratorUnsafe.kt         |  6 ++++++
 .../surroundWithNullCheck/iteratorUnsafe.kt.after   |  8 ++++++++
 .../kotlin/idea/quickfix/QuickFixTestGenerated.java |  6 ++++++
 6 files changed, 35 insertions(+)
 create mode 100644 idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt
 create mode 100644 idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt
index 5fcf19c7ded..f8144a3b914 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixRegistrar.kt
+++ b/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)
diff --git a/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithNullCheckFix.kt b/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithNullCheckFix.kt
index 4ff690ea40b..c986b631504 100644
--- a/idea/src/org/jetbrains/kotlin/idea/quickfix/SurroundWithNullCheckFix.kt
+++ b/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 {
diff --git a/idea/testData/quickfix/expressions/fixNullableWithExclExclAbsentWithBadIterator.kt b/idea/testData/quickfix/expressions/fixNullableWithExclExclAbsentWithBadIterator.kt
index 27ee2e07cb2..23187cbaec2 100644
--- a/idea/testData/quickfix/expressions/fixNullableWithExclExclAbsentWithBadIterator.kt
+++ b/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 {
diff --git a/idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt b/idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt
new file mode 100644
index 00000000000..7a3e33fb567
--- /dev/null
+++ b/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) {}
+}
\ No newline at end of file
diff --git a/idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt.after b/idea/testData/quickfix/surroundWithNullCheck/iteratorUnsafe.kt.after
new file mode 100644
index 00000000000..dd4c68728a9
--- /dev/null
+++ b/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) {}
+    }
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
index 8af10c24ae7..11f2808e241 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java
+++ b/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");