KT-14336 for to stdlib to takeWhile: propose conversion into takeWhile if break is located in else branch

#KT-14336 Fixed

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

@@ -155,16 +155,41 @@ abstract class FilterTransformationBase : SequenceTransformation {
 
         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
-            val then = ifStatement.then ?: return null
+            val thenBranch = ifStatement.then ?: return null
+            val elseBranch = ifStatement.`else`
 
+            if (elseBranch == null) {
+                return matchOneTransformation(state, condition, false, thenBranch, state.statements.drop(1))
+            }
+            else if (state.statements.size == 1) {
+                val thenStatement = thenBranch.blockExpressionsOrSingle().singleOrNull()
+                if (thenStatement is KtBreakExpression || thenStatement is KtContinueExpression) {
+                    return matchOneTransformation(state, condition, false, thenBranch, elseBranch.singletonList())
+                }
+
+                val elseStatement = elseBranch.blockExpressionsOrSingle().singleOrNull()
+                if (elseStatement is KtBreakExpression || elseStatement is KtContinueExpression) {
+                    return matchOneTransformation(state, condition, true, elseBranch, thenBranch.singletonList())
+                }
+            }
+
+            return null
+        }
+
+        private fun matchOneTransformation(
+                state: MatchingState,
+                condition: KtExpression,
+                negateCondition: Boolean,
+                then: KtExpression,
+                restStatements: List<KtExpression>
+        ): Pair<SequenceTransformation, MatchingState>? {
             // we do not allow filter() which uses neither input variable nor index variable (though is technically possible but looks confusing)
             // shouldUseInputVariables = false does not work for us because we sometimes return Result match in this matcher
             if (!state.inputVariable.hasUsages(condition) && (state.indexVariable == null || !state.indexVariable.hasUsages(condition))) return null
 
-            if (state.statements.size == 1) {
+            if (restStatements.isEmpty()) {
-                val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, Condition.create(condition))
+                val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, Condition.create(condition, negateCondition))
                 val newState = state.copy(statements = listOf(then))
                 return transformation to newState
             }
@@ -173,15 +198,15 @@ abstract class FilterTransformationBase : SequenceTransformation {
                 when (statement) {
                     is KtContinueExpression -> {
                         if (statement.targetLoop() != state.innerLoop) return null
-                        val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, Condition.create(condition, negated = true))
+                        val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, Condition.create(condition, !negateCondition))
-                        val newState = state.copy(statements = state.statements.drop(1))
+                        val newState = state.copy(statements = restStatements)
                         return transformation to newState
                     }
 
                     is KtBreakExpression -> {
                         if (statement.targetLoop() != state.outerLoop) return null
-                        val transformation = TakeWhileTransformation(state.outerLoop, state.inputVariable, condition.negate())
+                        val transformation = TakeWhileTransformation(state.outerLoop, state.inputVariable, if (negateCondition) condition else condition.negate())
-                        val newState = state.copy(statements = state.statements.drop(1))
+                        val newState = state.copy(statements = restStatements)
                         return transformation to newState
                     }
 

idea/testData/intentions/loopToCallChain/filter/ifElse.kt

@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'firstOrNull{}'"
+// IS_APPLICABLE_2: false
+fun foo(list: List<String>): String? {
+    <caret>for (s in list) {
+        if (s.isEmpty()) continue
+        else return s
+    }
+    return null
+}

idea/testData/intentions/loopToCallChain/filter/ifElse.kt.after

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'firstOrNull{}'"
+// IS_APPLICABLE_2: false
+fun foo(list: List<String>): String? {
+    return list.firstOrNull { !it.isEmpty() }
+}

idea/testData/intentions/loopToCallChain/takeWhile/ifElse1.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'takeWhile{}'"
+// IS_APPLICABLE_2: false
+fun foo(n: Int, list: List<Int>): List<Int>{
+    val result = mutableListOf<Int>()
+    <caret>for (i in list) {
+        if (i >= n)
+            result.add(i)
+        else
+            break
+    }
+    return result
+}

idea/testData/intentions/loopToCallChain/takeWhile/ifElse1.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'takeWhile{}'"
+// IS_APPLICABLE_2: false
+fun foo(n: Int, list: List<Int>): List<Int>{
+    val result = list.takeWhile { it >= n }
+    return result
+}

idea/testData/intentions/loopToCallChain/takeWhile/ifElse2.kt

@@ -0,0 +1,13 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'takeWhile{}'"
+// IS_APPLICABLE_2: false
+fun foo(n: Int, list: List<Int>): List<Int>{
+    val result = mutableListOf<Int>()
+    <caret>for (i in list) {
+        if (i < n)
+            break
+        else
+            result.add(i)
+    }
+    return result
+}

idea/testData/intentions/loopToCallChain/takeWhile/ifElse2.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with 'takeWhile{}'"
+// IS_APPLICABLE_2: false
+fun foo(n: Int, list: List<Int>): List<Int>{
+    val result = list.takeWhile { it >= n }
+    return result
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTest2Generated.java

@@ -574,6 +574,12 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             doTest(fileName);
         }
 
+        @TestMetadata("ifElse.kt")
+        public void testIfElse() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/filter/ifElse.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("inputVarNotUsed.kt")
         public void testInputVarNotUsed() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/filter/inputVarNotUsed.kt");
@@ -1444,6 +1450,18 @@ public class IntentionTest2Generated extends AbstractIntentionTest2 {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/loopToCallChain/takeWhile"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
         }
 
+        @TestMetadata("ifElse1.kt")
+        public void testIfElse1() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/takeWhile/ifElse1.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifElse2.kt")
+        public void testIfElse2() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/takeWhile/ifElse2.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("nestedLoop.kt")
         public void testNestedLoop() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/takeWhile/nestedLoop.kt");

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -8482,6 +8482,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("ifElse.kt")
+            public void testIfElse() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/filter/ifElse.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("inputVarNotUsed.kt")
             public void testInputVarNotUsed() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/filter/inputVarNotUsed.kt");
@@ -9352,6 +9358,18 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
                 KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/loopToCallChain/takeWhile"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), true);
             }
 
+            @TestMetadata("ifElse1.kt")
+            public void testIfElse1() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/takeWhile/ifElse1.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("ifElse2.kt")
+            public void testIfElse2() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/takeWhile/ifElse2.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("nestedLoop.kt")
             public void testNestedLoop() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/loopToCallChain/takeWhile/nestedLoop.kt");