From 95d4accda78a7fd8a412316483bca3ac03fafa13 Mon Sep 17 00:00:00 2001
From: Ilya Ryzhenkov <orangy@jetbrains.com>
Date: Tue, 13 Jan 2015 19:49:12 +0300
Subject: [PATCH] stdlib: add flatten()

---
 .../org/jetbrains/kotlin/utils/collections.kt |  4 ---
 ...MoveKotlinTopLevelDeclarationsProcessor.kt |  3 +-
 .../src/kotlin/collections/Operations.kt      | 34 +++++++++++++++++++
 .../stdlib/test/collections/IterableTests.kt  |  5 +++
 .../stdlib/test/collections/StreamTest.kt     |  7 ++++
 5 files changed, 47 insertions(+), 6 deletions(-)
 create mode 100644 libraries/stdlib/src/kotlin/collections/Operations.kt

diff --git a/core/util.runtime/src/org/jetbrains/kotlin/utils/collections.kt b/core/util.runtime/src/org/jetbrains/kotlin/utils/collections.kt
index c2e1047da14..09cdc6c3b83 100644
--- a/core/util.runtime/src/org/jetbrains/kotlin/utils/collections.kt
+++ b/core/util.runtime/src/org/jetbrains/kotlin/utils/collections.kt
@@ -79,10 +79,6 @@ public fun <K, V: Any> Iterable<K>.keysToMapExceptNulls(value: (K) -> V?): Map<K
 
 public fun <T, C: Collection<T>> C.ifEmpty(body: () -> C): C = if (isEmpty()) body() else this
 
-public fun <T> Iterable<Iterable<T>>.flatten(): List<T> {
-    return flatMapTo(ArrayList<T>(), {it})
-}
-
 public fun <T: Any> emptyOrSingletonList(item: T?): List<T> = if (item == null) listOf() else listOf(item)
 
 public fun <T: Any> MutableCollection<T>.addIfNotNull(t: T?) {
diff --git a/idea/src/org/jetbrains/kotlin/idea/refactoring/move/moveTopLevelDeclarations/MoveKotlinTopLevelDeclarationsProcessor.kt b/idea/src/org/jetbrains/kotlin/idea/refactoring/move/moveTopLevelDeclarations/MoveKotlinTopLevelDeclarationsProcessor.kt
index dabf5a5a69a..ece86647d65 100644
--- a/idea/src/org/jetbrains/kotlin/idea/refactoring/move/moveTopLevelDeclarations/MoveKotlinTopLevelDeclarationsProcessor.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/refactoring/move/moveTopLevelDeclarations/MoveKotlinTopLevelDeclarationsProcessor.kt
@@ -43,7 +43,6 @@ import org.jetbrains.kotlin.idea.refactoring.move.getFileNameAfterMove
 import org.jetbrains.kotlin.psi.JetNamedDeclaration
 import org.jetbrains.kotlin.asJava.toLightElements
 import java.util.HashMap
-import org.jetbrains.kotlin.utils.flatten
 import java.util.ArrayList
 import java.util.HashSet
 import com.intellij.psi.PsiReference
@@ -119,7 +118,7 @@ public class MoveKotlinTopLevelDeclarationsProcessor(
         val newPackageName = options.moveTarget.packageWrapper?.getQualifiedName() ?: ""
 
         fun collectUsages(): List<UsageInfo> {
-            return kotlinToLightElements.values().flatten().flatMap { lightElement ->
+            return kotlinToLightElements.values().flatMap { it }.flatMap { lightElement ->
                 val newFqName = StringUtil.getQualifiedName(newPackageName, lightElement.getName())
 
                 val foundReferences = HashSet<PsiReference>()
diff --git a/libraries/stdlib/src/kotlin/collections/Operations.kt b/libraries/stdlib/src/kotlin/collections/Operations.kt
new file mode 100644
index 00000000000..44b0a29dd06
--- /dev/null
+++ b/libraries/stdlib/src/kotlin/collections/Operations.kt
@@ -0,0 +1,34 @@
+package kotlin
+
+import java.util.ArrayList
+import kotlin.platform.platformName
+
+/**
+ * Returns a single list of all elements from all collections in original collection
+ */
+public fun <T> Iterable<Iterable<T>>.flatten(): List<T> {
+    val result = ArrayList<T>()
+    for (element in this) {
+        result.addAll(element)
+    }
+    return result
+}
+
+/**
+ * Returns a stream of all elements from all streams in original stream
+ */
+public fun <T> Stream<Stream<T>>.flatten(): Stream<T> {
+    return FlatteningStream(this, { it })
+}
+
+/**
+ * Returns a single list of all elements from all collections in original collection
+ */
+public fun <T> Array<Array<out T>>.flatten(): List<T> {
+    val result = ArrayList<T>(sumBy { it.size() })
+    for (element in this) {
+        result.addAll(element)
+    }
+    return result
+}
+
diff --git a/libraries/stdlib/test/collections/IterableTests.kt b/libraries/stdlib/test/collections/IterableTests.kt
index b034c1a09d1..d6d61b87033 100644
--- a/libraries/stdlib/test/collections/IterableTests.kt
+++ b/libraries/stdlib/test/collections/IterableTests.kt
@@ -164,6 +164,11 @@ abstract class IterableTests<T : Iterable<String>>(val data: T, val empty: T) {
         assertEquals(listOf(3, 3), lengths)
     }
 
+    Test
+    fun flatten() {
+        assertEquals(listOf(0, 1, 2, 3, 0, 1, 2, 3), data.map { 0..it.length() }.flatten())
+    }
+
     Test
     fun mapIndexed() {
         val shortened = data.mapIndexed {(index, value) -> value.substring(0..index) }
diff --git a/libraries/stdlib/test/collections/StreamTest.kt b/libraries/stdlib/test/collections/StreamTest.kt
index e379dbc7cf2..35f3be9493e 100644
--- a/libraries/stdlib/test/collections/StreamTest.kt
+++ b/libraries/stdlib/test/collections/StreamTest.kt
@@ -213,6 +213,13 @@ public class StreamTest {
         assertEquals(listOf(0, 1, 3, 4), result.toList())
     }
 
+    test
+    fun flatten() {
+        val data = streamOf(1, 2).map { streamOf(0..it) }
+        assertEquals(listOf(0, 1, 0, 1, 2), data.flatten().toList())
+    }
+
+
     /*
     test fun pairIterator() {
         val pairStr = (fibonacci() zip fibonacci().map { i -> i*2 }).joinToString(limit = 10)