KT-1478 Quick fix for using nullable collection in "for in" expression

#KT-1478 Fixed

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

@@ -23,13 +23,24 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.impl.source.tree.LeafPsiElement
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.JetBundle
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtPostfixExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
+import org.jetbrains.kotlin.util.OperatorChecks
+import org.jetbrains.kotlin.util.OperatorNameConventions
 
 public abstract class ExclExclCallFix : IntentionAction {
     override fun getFamilyName(): String = getText()
@@ -99,3 +110,37 @@ public class AddExclExclCallFix(val psiElement: PsiElement) : ExclExclCallFix()
                 = AddExclExclCallFix(diagnostic.getPsiElement())
     }
 }
+
+object MissingIteratorExclExclFixFactory : JetSingleIntentionActionFactory() {
+    override fun createAction(diagnostic: Diagnostic): IntentionAction? {
+        val element = diagnostic.psiElement
+        if (element !is KtExpression) return null
+        
+        val analyze = element.analyze(BodyResolveMode.PARTIAL)
+        val type = analyze.getType(element)
+        if (type == null || !TypeUtils.isNullableType(type)) return null
+        
+        val descriptor = type.constructor.declarationDescriptor
+
+        fun hasIteratorFunction(descriptor: ClassifierDescriptor?) : Boolean {
+            if (descriptor !is ClassDescriptor) return false
+
+            val memberScope = descriptor.unsubstitutedMemberScope
+            val functions = memberScope.getFunctions(OperatorNameConventions.ITERATOR, NoLookupLocation.FROM_IDE)
+
+            return functions.any { it.isOperator() && OperatorChecks.canBeOperator(it) }
+        }
+
+        when (descriptor) {
+            is TypeParameterDescriptor -> {
+                if (descriptor.upperBounds.none { hasIteratorFunction(it.constructor.declarationDescriptor) }) return null
+            }
+            is ClassifierDescriptor -> {
+                if (!hasIteratorFunction(descriptor)) return null
+            }
+            else -> return null
+        }
+
+        return AddExclExclCallFix(element)
+    }
+}

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

@@ -291,6 +291,7 @@ public class QuickFixRegistrar : QuickFixContributor {
         NEXT_MISSING.registerFactory(CreateNextFunctionActionFactory)
         NEXT_NONE_APPLICABLE.registerFactory(CreateNextFunctionActionFactory)
         ITERATOR_MISSING.registerFactory(CreateIteratorFunctionActionFactory)
+        ITERATOR_MISSING.registerFactory(MissingIteratorExclExclFixFactory)
         COMPONENT_FUNCTION_MISSING.registerFactory(CreateComponentFunctionActionFactory)
 
         DELEGATE_SPECIAL_FUNCTION_MISSING.registerFactory(CreatePropertyDelegateAccessorsActionFactory)

idea/testData/quickfix/expressions/fixNullableIterableGenericWithExclExcl.kt

@@ -0,0 +1,4 @@
+// "Add non-null asserted (!!) call" "true"
+fun <T: Collection<Int>?> foo(c: T) {
+    for (i in <caret>c) { }
+}

idea/testData/quickfix/expressions/fixNullableIterableGenericWithExclExcl.kt.after

@@ -0,0 +1,4 @@
+// "Add non-null asserted (!!) call" "true"
+fun <T: Collection<Int>?> foo(c: T) {
+    for (i in c!!) { }
+}

idea/testData/quickfix/expressions/fixNullableIterableWithExclExcl.kt

@@ -0,0 +1,5 @@
+// "Add non-null asserted (!!) call" "true"
+fun foo() {
+    val test : Collection<Int>? = null!!
+    for (i in <caret>test) { }
+}

idea/testData/quickfix/expressions/fixNullableIterableWithExclExcl.kt.after

@@ -0,0 +1,5 @@
+// "Add non-null asserted (!!) call" "true"
+fun foo() {
+    val test : Collection<Int>? = null!!
+    for (i in <caret>test!!) { }
+}

idea/testData/quickfix/expressions/fixNullableWithExclExclAbsentWithBadIterator.kt

@@ -0,0 +1,12 @@
+// "Add non-null asserted (!!) call" "false"
+// ACTION: Replace with a 'forEach' function call
+// ERROR: For-loop range must have an iterator() method
+
+class Some {
+    fun iterator(): Iterator<Int> = null!!
+}
+
+fun foo() {
+    val test: Some? = Some()
+    for (i in <caret>test) { }
+}

idea/testData/quickfix/expressions/fixNullableWithIteratorWithExclExcl.kt

@@ -0,0 +1,9 @@
+// "Add non-null asserted (!!) call" "true"
+class Some {
+    operator fun iterator(): Iterator<Int> = null!!
+}
+
+fun foo() {
+    val test: Some? = Some()
+    for (i in <caret>test) { }
+}

idea/testData/quickfix/expressions/fixNullableWithIteratorWithExclExcl.kt.after

@@ -0,0 +1,9 @@
+// "Add non-null asserted (!!) call" "true"
+class Some {
+    operator fun iterator(): Iterator<Int> = null!!
+}
+
+fun foo() {
+    val test: Some? = Some()
+    for (i in test!!) { }
+}

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

@@ -109,10 +109,14 @@ public abstract class AbstractQuickFixTest extends KotlinLightQuickFixTestCase {
                 catch (FileComparisonFailure e) {
                     throw e;
                 }
+                catch (AssertionError e) {
+                    throw e;
+                }
                 catch (Throwable e) {
                     e.printStackTrace();
                     fail(testName);
-                } finally {
+                }
+                finally {
                     ConfigLibraryUtil.unconfigureLibrariesByDirective(getModule(), fileText);
                 }
             }

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

@@ -4070,6 +4070,30 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/expressions"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }
 
+        @TestMetadata("fixNullableIterableGenericWithExclExcl.kt")
+        public void testFixNullableIterableGenericWithExclExcl() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/expressions/fixNullableIterableGenericWithExclExcl.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("fixNullableIterableWithExclExcl.kt")
+        public void testFixNullableIterableWithExclExcl() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/expressions/fixNullableIterableWithExclExcl.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("fixNullableWithExclExclAbsentWithBadIterator.kt")
+        public void testFixNullableWithExclExclAbsentWithBadIterator() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/expressions/fixNullableWithExclExclAbsentWithBadIterator.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("fixNullableWithIteratorWithExclExcl.kt")
+        public void testFixNullableWithIteratorWithExclExcl() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/expressions/fixNullableWithIteratorWithExclExcl.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("removeUselessCast.kt")
         public void testRemoveUselessCast() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/quickfix/expressions/removeUselessCast.kt");