Merging subsequent filtering checks even when index variable is used

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

@@ -80,10 +80,7 @@ fun match(loop: KtForExpression): MatchResult? {
 
     MatchLoop@
     while (true) {
-        val block = state.statements.singleOrNull() as? KtBlockExpression
+        state = state.unwrapBlock()
-        if (block != null) {
-            state = state.copy(statements = block.statements)
-        }
 
         val inputVariableUsed = state.inputVariable.hasUsages(state.statements)
 

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/AddToCollectionTransformation.kt

@@ -22,7 +22,11 @@ import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
-import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.*
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterNotNullTransformation
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FlatMapTransformation
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.MapTransformation
+import org.jetbrains.kotlin.idea.intentions.negate
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
@@ -37,11 +41,7 @@ class AddToCollectionTransformation(
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         return when (previousTransformation) {
             is FilterTransformation -> {
-                FilterToTransformation.create(loop, previousTransformation.inputVariable, targetCollection, previousTransformation.condition, previousTransformation.isInverse)
+                FilterToTransformation.create(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, targetCollection, previousTransformation.condition, previousTransformation.isInverse)
-            }
-
-            is FilterIndexedTransformation -> {
-                FilterIndexedToTransformation.create(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, targetCollection, previousTransformation.condition)
             }
 
             is FilterNotNullTransformation -> {
@@ -189,18 +189,28 @@ class AddToCollectionTransformation(
 class FilterToTransformation private constructor(
         loop: KtForExpression,
         private val inputVariable: KtCallableDeclaration,
+        private val indexVariable: KtCallableDeclaration?,
         private val targetCollection: KtExpression,
-        private val filter: KtExpression,
+        private val condition: KtExpression,
-        isInverse: Boolean
+        private val isInverse: Boolean
 ) : ReplaceLoopResultTransformation(loop) {
 
-    private val functionName = if (isInverse) "filterNotTo" else "filterTo"
+    private fun effectiveCondition() = if (isInverse) condition.negate() else condition
+
+    private val functionName = when {
+        indexVariable != null -> "filterIndexedTo"
+        isInverse -> "filterNotTo"
+        else -> "filterTo"
+    }
 
     override val presentation: String
         get() = "$functionName(){}"
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(inputVariable, filter)
+        val lambda = if (indexVariable != null)
+            generateLambda(effectiveCondition(), indexVariable, inputVariable)
+        else
+            generateLambda(inputVariable, condition)
         return chainedCallGenerator.generate("$functionName($0) $1:'{}'", targetCollection, lambda)
     }
 
@@ -208,53 +218,18 @@ class FilterToTransformation private constructor(
         fun create(
                 loop: KtForExpression,
                 inputVariable: KtCallableDeclaration,
+                indexVariable: KtCallableDeclaration?,
                 targetCollection: KtExpression,
-                filter: KtExpression,
+                condition: KtExpression,
                 isInverse: Boolean
         ): ResultTransformation {
             val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
             if (initialization != null && initialization.initializer.hasNoSideEffect()) {
-                val transformation = FilterToTransformation(loop, inputVariable, initialization.initializer, filter, isInverse)
+                val transformation = FilterToTransformation(loop, inputVariable, indexVariable, initialization.initializer, condition, isInverse)
                 return AssignToVariableResultTransformation.createDelegated(transformation, initialization)
             }
             else {
-                return FilterToTransformation(loop, inputVariable, targetCollection, filter, isInverse)
+                return FilterToTransformation(loop, inputVariable, indexVariable, targetCollection, condition, isInverse)
-            }
-        }
-    }
-}
-
-class FilterIndexedToTransformation private constructor(
-        loop: KtForExpression,
-        private val inputVariable: KtCallableDeclaration,
-        private val indexVariable: KtCallableDeclaration,
-        private val targetCollection: KtExpression,
-        private val filter: KtExpression
-) : ReplaceLoopResultTransformation(loop) {
-
-    override val presentation: String
-        get() = "filterIndexedTo(){}"
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(filter, indexVariable, inputVariable)
-        return chainedCallGenerator.generate("filterIndexedTo($0) $1:'{}'", targetCollection, lambda)
-    }
-
-    companion object {
-        fun create(
-                loop: KtForExpression,
-                inputVariable: KtCallableDeclaration,
-                indexVariable: KtCallableDeclaration,
-                targetCollection: KtExpression,
-                filter: KtExpression
-        ): ResultTransformation {
-            val initialization = targetCollection.detectInitializationBeforeLoop(loop, checkNoOtherUsagesInLoop = true)
-            if (initialization != null && initialization.initializer.hasNoSideEffect()) {
-                val transformation = FilterIndexedToTransformation(loop, inputVariable, indexVariable, initialization.initializer, filter)
-                return AssignToVariableResultTransformation.createDelegated(transformation, initialization)
-            }
-            else {
-                return FilterIndexedToTransformation(loop, inputVariable, indexVariable, targetCollection, filter)
             }
         }
     }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/CountTransformation.kt

@@ -32,6 +32,7 @@ class CountTransformation(
 
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         if (previousTransformation !is FilterTransformation) return null
+        if (previousTransformation.indexVariable != null) return null
         assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"}
         return CountTransformation(loop, previousTransformation.inputVariable, initialization, previousTransformation.effectiveCondition())
     }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndAssignTransformation.kt

@@ -33,6 +33,7 @@ class FindAndAssignTransformation(
 
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         if (previousTransformation !is FilterTransformation) return null
+        if (previousTransformation.indexVariable != null) return null
         assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"}
         return FindAndAssignTransformation(loop, generator, initialization, previousTransformation.effectiveCondition())
     }

idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/result/FindAndReturnTransformation.kt

@@ -32,6 +32,7 @@ class FindAndReturnTransformation(
 
     override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
         if (previousTransformation !is FilterTransformation) return null
+        if (previousTransformation.indexVariable != null) return null
         assert(filter == null) { "Should not happen because no 2 consecutive FilterTransformation's possible"}
         return FindAndReturnTransformation(loop, generator, endReturn, previousTransformation.effectiveCondition())
     }

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

@@ -26,22 +26,19 @@ import org.jetbrains.kotlin.psi.psiUtil.blockExpressionsOrSingle
 class FilterTransformation(
         override val loop: KtForExpression,
         val inputVariable: KtCallableDeclaration,
+        val indexVariable: KtCallableDeclaration?,
         val condition: KtExpression,
         val isInverse: Boolean
 ) : SequenceTransformation {
 
     fun effectiveCondition() = if (isInverse) condition.negate() else condition
 
-    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? {
+    private val functionName = when {
-        if (previousTransformation !is FilterTransformation) return null
+        indexVariable != null -> "filterIndexed"
-        assert(previousTransformation.inputVariable == inputVariable)
+        isInverse -> "filterNot"
-        val mergedCondition = KtPsiFactory(condition).createExpressionByPattern(
+        else -> "filter"
-                "$0 && $1", previousTransformation.effectiveCondition(), effectiveCondition())
-        return FilterTransformation(loop, inputVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases?
     }
 
-    private val functionName = if (isInverse) "filterNot" else "filter"
-
     override val affectsIndex: Boolean
         get() = true
 
@@ -49,7 +46,10 @@ class FilterTransformation(
         get() = "$functionName{}"
 
     override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(inputVariable, condition)
+        val lambda = if (indexVariable != null)
+            generateLambda(effectiveCondition(), indexVariable, inputVariable)
+        else
+            generateLambda(inputVariable, condition)
         return chainedCallGenerator.generate("$0$1:'{}'", functionName, lambda)
     }
 
@@ -76,22 +76,29 @@ class FilterTransformation(
      */
     object Matcher : SequenceTransformationMatcher {
         override fun match(state: MatchingState): SequenceTransformationMatch? {
-            var (transformation, currentState) = matchOneTransformation(state, takeWhileAllowed = true) ?: return null
+            // we merge filter transformations here instead of FilterTransformation.mergeWithPrevious() because of filterIndexed that won't merge otherwise
 
-/*
+            var (transformation, currentState) = matchOneTransformation(state) ?: return null
-            if (transformation is FilterTransformation || transformation is FilterIndexedTransformation) {
+
+            if (transformation is FilterTransformation) {
                 while (true) {
-                    val (nextTransformation, nextState) = matchOneTransformation(currentState, takeWhileAllowed = false) ?: break
+                    currentState = currentState.unwrapBlock()
-                    if (nextTransformation !is FilterTransformation && nextTransformation !is FilterIndexedTransformation) break
 
+                    val (nextTransformation, nextState) = matchOneTransformation(currentState) ?: break
+                    if (nextTransformation !is FilterTransformation) break
+
+                    val indexVariable = transformation.indexVariable ?: nextTransformation.indexVariable
+                    val mergedCondition = KtPsiFactory(state.outerLoop).createExpressionByPattern(
+                            "$0 && $1", transformation.effectiveCondition(), nextTransformation.effectiveCondition())
+                    transformation = FilterTransformation(state.outerLoop, transformation.inputVariable, indexVariable, mergedCondition, isInverse = false) //TODO: build filterNot in some cases?
+                    currentState = nextState
                 }
             }
-*/
 
             return SequenceTransformationMatch(transformation, currentState)
         }
 
-        private fun matchOneTransformation(state: MatchingState, takeWhileAllowed: Boolean): Pair<SequenceTransformation, MatchingState>? {
+        private fun matchOneTransformation(state: MatchingState): Pair<SequenceTransformation, MatchingState>? {
             val ifStatement = state.statements.firstOrNull() as? KtIfExpression ?: return null
             if (ifStatement.`else` != null) return null
             val condition = ifStatement.condition ?: return null
@@ -113,7 +120,6 @@ class FilterTransformation(
                     }
 
                     is KtBreakExpression -> {
-                        if (!takeWhileAllowed) return null
                         if (statement.targetLoop() != state.outerLoop) return null
                         val transformation = TakeWhileTransformation(state.outerLoop, state.inputVariable, condition.negate())
                         val newState = state.copy(statements = state.statements.drop(1))
@@ -125,8 +131,6 @@ class FilterTransformation(
             }
         }
 
-//        private fun mergeFilterTransformations(transformation1: )
-
         //TODO: choose filter or filterNot depending on condition
         private fun createFilterTransformation(
                 loop: KtForExpression,
@@ -139,7 +143,7 @@ class FilterTransformation(
             val effectiveCondition = if (isInverse) condition.negate() else condition
 
             if (indexVariable != null && indexVariable.hasUsages(condition)) {
-                return FilterIndexedTransformation(loop, inputVariable, indexVariable, effectiveCondition)
+                return FilterTransformation(loop, inputVariable, indexVariable, effectiveCondition, isInverse = false)
             }
 
             if (effectiveCondition is KtIsExpression
@@ -160,7 +164,7 @@ class FilterTransformation(
                 return FilterNotNullTransformation(loop)
             }
 
-            return FilterTransformation(loop, inputVariable, condition, isInverse)
+            return FilterTransformation(loop, inputVariable, null, condition, isInverse)
         }
     }
 }
@@ -200,27 +204,6 @@ class FilterNotNullTransformation(override val loop: KtForExpression) : Sequence
     }
 }
 
-class FilterIndexedTransformation(
-        override val loop: KtForExpression,
-        val inputVariable: KtCallableDeclaration,
-        val indexVariable: KtCallableDeclaration,
-        val condition: KtExpression
-) : SequenceTransformation {
-
-    //TODO: how to handle multiple if's using index?
-
-    override val affectsIndex: Boolean
-        get() = true
-
-    override val presentation: String
-        get() = "filterIndexed{}"
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
-        val lambda = generateLambda(condition, indexVariable, inputVariable)
-        return chainedCallGenerator.generate("filterIndexed $0:'{}'", lambda)
-    }
-}
-
 class TakeWhileTransformation(
         override val loop: KtForExpression,
         val inputVariable: KtCallableDeclaration,

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

@@ -306,4 +306,9 @@ private fun isEmbeddedBreakOrContinue(expression: KtExpressionWithLabel): Boolea
 
         else -> return true
     }
-}
+}
+
+fun MatchingState.unwrapBlock(): MatchingState {
+    val block = statements.singleOrNull() as? KtBlockExpression ?: return this
+    return this.copy(statements = block.statements)
+}

idea/testData/intentions/loopToCallChain/filterIndexed_merge1.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'"
+fun foo(list: List<String>): String? {
+    <caret>for ((index, s) in list.withIndex()) {
+        if (s.length > index * 10) continue
+        if (s.length > index) {
+            return s
+        }
+    }
+    return null
+}

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

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'"
+fun foo(list: List<String>): String? {
+    <caret>return list
+            .filterIndexed { index, s -> s.length <= index * 10 && s.length > index }
+            .firstOrNull()
+}

idea/testData/intentions/loopToCallChain/filterIndexed_merge2.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'"
+fun foo(list: List<String>): String? {
+    <caret>for ((index, s) in list.withIndex()) {
+        if (s.isBlank()) continue
+        if (s.length > index) {
+            return s
+        }
+    }
+    return null
+}

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

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'"
+fun foo(list: List<String>): String? {
+    <caret>return list
+            .filterIndexed { index, s -> !s.isBlank() && s.length > index }
+            .firstOrNull()
+}

idea/testData/intentions/loopToCallChain/filterIndexed_merge3.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'"
+fun foo(list: List<String>): String? {
+    <caret>for ((index, s) in list.withIndex()) {
+        if (s.length > index) continue
+        if (s.isNotBlank()) {
+            return s
+        }
+    }
+    return null
+}

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

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'filterIndexed{}.firstOrNull()'"
+fun foo(list: List<String>): String? {
+    <caret>return list
+            .filterIndexed { index, s -> s.length <= index && s.isNotBlank() }
+            .firstOrNull()
+}