Prohibiting incorrect transformations of loops with expression-embedded break or continue + allowed "?: continue" pattern for mapNotNull

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

@@ -100,10 +100,19 @@ data class MatchingState(
         val initializationStatementsToDelete: Collection<KtExpression> = emptyList()
 )
 
+interface BaseMatcher {
+    /**
+     * Implementors should return true if they match some constructs with expression-embedded break or continue.
+     * In this case they are obliged to deal with them and filter out invalid cases.
+     */
+    val embeddedBreakOrContinuePossible: Boolean
+        get() = false
+}
+
 /**
  * A matcher that can recognize one or more [SequenceTransformation]'s
  */
-interface SequenceTransformationMatcher {
+interface SequenceTransformationMatcher : BaseMatcher {
     fun match(state: MatchingState): SequenceTransformationMatch?
 }
 
@@ -118,7 +127,7 @@ class SequenceTransformationMatch(
  * A matcher that can recognize a [ResultTransformation] (optionally prepended by some [SequenceTransformation]'s).
  * Should match the whole rest part of the loop.
  */
-interface ResultTransformationMatcher {
+interface ResultTransformationMatcher : BaseMatcher {
     fun match(state: MatchingState): ResultTransformationMatch?
 
     val indexVariableUsePossible: Boolean

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

@@ -66,6 +66,9 @@ data class MatchResult(
 fun match(loop: KtForExpression): MatchResult? {
     val (inputVariable, indexVariable, sequenceExpression) = extractLoopData(loop) ?: return null
 
+    // used just as optimization to avoid unnecessary checks
+    val loopContainsEmbeddedBreakOrContinue = loop.containsEmbeddedBreakOrContinue()
+
     val sequenceTransformations = ArrayList<SequenceTransformation>()
     var state = MatchingState(
             outerLoop = loop,
@@ -89,8 +92,11 @@ fun match(loop: KtForExpression): MatchResult? {
             state = state.copy(indexVariable = null)
         }
 
+        val restContainsEmbeddedBreakOrContinue = loopContainsEmbeddedBreakOrContinue && state.statements.any { it.containsEmbeddedBreakOrContinue() }
+
         for (matcher in MatcherRegistrar.resultMatchers) {
             if (state.indexVariable != null && !matcher.indexVariableUsePossible) continue
+            if (restContainsEmbeddedBreakOrContinue && !matcher.embeddedBreakOrContinuePossible) continue
 
             val match = matcher.match(state)
             if (match != null) {
@@ -115,8 +121,8 @@ fun match(loop: KtForExpression): MatchResult? {
                     return null
                 }
 
-                var newState = match.newState
                 // check that old input variable is not needed anymore
+                var newState = match.newState
                 if (state.inputVariable != newState.inputVariable && state.inputVariable.hasUsages(newState.statements)) return null
 
                 if (state.indexVariable != null && match.transformations.any { it.affectsIndex }) {
@@ -125,6 +131,12 @@ fun match(loop: KtForExpression): MatchResult? {
                     newState = newState.copy(indexVariable = null)
                 }
 
+                if (restContainsEmbeddedBreakOrContinue && !matcher.embeddedBreakOrContinuePossible) {
+                    val countBefore = state.statements.sumBy { it.countEmbeddedBreaksAndContinues() }
+                    val countAfter = newState.statements.sumBy { it.countEmbeddedBreaksAndContinues() }
+                    if (countAfter != countBefore) continue // some embedded break or continue in the matched part
+                }
+
                 sequenceTransformations.addAll(match.transformations)
                 state = newState
                 continue@MatchLoop

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

@@ -17,10 +17,8 @@
 package org.jetbrains.kotlin.idea.intentions.loopToCallChain.sequence
 
 import org.jetbrains.kotlin.idea.intentions.loopToCallChain.*
-import org.jetbrains.kotlin.psi.KtCallableDeclaration
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtForExpression
-import org.jetbrains.kotlin.psi.KtProperty
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.*
 
 class MapTransformation(
         override val loop: KtForExpression,
@@ -47,12 +45,32 @@ class MapTransformation(
      *     }
      */
     object Matcher : SequenceTransformationMatcher {
+        override val embeddedBreakOrContinuePossible: Boolean
+            get() = true
+
         override fun match(state: MatchingState): SequenceTransformationMatch? {
             val declaration = state.statements.firstOrNull() as? KtProperty ?: return null //TODO: support multi-variables
             val initializer = declaration.initializer ?: return null
             if (declaration.hasWriteUsages()) return null
             val restStatements = state.statements.drop(1)
 
+            if (initializer is KtBinaryExpression && initializer.operationToken == KtTokens.ELVIS) {
+                val continueExpression = initializer.right as? KtContinueExpression ?: return null
+                if (continueExpression.targetLoop() != state.innerLoop) return null
+
+                val mapping = initializer.left ?: return null
+                if (mapping.containsEmbeddedBreakOrContinue()) return null
+
+                val transformation = if (state.indexVariable != null && state.indexVariable.hasUsages(mapping))
+                    MapIndexedNotNullTransformation(state.outerLoop, state.inputVariable, state.indexVariable, mapping)
+                else
+                    MapNotNullTransformation(state.outerLoop, state.inputVariable, mapping)
+                val newState = state.copy(statements = restStatements, inputVariable = declaration)
+                return SequenceTransformationMatch(transformation, newState)
+            }
+
+            if (initializer.containsEmbeddedBreakOrContinue()) return null
+
             val transformation = if (state.indexVariable != null && state.indexVariable.hasUsages(initializer))
                 MapIndexedTransformation(state.outerLoop, state.inputVariable, state.indexVariable, initializer)
             else
@@ -81,3 +99,40 @@ class MapIndexedTransformation(
         return chainedCallGenerator.generate("mapIndexed $0:'{}'", lambda)
     }
 }
+
+class MapNotNullTransformation(
+        override val loop: KtForExpression,
+        val inputVariable: KtCallableDeclaration,
+        val mapping: KtExpression
+) : SequenceTransformation {
+
+    override val affectsIndex: Boolean
+        get() = true
+
+    override val presentation: String
+        get() = "mapNotNull{}"
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, mapping)
+        return chainedCallGenerator.generate("mapNotNull$0:'{}'", lambda)
+    }
+}
+
+class MapIndexedNotNullTransformation(
+        override val loop: KtForExpression,
+        val inputVariable: KtCallableDeclaration,
+        val indexVariable: KtCallableDeclaration,
+        val mapping: KtExpression
+) : SequenceTransformation {
+
+    override val affectsIndex: Boolean
+        get() = false
+
+    override val presentation: String
+        get() = "mapIndexedNotNull{}"
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(mapping, indexVariable, inputVariable)
+        return chainedCallGenerator.generate("mapIndexedNotNull $0:'{}'", lambda)
+    }
+}

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

@@ -283,4 +283,27 @@ fun KtExpression.isStableInLoop(loop: KtLoopExpression, checkNoOtherUsagesInLoop
 
         else -> return false
     }
+}
+
+fun KtExpression.containsEmbeddedBreakOrContinue(): Boolean {
+    return anyDescendantOfType<KtExpressionWithLabel>(::isEmbeddedBreakOrContinue)
+}
+
+fun KtExpression.countEmbeddedBreaksAndContinues(): Int {
+    return collectDescendantsOfType<KtExpressionWithLabel>(::isEmbeddedBreakOrContinue).size
+}
+
+private fun isEmbeddedBreakOrContinue(expression: KtExpressionWithLabel): Boolean {
+    if (expression !is KtBreakExpression && expression !is KtContinueExpression) return false
+    val parent = expression.parent
+    when (parent) {
+        is KtBlockExpression -> return false
+
+        is KtContainerNode -> {
+            val containerExpression = parent.parent as KtExpression
+            return containerExpression.analyze(BodyResolveMode.PARTIAL)[BindingContext.USED_AS_EXPRESSION, containerExpression] == true
+        }
+
+        else -> return true
+    }
 }

idea/testData/intentions/loopToCallChain/embeddedBreak1.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String?>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        val length = s?.length ?: break
+        target.add(length.toString())
+    }
+}

idea/testData/intentions/loopToCallChain/embeddedBreak2.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        val length = if (s.isNotEmpty()) s.length else break
+        target.add(length.toString())
+    }
+}

idea/testData/intentions/loopToCallChain/embeddedContinue.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        val length = if (s.isNotEmpty()) s.length else continue
+        target.add(length.toString())
+    }
+}

idea/testData/intentions/loopToCallChain/mapIndexedNotNull_elvisContinue.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedNotNull{}.mapTo(){}'"
+fun foo(list: List<String?>, target: MutableList<String>) {
+    <caret>for ((index, s) in list.withIndex()) {
+        val length = s?.substring(index)?.length ?: continue
+        target.add(length.toString())
+    }
+}

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

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapIndexedNotNull{}.mapTo(){}'"
+fun foo(list: List<String?>, target: MutableList<String>) {
+    <caret>list
+            .mapIndexedNotNull { index, s -> s?.substring(index)?.length }
+            .mapTo(target) { it.toString() }
+}

idea/testData/intentions/loopToCallChain/mapNotNull_elvisContinue.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapNotNull{}.mapTo(){}'"
+fun foo(list: List<String?>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        val length = s?.length ?: continue
+        target.add(length.toString())
+    }
+}

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

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'mapNotNull{}.mapTo(){}'"
+fun foo(list: List<String?>, target: MutableList<String>) {
+    <caret>list
+            .mapNotNull { it?.length }
+            .mapTo(target) { it.toString() }
+}

idea/testData/intentions/loopToCallChain/mapNotNull_nestedLoopElvisContinue.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'flatMap{}.mapNotNull{}.mapTo(){}'"
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>for (s in list) {
+        for (i in s.indices) {
+            val v = bar(i) ?: continue
+            target.add(v.substring(1))
+        }
+    }
+}
+
+fun bar(p: Int): String? = null

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

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'flatMap{}.mapNotNull{}.mapTo(){}'"
+fun foo(list: List<String>, target: MutableList<String>) {
+    <caret>list
+            .flatMap { it.indices }
+            .mapNotNull { bar(it) }
+            .mapTo(target) { it.substring(1) }
+}
+
+fun bar(p: Int): String? = null

idea/testData/intentions/loopToCallChain/mapNotNull_wrongElvisContinue.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String>, target: MutableList<String>) {
+    Loop@
+    <caret>for (s in list) {
+        for (i in s.indices) {
+            val v = bar(i) ?: continue@Loop
+            target.add(v.substring(1))
+        }
+    }
+}
+
+fun bar(p: Int): String? = null