From 8c7cdf1bf8de5d582ffb9ae3215b72491ffbae7f Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Fri, 22 Apr 2016 11:48:10 +0300
Subject: [PATCH] Supported other patterns for mapNotNull

---
 .../loopToCallChain/sequence/FilterTransformation.kt     | 7 +++++++
 .../intentions/loopToCallChain/mapIndexedNotNull.kt      | 9 +++++++++
 .../loopToCallChain/mapIndexedNotNull.kt.after           | 7 +++++++
 idea/testData/intentions/loopToCallChain/mapNotNull.kt   | 9 +++++++++
 .../intentions/loopToCallChain/mapNotNull.kt.after       | 7 +++++++
 5 files changed, 39 insertions(+)
 create mode 100644 idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt
 create mode 100644 idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt.after
 create mode 100644 idea/testData/intentions/loopToCallChain/mapNotNull.kt
 create mode 100644 idea/testData/intentions/loopToCallChain/mapNotNull.kt.after

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 ac41ade66f6..28b8048d51e 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
@@ -163,6 +163,13 @@ class FilterIsInstanceTransformation(
 }
 
 class FilterNotNullTransformation(override val loop: KtForExpression) : SequenceTransformation {
+    override fun mergeWithPrevious(previousTransformation: SequenceTransformation): SequenceTransformation? {
+        return when (previousTransformation) {
+            is MapTransformation -> MapNotNullTransformation(loop, previousTransformation.inputVariable, previousTransformation.mapping)
+            is MapIndexedTransformation -> MapIndexedNotNullTransformation(loop, previousTransformation.inputVariable, previousTransformation.indexVariable, previousTransformation.mapping)
+            else -> null
+        }
+    }
 
     override val affectsIndex: Boolean
         get() = true
diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt b/idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt
new file mode 100644
index 00000000000..57380912bd9
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt
@@ -0,0 +1,9 @@
+// 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
+        if (length == null) continue
+        target.add(length.toString())
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt.after b/idea/testData/intentions/loopToCallChain/mapIndexedNotNull.kt.after
new file mode 100644
index 00000000000..1f5245bb3eb
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/mapIndexedNotNull.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() }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/mapNotNull.kt b/idea/testData/intentions/loopToCallChain/mapNotNull.kt
new file mode 100644
index 00000000000..652693bc8df
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/mapNotNull.kt
@@ -0,0 +1,9 @@
+// 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
+        if (length == null) continue
+        target.add(length.toString())
+    }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/loopToCallChain/mapNotNull.kt.after b/idea/testData/intentions/loopToCallChain/mapNotNull.kt.after
new file mode 100644
index 00000000000..b9a82c98c29
--- /dev/null
+++ b/idea/testData/intentions/loopToCallChain/mapNotNull.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() }
+}
\ No newline at end of file