diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt index 19e1ae97c95..bbc20569505 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/matchAndConvert.kt +++ b/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 + 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) diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt index 77d33ba6476..c8e1b8f24bd 100644 --- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FlatMapTransformation.kt +++ b/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 diff --git a/idea/testData/intentions/loopToCallChain/array.kt b/idea/testData/intentions/loopToCallChain/array.kt new file mode 100644 index 00000000000..c09b8777d0c --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/array.kt @@ -0,0 +1,9 @@ +// WITH_RUNTIME +fun foo(array: Array): String? { + for (s in array) { + if (s.isNotBlank()) { + return s + } + } + return null +} diff --git a/idea/testData/intentions/loopToCallChain/array.kt.after b/idea/testData/intentions/loopToCallChain/array.kt.after new file mode 100644 index 00000000000..6ec56d095fd --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/array.kt.after @@ -0,0 +1,4 @@ +// WITH_RUNTIME +fun foo(array: Array): String? { + return array.firstOrNull { it.isNotBlank() } +} diff --git a/idea/testData/intentions/loopToCallChain/customTypeWithIterator.kt b/idea/testData/intentions/loopToCallChain/customTypeWithIterator.kt new file mode 100644 index 00000000000..61f733da609 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/customTypeWithIterator.kt @@ -0,0 +1,16 @@ +// IS_APPLICABLE: false +// WITH_RUNTIME +class X { + operator fun iterator(): Iterator{ + return emptyList().iterator() + } +} + +fun foo(x: X): String? { + for (s in x) { + if (s.isNotBlank()) { + return s + } + } + return null +} diff --git a/idea/testData/intentions/loopToCallChain/intArray.kt b/idea/testData/intentions/loopToCallChain/intArray.kt new file mode 100644 index 00000000000..01c6ca99af3 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/intArray.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(array: IntArray): List { + val result = ArrayList() + for (i in array) { + if (i % 3 == 0) { + result.add(i) + } + } + return result +} diff --git a/idea/testData/intentions/loopToCallChain/intArray.kt.after b/idea/testData/intentions/loopToCallChain/intArray.kt.after new file mode 100644 index 00000000000..059579a8e3d --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/intArray.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +import java.util.ArrayList + +fun foo(array: IntArray): List { + val result = array.filter { it % 3 == 0 } + return result +} diff --git a/idea/testData/intentions/loopToCallChain/range.kt b/idea/testData/intentions/loopToCallChain/range.kt new file mode 100644 index 00000000000..d4765d06d81 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/range.kt @@ -0,0 +1,12 @@ +// WITH_RUNTIME +import java.util.* + +fun foo(): List { + val result = ArrayList() + for (i in 1..10) { + if (i % 3 == 0) { + result.add(i) + } + } + return result +} diff --git a/idea/testData/intentions/loopToCallChain/range.kt.after b/idea/testData/intentions/loopToCallChain/range.kt.after new file mode 100644 index 00000000000..34fcb781587 --- /dev/null +++ b/idea/testData/intentions/loopToCallChain/range.kt.after @@ -0,0 +1,7 @@ +// WITH_RUNTIME +import java.util.* + +fun foo(): List { + val result = (1..10).filter { it % 3 == 0 } + return result +}