Checking type of the collection being iterated

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain
 
 import com.intellij.openapi.util.Key
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
@@ -27,7 +29,9 @@ import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTrans
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FlatMapTransformation
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.IntroduceIndexMatcher
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.MapTransformation
+import org.jetbrains.kotlin.idea.project.builtIns
 import org.jetbrains.kotlin.idea.util.CommentSaver
+import org.jetbrains.kotlin.idea.util.FuzzyType
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
@@ -174,16 +178,37 @@ private fun extractLoopData(loop: KtForExpression): LoopData? {
         val qualifiedExpression = loopRange as? KtDotQualifiedExpression
         if (qualifiedExpression != null) {
             val call = qualifiedExpression.selectorExpression as? KtCallExpression
-            //TODO: check that it's the correct "withIndex"
             if (call != null && call.calleeExpression.isSimpleName(Name.identifier("withIndex")) && call.valueArguments.isEmpty()) {
-                return LoopData(destructuringParameter.entries[1], destructuringParameter.entries[0], qualifiedExpression.receiverExpression)
+                val receiver = qualifiedExpression.receiverExpression
+                if (!isExpressionTypeSupported(receiver)) return null
+                return LoopData(destructuringParameter.entries[1], destructuringParameter.entries[0], receiver)
             }
         }
     }
 
+    if (!isExpressionTypeSupported(loopRange)) return null
+
     return LoopData(loop.loopParameter ?: return null, null, loopRange)
 }
 
+private fun isExpressionTypeSupported(expression: KtExpression): Boolean {
+    val type = expression.analyze(BodyResolveMode.PARTIAL).getType(expression) ?: return false
+
+    val builtIns = expression.builtIns
+    return when {
+        type.isSubtypeOf(builtIns.iterable) -> true
+        type.isSubtypeOf(builtIns.array) -> true
+        KotlinBuiltIns.isPrimitiveArray(type) -> true
+        // TODO: support Sequence<T>
+        else -> false
+    }
+}
+
+private fun KotlinType.isSubtypeOf(classDescriptor: ClassDescriptor): Boolean {
+    val fuzzyType = FuzzyType(classDescriptor.defaultType, classDescriptor.declaredTypeParameters)
+    return fuzzyType.checkIsSuperTypeOf(this) != null
+}
+
 private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: MatchResult): Boolean {
     val bindingContext = loop.analyze(BodyResolveMode.FULL)
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt

@@ -17,8 +17,8 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence
 
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
-import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
+import org.jetbrains.kotlin.idea.project.builtIns
 import org.jetbrains.kotlin.idea.util.FuzzyType
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
@@ -59,7 +59,7 @@ class FlatMapTransformation(
             val transform = nestedLoop.loopRange ?: return null
             // check that we iterate over Iterable
             val nestedSequenceType = transform.analyze(BodyResolveMode.PARTIAL).getType(transform) ?: return null
-            val builtIns = transform.getResolutionFacade().moduleDescriptor.builtIns
+            val builtIns = transform.builtIns
             val iterableType = FuzzyType(builtIns.iterableType, builtIns.iterable.declaredTypeParameters)
             if (iterableType.checkIsSuperTypeOf(nestedSequenceType) == null) return null
 

idea/testData/intentions/loopToCallChain/array.kt

@@ -0,0 +1,9 @@
+// WITH_RUNTIME
+fun foo(array: Array<String>): String? {
+    <caret>for (s in array) {
+        if (s.isNotBlank()) {
+            return s
+        }
+    }
+    return null
+}

idea/testData/intentions/loopToCallChain/array.kt.after

@@ -0,0 +1,4 @@
+// WITH_RUNTIME
+fun foo(array: Array<String>): String? {
+    <caret>return array.firstOrNull { it.isNotBlank() }
+}

idea/testData/intentions/loopToCallChain/customTypeWithIterator.kt

@@ -0,0 +1,16 @@
+// IS_APPLICABLE: false
+// WITH_RUNTIME
+class X {
+    operator fun iterator(): Iterator<String>{
+        return emptyList<String>().iterator()
+    }
+}
+
+fun foo(x: X): String? {
+    <caret>for (s in x) {
+        if (s.isNotBlank()) {
+            return s
+        }
+    }
+    return null
+}

idea/testData/intentions/loopToCallChain/intArray.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(array: IntArray): List<Int> {
+    val result = ArrayList<Int>()
+    <caret>for (i in array) {
+        if (i % 3 == 0) {
+            result.add(i)
+        }
+    }
+    return result
+}

idea/testData/intentions/loopToCallChain/intArray.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+import java.util.ArrayList
+
+fun foo(array: IntArray): List<Int> {
+    val result = array.filter { it % 3 == 0 }
+    return result
+}

idea/testData/intentions/loopToCallChain/range.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+import java.util.*
+
+fun foo(): List<Int> {
+    val result = ArrayList<Int>()
+    <caret>for (i in 1..10) {
+        if (i % 3 == 0) {
+            result.add(i)
+        }
+    }
+    return result
+}

idea/testData/intentions/loopToCallChain/range.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+import java.util.*
+
+fun foo(): List<Int> {
+    val <caret>result = (1..10).filter { it % 3 == 0 }
+    return result
+}