From 3f563f7058cfe5aefa20d56c6713518a0fbb0a24 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Thu, 21 Apr 2016 22:06:09 +0300
Subject: [PATCH] takeWhile supported

---
 .../sequence/FilterTransformation.kt          | 48 +++++++++++++++++--
 .../intentions/loopToCallChain/takeWhile.kt   |  8 ++++
 .../loopToCallChain/takeWhile.kt.after        |  5 ++
 .../loopToCallChain/takeWhile_nestedLoop.kt   | 11 +++++
 .../takeWhile_nestedLoop.kt.after             |  7 +++
 .../loopToCallChain/takeWhile_wrongBreak.kt   | 10 ++++
 6 files changed, 84 insertions(+), 5 deletions(-)
 create mode 100644 idea/testData/intentions/loopToCallChain/takeWhile.kt
 create mode 100644 idea/testData/intentions/loopToCallChain/takeWhile.kt.after
 create mode 100644 idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt
 create mode 100644 idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt.after
 create mode 100644 idea/testData/intentions/loopToCallChain/takeWhile_wrongBreak.kt

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt
index 7b568b74ed0..ac41ade66f6 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/loopToCallChain/sequence/FilterTransformation.kt
@@ -67,6 +67,12 @@ class FilterTransformation(
      *         if (<condition>) continue
      *         ...
      *     }
+     * or
+     *
+     *     for (...) {
+     *         if (<condition>) break
+     *         ...
+     *     }
      */
     object Matcher : SequenceTransformationMatcher {
         override fun match(state: MatchingState): SequenceTransformationMatch? {
@@ -81,11 +87,24 @@ class FilterTransformation(
                 return SequenceTransformationMatch(transformation, newState)
             }
             else {
-                val continueExpression = then.blockExpressionsOrSingle().singleOrNull() as? KtContinueExpression ?: return null
-                if (continueExpression.targetLoop() != state.innerLoop) return null
-                val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, condition, isInverse = true)
-                val newState = state.copy(statements = state.statements.drop(1))
-                return SequenceTransformationMatch(transformation, newState)
+                val statement = then.blockExpressionsOrSingle().singleOrNull() ?: return null
+                when (statement) {
+                    is KtContinueExpression -> {
+                        if (statement.targetLoop() != state.innerLoop) return null
+                        val transformation = createFilterTransformation(state.outerLoop, state.inputVariable, state.indexVariable, condition, isInverse = true)
+                        val newState = state.copy(statements = state.statements.drop(1))
+                        return SequenceTransformationMatch(transformation, newState)
+                    }
+
+                    is KtBreakExpression -> {
+                        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))
+                        return SequenceTransformationMatch(transformation, newState)
+                    }
+
+                    else -> return null
+                }
             }
         }
 
@@ -177,3 +196,22 @@ class FilterIndexedTransformation(
     }
 }
 
+class TakeWhileTransformation(
+        override val loop: KtForExpression,
+        val inputVariable: KtCallableDeclaration,
+        val condition: KtExpression
+) : SequenceTransformation {
+
+    //TODO: merge multiple
+
+    override val affectsIndex: Boolean
+        get() = false
+
+    override val presentation: String
+        get() = "takeWhile{}"
+
+    override fun generateCode(chainedCallGenerator: ChainedCallGenerator): KtExpression {
+        val lambda = generateLambda(inputVariable, condition)
+        return chainedCallGenerator.generate("takeWhile$0:'{}'", lambda)
+    }
+}
diff --git a/idea/testData/intentions/loopToCallChain/takeWhile.kt b/idea/testData/intentions/loopToCallChain/takeWhile.kt
new file mode 100644
index 00000000000..91d7368f45b
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/takeWhile.kt
@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with '+= takeWhile{}'"
+fun foo(list: List<String>, target: MutableCollection<String>) {
+    <caret>for (s in list) {
+        if (s.isEmpty()) break
+        target.add(s)
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/takeWhile.kt.after b/idea/testData/intentions/loopToCallChain/takeWhile.kt.after
new file mode 100644
index 00000000000..f2d1c963b8a
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/takeWhile.kt.after
@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with '+= takeWhile{}'"
+fun foo(list: List<String>, target: MutableCollection<String>) {
+    <caret>target += list.takeWhile { !it.isEmpty() }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt b/idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt
new file mode 100644
index 00000000000..a7922bcb5cb
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt
@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with '+= flatMap{}.takeWhile{}'"
+fun foo(list: List<String>, target: MutableCollection<Int>) {
+    Outer@
+    <caret>for (s in list) {
+        for (i in s.indices) {
+            if (i > 1000) break@Outer
+            target.add(i)
+        }
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt.after b/idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt.after
new file mode 100644
index 00000000000..5121e06e5ff
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/takeWhile_nestedLoop.kt.after
@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+// INTENTION_TEXT: "Replace with '+= flatMap{}.takeWhile{}'"
+fun foo(list: List<String>, target: MutableCollection<Int>) {
+    <caret>target += list
+            .flatMap { it.indices }
+            .takeWhile { it <= 1000 }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/takeWhile_wrongBreak.kt b/idea/testData/intentions/loopToCallChain/takeWhile_wrongBreak.kt
new file mode 100644
index 00000000000..4f45bbe5f96
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/takeWhile_wrongBreak.kt
@@ -0,0 +1,10 @@
+// WITH_RUNTIME
+// IS_APPLICABLE: false
+fun foo(list: List<String>, target: MutableCollection<Int>) {
+    <caret>for (s in list) {
+        for (i in s.indices) {
+            if (i > 1000) break
+            target.add(i)
+        }
+    }
+}
\ No newline at end of file