From 32c3cd4f48d27a0adf15a7c73f2498d3e2418fb5 Mon Sep 17 00:00:00 2001
From: mschauer <moritzschauer@web.de>
Date: Wed, 1 Jun 2016 21:54:50 +0200
Subject: [PATCH] Concat iterator that concatenates the elements of an iterator

---
 base/iterator.jl   | 57 ++++++++++++++++++++++++++++++++++++++++++++++
 test/functional.jl | 14 ++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/base/iterator.jl b/base/iterator.jl
index c3479f8fab74b..f48391f001b3a 100644
--- a/base/iterator.jl
+++ b/base/iterator.jl
@@ -559,3 +559,60 @@ function next(itr::PartitionIterator, state)
     end
     return resize!(v, i), state
 end
+
+
+immutable Concat{I,T,S}
+    first::T
+    itr::I
+    st::S
+end
+
+function Concat(c)
+    s = start(c)
+    done(c, s) && throw(ArgumentError("argument to Concat must contain at least one element"))
+    head, s = next(c, s)
+    Concat(head, c, s)
+end
+
+"""
+    concat(iter)
+
+Given an iterator with trait `HasShape()` that yields iterators with trait
+`HasShape()`, return an iterator that yields the elements of those iterators
+concatenated.
+
+```jldoctest
+julia> collect(Base.concat(((3*(i-1)+j for j in 1:3) for i in 1:4)))
+3×4 Array{Int64,2}:
+ 1  4  7  10
+ 2  5  8  11
+ 3  6  9  12
+```
+"""
+concat(it) = Concat(it)
+
+eltype{I,T,S}(::Type{Concat{I,T,S}}) = eltype(T)
+iteratorsize{I,T,S}(::Type{Concat{I,T,S}}) = HasShape()
+iteratoreltype{I,T,S}(::Type{Concat{I,T,S}}) = iteratoreltype(T)
+size(it::Concat) = (size(it.first)...,size(it.itr)...)
+length(it::Concat) = prod(size(it))
+
+function start(c::Concat)
+    return c.st, c.first, start(c.first)
+end
+
+function next(c::Concat, state)
+    s, inner, s2 = state
+    val, s2 = next(inner, s2)
+    while done(inner, s2) && !done(c.itr, s)
+        inner, s = next(c.itr, s)
+        size(inner) == size(c.first) || throw(DimensionMismatch("elements of different size in argument to Concat"))
+        s2 = start(inner)
+    end
+    return val, (s, inner, s2)
+end
+
+@inline function done(c::Concat, state)
+    s, inner, s2 = state
+    return done(c.itr, s) && done(inner, s2)
+end
diff --git a/test/functional.jl b/test/functional.jl
index 754442d8bfbeb..2d8d7a118f67c 100644
--- a/test/functional.jl
+++ b/test/functional.jl
@@ -374,6 +374,20 @@ import Base.flatten
 @test eltype(flatten(UnitRange{Int8}[1:2, 3:4])) == Int8
 @test_throws ArgumentError collect(flatten(Any[]))
 
+
+# concat
+# ------
+
+import Base.concat
+
+@test collect((concat((i:i+10 for i in 1:3)))) == [1:11 2:12 3:13]
+@test typeof(collect((concat((i:i+10 for i in 1:3))))) == Array{Int,2}
+@test_throws DimensionMismatch collect(concat((i:10 for i in 1:3)))
+@test_throws ArgumentError collect(flatten(Any[]))
+
+@test [reshape(1:6,3,2)[k,l]*i+j for k in 1:3, l in 1:2, i in 1:4, j in 1:5] == collect(Base.concat([reshape(1:6,3,2)*i+j for i in 1:4, j in 1:5]))
+
+
 # foreach
 let
     a = []