From e4bdcacaff00508052942a9df5db0d545f4a02a4 Mon Sep 17 00:00:00 2001
From: Stefan Karpinski <stefan@karpinski.org>
Date: Wed, 9 Oct 2019 14:11:39 -0400
Subject: [PATCH] unzip: the inverse of zip

---
 base/exports.jl   |  1 +
 base/iterators.jl | 65 ++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/base/exports.jl b/base/exports.jl
index 49063720c14c4..590063e2b3c5a 100644
--- a/base/exports.jl
+++ b/base/exports.jl
@@ -628,6 +628,7 @@ export
 
     enumerate,  # re-exported from Iterators
     zip,
+    unzip,
     only,
 
 # object identity and equality
diff --git a/base/iterators.jl b/base/iterators.jl
index 9c813d7e21632..b30099eef87f6 100644
--- a/base/iterators.jl
+++ b/base/iterators.jl
@@ -22,7 +22,8 @@ import .Base:
     getindex, setindex!, get, iterate,
     popfirst!, isdone, peek
 
-export enumerate, zip, rest, countfrom, take, drop, cycle, repeated, product, flatten, partition
+export enumerate, zip, unzip, rest, countfrom, take, drop,
+    cycle, repeated, product, flatten, partition
 
 tail_if_any(::Tuple{}) = ()
 tail_if_any(x::Tuple) = tail(x)
@@ -390,6 +391,68 @@ _zip_iterator_eltype(::Type{Tuple{}}) = HasEltype()
 
 reverse(z::Zip) = Zip(map(reverse, z.is))
 
+# unzip
+
+"""
+    unzip(itrs) -> Vector{<:Vector}
+
+The `unzip` function takes an iterator of iterators and returns a vector of
+vectors such that the first vector contains the first element yielded by each
+iterator, the second vector the second element yielded by each iterator, etc.
+`unzip` is sort of an inverse to the `zip` operation, as the name suggests.
+In particular, if we define
+
+    ≐(a, b) = collect(collect.(a)) == collect(collect.(b))
+
+Then the following identities relating `zip` and `unzip` hold:
+
+    unzip(zip(itrs...)) ≐ itrs
+
+    zip(unzip(itrs)...) ≐ itrs
+
+Note that `unzip` does not return an iterator: it always consumes all of
+its argument and all of each iterator yielded by its argument. It is only
+associated with iteration beacuse it is the inverse of `zip`.
+
+# Examples
+
+```jldoctest
+julia> unzip(enumerate("Hello"))
+2-element Array{Array{T,1} where T,1}:
+ [1, 2, 3]
+ ['a', 'b', 'c']
+
+julia> unzip([[1, 'a'], [2.5, 'z'], [0, 'x']])
+2-element Array{Array{T,1} where T,1}:
+ Real[1, 2.5, 0]
+ ['a', 'z', 'x']
+```
+"""
+function unzip(itrs)
+    n = Base.haslength(itrs) ? length(itrs) : -1
+    vecs = Vector[]
+    for itr in itrs
+        for (j, x) in enumerate(itr)
+            if length(vecs) < j
+                v = [x]
+                push!(vecs, v)
+                n ≥ 0 && sizehint!(v, n)
+            else
+                v = vecs[j]
+                if !(x isa eltype(v))
+                    T = Base.promote_typejoin(typeof(x), eltype(v))
+                    v = vecs[j] = copyto!(similar(v, T), v)
+                    n ≥ 0 && sizehint!(v, n)
+                end
+                push!(v, x)
+            end
+        end
+        length(first(vecs)) == length(last(vecs)) ||
+            throw(ArgumentError("unzip called with uneven iterators"))
+    end
+    return vecs
+end
+
 # filter
 
 struct Filter{F,I}