Check for continue-statement of the inner loop!

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

@@ -54,10 +54,11 @@ interface ResultTransformation : Transformation {
 data class FilterOrMap(val expression: KtExpression, val workingVariable: KtCallableDeclaration)
 
 data class MatchingState(
+        val outerLoop: KtForExpression,
+        val innerLoop: KtForExpression,
         val statements: Collection<KtExpression>,
         val workingVariable: KtCallableDeclaration,
-        val indexVariable: KtCallableDeclaration?,
+        val indexVariable: KtCallableDeclaration?
-        val loop: KtForExpression
 )
 
 interface SequenceTransformationMatcher {

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

@@ -54,10 +54,11 @@ object MatcherRegistrar {
 fun match(loop: KtForExpression): ResultTransformationMatch? {
     val sequenceTransformations = ArrayList<SequenceTransformation>()
     var state = MatchingState(
+            outerLoop = loop,
+            innerLoop = loop,
             statements = listOf(loop.body ?: return null),
             workingVariable = loop.loopParameter ?: return null,
-            indexVariable = null,
+            indexVariable = null
-            loop = loop
     )
 
     MatchLoop@

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

@@ -68,7 +68,7 @@ class FindAndAssignTransformation(
             if (state.statements.size != 2) return null
 
             val breakExpression = state.statements.last() as? KtBreakExpression ?: return null
-            if (!breakExpression.isBreakOrContinueOfLoop(state.loop)) return null
+            if (!breakExpression.isBreakOrContinueOfLoop(state.outerLoop)) return null
 
             val binaryExpression = state.statements.first() as? KtBinaryExpression ?: return null
             if (binaryExpression.operationToken != KtTokens.EQ) return null
@@ -76,16 +76,16 @@ class FindAndAssignTransformation(
             val right = binaryExpression.right ?: return null
 
             //TODO: support also assignment instead of declaration
-            val declarationBeforeLoop = state.loop.previousStatement() as? KtProperty ?: return null
+            val declarationBeforeLoop = state.outerLoop.previousStatement() as? KtProperty ?: return null
             val initializer = declarationBeforeLoop.initializer ?: return null
             if (!left.isVariableReference(declarationBeforeLoop)) return null
 
-            val usageCountInLoop = ReferencesSearch.search(declarationBeforeLoop, LocalSearchScope(state.loop)).count()
+            val usageCountInLoop = ReferencesSearch.search(declarationBeforeLoop, LocalSearchScope(state.outerLoop)).count()
             if (usageCountInLoop != 1) return null // this should be the only usage of this variable inside the loop
 
             val stdlibFunName = stdlibFunNameForFind(right, initializer, state.workingVariable) ?: return null
 
-            val transformation = FindAndAssignTransformation(state.loop, state.workingVariable, stdlibFunName, declarationBeforeLoop)
+            val transformation = FindAndAssignTransformation(state.outerLoop, state.workingVariable, stdlibFunName, declarationBeforeLoop)
             return ResultTransformationMatch(transformation)
         }
     }

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

@@ -62,7 +62,7 @@ class FindAndReturnTransformation(
             if (state.indexVariable != null) return null
 
             val returnInLoop = state.statements.singleOrNull() as? KtReturnExpression ?: return null
-            val returnAfterLoop = state.loop.nextStatement() as? KtReturnExpression ?: return null
+            val returnAfterLoop = state.outerLoop.nextStatement() as? KtReturnExpression ?: return null
             if (returnInLoop.getLabelName() != returnAfterLoop.getLabelName()) return null
 
             val returnValueInLoop = returnInLoop.returnedExpression ?: return null
@@ -70,7 +70,7 @@ class FindAndReturnTransformation(
 
             val stdlibFunName = stdlibFunNameForFind(returnValueInLoop, returnValueAfterLoop, state.workingVariable) ?: return null
 
-            val transformation = FindAndReturnTransformation(state.loop, state.workingVariable, stdlibFunName, returnAfterLoop)
+            val transformation = FindAndReturnTransformation(state.outerLoop, state.workingVariable, stdlibFunName, returnAfterLoop)
             return ResultTransformationMatch(transformation)
         }
     }

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

@@ -59,7 +59,7 @@ class FilterTransformation(
             }
             else {
                 val continueExpression = then.blockExpressionsOrSingle().singleOrNull() as? KtContinueExpression ?: return null
-                if (!continueExpression.isBreakOrContinueOfLoop(state.loop)) return null
+                if (!continueExpression.isBreakOrContinueOfLoop(state.innerLoop)) return null
                 val transformation = createFilterTransformation(state.workingVariable, condition, isInverse = true)
                 val newState = state.copy(statements = state.statements.drop(1))
                 return SequenceTransformationMatch(transformation, newState)

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

@@ -59,6 +59,7 @@ class FlatMapTransformation(
             val loopBody = nestedLoop.body ?: return null
             val transformation = FlatMapTransformation(state.workingVariable, transform)
             val newState = state.copy(
+                    innerLoop = nestedLoop,
                     statements = listOf(loopBody),
                     workingVariable = newWorkingVariable
             )

idea/testData/intentions/loopToCallChain/flatMapWithContinue.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+fun foo(list: List<String>): String? {
+    <caret>for (s in list) {
+        for (line in s.lines()) {
+            if (line.isBlank()) continue
+            return line
+        }
+    }
+    return null
+}

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

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+fun foo(list: List<String>): String? {
+    <caret>return list
+            .flatMap { it.lines() }
+            .firstOrNull { !it.isBlank() }
+}

idea/testData/intentions/loopToCallChain/flatMap_wrongContinue.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String>): String? {
+    OuterLoop@
+    <caret>for (s in list) {
+        for (line in s.lines()) {
+            if (line.isBlank()) continue@OuterLoop
+            return line
+        }
+    }
+    return null
+}