Refactored API to allow merging ResultTransformation with any kind of previous transformation

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

@@ -36,23 +36,16 @@ interface SequenceTransformation : Transformation {
 }
 
 interface ResultTransformation : Transformation {
-    val canIncludeFilter: Boolean
-    val canIncludeMap: Boolean
+    fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? = null
 
     val commentSavingRange: PsiChildRange
     val commentRestoringRange: PsiChildRange
 
-    fun generateCode(
-            chainedCallGenerator: ChainedCallGenerator,
-            filter: FilterOrMap?,
-            map: FilterOrMap?
-    ): KtExpression
+    fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression
 
     fun convertLoop(resultCallChain: KtExpression): KtExpression
 }
 
-data class FilterOrMap(val expression: KtExpression, val workingVariable: KtCallableDeclaration)
-
 data class MatchingState(
         val outerLoop: KtForExpression,
         val innerLoop: KtForExpression,

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

@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.idea.intentions.loopToCallChain.result.FindAndReturn
 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.CommentSaver
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.*
@@ -156,12 +155,10 @@ private fun checkSmartCastsPreserved(loop: KtForExpression, matchResult: ResultT
 
 private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression): KtExpression {
     var sequenceTransformations = sequenceTransformations
-    val last = sequenceTransformations.lastOrNull()
-    var lastFilter: FilterOrMap? = null
-    val lastMap: FilterOrMap? = null //TODO
-    if (last is FilterTransformation && resultTransformation.canIncludeFilter) {
-        val condition = if (last.isInverse) last.condition.negate() else last.condition
-        lastFilter = FilterOrMap(condition, last.inputVariable)
+    var resultTransformation = resultTransformation
+    while(true) {
+        val last = sequenceTransformations.lastOrNull() ?: break
+        resultTransformation = resultTransformation.mergeWithPrevious(last) ?: break
         sequenceTransformations = sequenceTransformations.dropLast(1)
     }
 
@@ -181,7 +178,7 @@ private fun ResultTransformationMatch.generateCallChain(loop: KtForExpression):
         callChain = transformation.generateCode(chainedCallGenerator)
     }
 
-    callChain = resultTransformation.generateCode(chainedCallGenerator, lastFilter, lastMap)
+    callChain = resultTransformation.generateCode(chainedCallGenerator)
     return callChain
 }
 

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

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result
 import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
@@ -27,24 +28,25 @@ class FindAndAssignTransformation(
         private val loop: KtForExpression,
         override val inputVariable: KtCallableDeclaration,
         private val stdlibFunName: String,
-        private val initialDeclaration: KtProperty
+        private val initialDeclaration: KtProperty,
+        private val filter: KtExpression? = null
 ) : ResultTransformation {
 
+    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
+        if (previousTransformation !is FilterTransformation) return null
+        if (filter != null) return null //TODO
+        return FindAndAssignTransformation(loop, previousTransformation.inputVariable, stdlibFunName, initialDeclaration, previousTransformation.buildRealCondition())
+    }
+
     override val commentSavingRange = PsiChildRange(initialDeclaration, loop.unwrapIfLabeled())
     override val commentRestoringRange = commentSavingRange.withoutLastStatement()
 
-    override val canIncludeFilter: Boolean
-        get() = true
-
-    override val canIncludeMap: Boolean
-        get() = false
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator, filter: FilterOrMap?, map: FilterOrMap?): KtExpression {
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return if (filter == null) {
             chainedCallGenerator.generate("$stdlibFunName()")
         }
         else {
-            val lambda = generateLambda(filter.workingVariable, filter.expression)
+            val lambda = generateLambda(inputVariable, filter)
             chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda)
         }
     }

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain.result
 
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
+import org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence.FilterTransformation
 import org.jetbrains.kotlin.psi.KtCallableDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtForExpression
@@ -27,25 +28,26 @@ class FindAndReturnTransformation(
         private val loop: KtForExpression,
         override val inputVariable: KtCallableDeclaration,
         private val stdlibFunName: String,
-        private val endReturn: KtReturnExpression
+        private val endReturn: KtReturnExpression,
+        private val filter: KtExpression? = null
 ) : ResultTransformation {
 
+    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): ResultTransformation? {
+        if (previousTransformation !is FilterTransformation) return null
+        if (filter != null) return null //TODO
+        return FindAndReturnTransformation(loop, previousTransformation.inputVariable, stdlibFunName, endReturn, previousTransformation.buildRealCondition())
+    }
+
     override val commentSavingRange = PsiChildRange(loop.unwrapIfLabeled(), endReturn)
 
     override val commentRestoringRange = commentSavingRange.withoutFirstStatement()
 
-    override val canIncludeFilter: Boolean
-        get() = true
-
-    override val canIncludeMap: Boolean
-        get() = false
-
-    override fun generateCode(chainedCallGenerator: ChainedCallGenerator, filter: FilterOrMap?, map: FilterOrMap?): KtExpression {
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
         return if (filter == null) {
             chainedCallGenerator.generate("$stdlibFunName()")
         }
         else {
-            val lambda = generateLambda(filter.workingVariable, filter.expression)
+            val lambda = generateLambda(inputVariable, filter)
             chainedCallGenerator.generate("$stdlibFunName $0:'{}'", lambda)
         }
     }

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

@@ -33,6 +33,8 @@ class FilterTransformation(
         assert(condition.isPhysical)
     }
 
+    fun buildRealCondition() = if (isInverse) condition.negate() else condition
+
     override val affectsIndex: Boolean
         get() = true