From da3ec891d0b476b26881bd122dcb63c24d42b711 Mon Sep 17 00:00:00 2001
From: Ilya Gorbunov <ilya.gorbunov@jetbrains.com>
Date: Tue, 14 Jul 2015 16:23:32 +0300
Subject: [PATCH] Provide `unzip` method for Iterables, Arrays and Sequences of
 pairs. #KT-5793 Fixed

---
 .../src/kotlin/collections/Operations.kt      | 45 +++++++++++++++++++
 .../stdlib/test/collections/CollectionTest.kt | 13 ++++++
 .../stdlib/test/collections/SequenceTest.kt   |  9 +++-
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/libraries/stdlib/src/kotlin/collections/Operations.kt b/libraries/stdlib/src/kotlin/collections/Operations.kt
index 0a2b757248d..bbb32f03eb6 100644
--- a/libraries/stdlib/src/kotlin/collections/Operations.kt
+++ b/libraries/stdlib/src/kotlin/collections/Operations.kt
@@ -32,3 +32,48 @@ public fun <T> Array<Array<out T>>.flatten(): List<T> {
     return result
 }
 
+/**
+ * Returns a pair of lists, where
+ * *first* list is built from the first values of each pair from this collection,
+ * *second* list is built from the second values of each pair from this collection.
+ */
+public fun <T, R> Iterable<Pair<T, R>>.unzip(): Pair<List<T>, List<R>> {
+    val expectedSize = collectionSizeOrDefault(10)
+    val listT = ArrayList<T>(expectedSize)
+    val listR = ArrayList<R>(expectedSize)
+    for (pair in this) {
+        listT.add(pair.first)
+        listR.add(pair.second)
+    }
+    return listT to listR
+}
+
+/**
+ * Returns a pair of lists, where
+ * *first* list is built from the first values of each pair from this array,
+ * *second* list is built from the second values of each pair from this array.
+ */
+public fun <T, R> Array<Pair<T, R>>.unzip(): Pair<List<T>, List<R>> {
+    val listT = ArrayList<T>(size())
+    val listR = ArrayList<R>(size())
+    for (pair in this) {
+        listT.add(pair.first)
+        listR.add(pair.second)
+    }
+    return listT to listR
+}
+
+/**
+ * Returns a pair of lists, where
+ * *first* list is built from the first values of each pair from this sequence,
+ * *second* list is built from the second values of each pair from this sequence.
+ */
+public fun <T, R> Sequence<Pair<T, R>>.unzip(): Pair<List<T>, List<R>> {
+    val listT = ArrayList<T>()
+    val listR = ArrayList<R>()
+    for (pair in this) {
+        listT.add(pair.first)
+        listR.add(pair.second)
+    }
+    return listT to listR
+}
\ No newline at end of file
diff --git a/libraries/stdlib/test/collections/CollectionTest.kt b/libraries/stdlib/test/collections/CollectionTest.kt
index 59f593718a5..a213fd01634 100644
--- a/libraries/stdlib/test/collections/CollectionTest.kt
+++ b/libraries/stdlib/test/collections/CollectionTest.kt
@@ -621,6 +621,19 @@ class CollectionTest {
         assertEquals(e, 5)
     }
 
+    test fun unzipList() {
+        val list = listOf(1 to 'a', 2 to 'b', 3 to 'c')
+        val (ints, chars) = list.unzip()
+        assertEquals(listOf(1, 2, 3), ints)
+        assertEquals(listOf('a', 'b', 'c'), chars)
+    }
+
+    test fun unzipArray() {
+        val array = arrayOf(1 to 'a', 2 to 'b', 3 to 'c')
+        val (ints, chars) = array.unzip()
+        assertEquals(listOf(1, 2, 3), ints)
+        assertEquals(listOf('a', 'b', 'c'), chars)
+    }
 
     test fun specialLists() {
         compare(arrayListOf<Int>(), listOf<Int>()) { listBehavior() }
diff --git a/libraries/stdlib/test/collections/SequenceTest.kt b/libraries/stdlib/test/collections/SequenceTest.kt
index a1739f3e988..d14b6b87a61 100644
--- a/libraries/stdlib/test/collections/SequenceTest.kt
+++ b/libraries/stdlib/test/collections/SequenceTest.kt
@@ -291,8 +291,7 @@ public class SequenceTest {
         assertEquals(listOf(0, 1, 3, 4), result.toList())
     }
 
-    test
-    fun flatten() {
+    test fun flatten() {
         val data = sequenceOf(1, 2).map { sequenceOf(0..it) }
         assertEquals(listOf(0, 1, 0, 1, 2), data.flatten().toList())
     }
@@ -307,6 +306,12 @@ public class SequenceTest {
         assertEquals(listOf(13, 34, 55, 144), sequence.distinctBy { it % 4 }.toList())
     }
 
+    test fun unzip() {
+        val seq = sequenceOf(1 to 'a', 2 to 'b', 3 to 'c')
+        val (ints, chars) = seq.unzip()
+        assertEquals(listOf(1, 2, 3), ints)
+        assertEquals(listOf('a', 'b', 'c'), chars)
+    }
 
     /*
     test fun pairIterator() {