make map and collect more general and uniform by adding iterator traits

Author: JeffBezanson

base/abstractarray.jl

@@ -1087,34 +1087,6 @@ foreach(f) = (f(); nothing)
 foreach(f, itr) = (for x in itr; f(x); end; nothing)
 foreach(f, itrs...) = (for z in zip(itrs...); f(z...); end; nothing)
 
-# generic map on any iterator
-function map(f, iters...)
-    result = []
-    len = length(iters)
-    states = [start(iters[idx]) for idx in 1:len]
-    nxtvals = cell(len)
-    cont = true
-    for idx in 1:len
-        if done(iters[idx], states[idx])
-            cont = false
-            break
-        end
-    end
-    while cont
-        for idx in 1:len
-            nxtvals[idx],states[idx] = next(iters[idx], states[idx])
-        end
-        push!(result, f(nxtvals...))
-        for idx in 1:len
-            if done(iters[idx], states[idx])
-                cont = false
-                break
-            end
-        end
-    end
-    result
-end
-
 ## map over arrays ##
 
 ## transform any set of dimensions
@@ -1175,39 +1147,6 @@ function mapslices(f, A::AbstractArray, dims::AbstractVector)
     return R
 end
 
-
-# using promote_type
-function promote_to!{T,F}(f::F, offs, dest::AbstractArray{T}, A::AbstractArray)
-    # map to dest array, checking the type of each result. if a result does not
-    # match, do a type promotion and re-dispatch.
-    for i = offs:length(A)
-        @inbounds Ai = A[i]
-        el = f(Ai)
-        S = typeof(el)
-        if S === T || S <: T
-            @inbounds dest[i] = el::T
-        else
-            R = promote_type(T, S)
-            if R !== T
-                new = similar(dest, R)
-                copy!(new,1, dest,1, i-1)
-                new[i] = el
-                return promote_to!(f, i+1, new, A)
-            end
-            @inbounds dest[i] = el
-        end
-    end
-    return dest
-end
-
-function map_promote(f, A::AbstractArray)
-    if isempty(A); return similar(A, Bottom); end
-    first = f(A[1])
-    dest = similar(A, typeof(first))
-    dest[1] = first
-    return promote_to!(f, 2, dest, A)
-end
-
 # These are needed because map(eltype, As) is not inferrable
 promote_eltype_op(::Any) = (@_pure_meta; Bottom)
 promote_eltype_op{T}(op, ::AbstractArray{T}) = (@_pure_meta; promote_op(op, T))
@@ -1226,39 +1165,7 @@ function map!{F}(f::F, dest::AbstractArray, A::AbstractArray)
     return dest
 end
 
-function map_to!{T,F}(f::F, offs, st, dest::AbstractArray{T}, A)
-    # map to dest array, checking the type of each result. if a result does not
-    # match, widen the result type and re-dispatch.
-    i = offs
-    while !done(A, st)
-        @inbounds Ai, st = next(A, st)
-        el = f(Ai)
-        S = typeof(el)
-        if S === T || S <: T
-            @inbounds dest[i] = el::T
-            i += 1
-        else
-            R = typejoin(T, S)
-            new = similar(dest, R)
-            copy!(new,1, dest,1, i-1)
-            @inbounds new[i] = el
-            return map_to!(f, i+1, st, new, A)
-        end
-    end
-    return dest
-end
-
-function map(f, A::AbstractArray)
-    if isempty(A)
-        return isa(f,Type) ? similar(A,f) : similar(A)
-    end
-    st = start(A)
-    A1, st = next(A, st)
-    first = f(A1)
-    dest = similar(A, typeof(first))
-    dest[1] = first
-    return map_to!(f, 2, st, dest, A)
-end
+map{F}(f::F, A::AbstractArray) = collect(Generator(f,A))
 
 ## 2 argument
 function map!{F}(f::F, dest::AbstractArray, A::AbstractArray, B::AbstractArray)
@@ -1268,34 +1175,6 @@ function map!{F}(f::F, dest::AbstractArray, A::AbstractArray, B::AbstractArray)
     return dest
 end
 
-function map_to!{T,F}(f::F, offs, dest::AbstractArray{T}, A::AbstractArray, B::AbstractArray)
-    for i = offs:length(A)
-        @inbounds Ai, Bi = A[i], B[i]
-        el = f(Ai, Bi)
-        S = typeof(el)
-        if (S !== T) && !(S <: T)
-            R = typejoin(T, S)
-            new = similar(dest, R)
-            copy!(new,1, dest,1, i-1)
-            @inbounds new[i] = el
-            return map_to!(f, i+1, new, A, B)
-        end
-        @inbounds dest[i] = el::T
-    end
-    return dest
-end
-
-function map(f, A::AbstractArray, B::AbstractArray)
-    shp = promote_shape(size(A),size(B))
-    if prod(shp) == 0
-        return similar(A, promote_type(eltype(A),eltype(B)), shp)
-    end
-    first = f(A[1], B[1])
-    dest = similar(A, typeof(first), shp)
-    dest[1] = first
-    return map_to!(f, 2, dest, A, B)
-end
-
 ## N argument
 
 ith_all(i, ::Tuple{}) = ()
@@ -1311,32 +1190,9 @@ end
 
 map!{F}(f::F, dest::AbstractArray, As::AbstractArray...) = map_n!(f, dest, As)
 
-function map_to_n!{T,F}(f::F, offs, dest::AbstractArray{T}, As)
-    for i = offs:length(As[1])
-        el = f(ith_all(i, As)...)
-        S = typeof(el)
-        if (S !== T) && !(S <: T)
-            R = typejoin(T, S)
-            new = similar(dest, R)
-            copy!(new,1, dest,1, i-1)
-            @inbounds new[i] = el
-            return map_to_n!(f, i+1, new, As)
-        end
-        @inbounds dest[i] = el::T
-    end
-    return dest
-end
+spread(f) = (args)->f(args...)
 
-function map(f, As::AbstractArray...)
-    shape = mapreduce(size, promote_shape, As)
-    if prod(shape) == 0
-        return similar(As[1], promote_eltype(As...), shape)
-    end
-    first = f(map(a->a[1], As)...)
-    dest = similar(As[1], typeof(first), shape)
-    dest[1] = first
-    return map_to_n!(f, 2, dest, As)
-end
+map(f, iters...) = collect(Generator(spread(f),zip(iters...)))
 
 # multi-item push!, unshift! (built on top of type-specific 1-item version)
 # (note: must not cause a dispatch loop when 1-item case is not defined)

base/array.jl

@@ -204,30 +204,93 @@ promote_rule{T,n,S}(::Type{Array{T,n}}, ::Type{Array{S,n}}) = Array{promote_type
 
 Return an array of type `Array{element_type,1}` of all items in a collection.
 """
-function collect{T}(::Type{T}, itr)
-    if applicable(length, itr)
-        # when length() isn't defined this branch might pollute the
-        # type of the other.
-        a = Array(T,length(itr)::Integer)
-        i = 0
-        for x in itr
-            a[i+=1] = x
-        end
-    else
-        a = Array(T,0)
-        for x in itr
-            push!(a,x)
-        end
-    end
-    return a
-end
+collect{T}(::Type{T}, itr) = _collect_t(T, itr, iteratorsize(itr))
 
 """
     collect(collection)
 
 Return an array of all items in a collection. For associative collections, returns Pair{KeyType, ValType}.
 """
-collect(itr) = collect(eltype(itr), itr)
+collect(itr) = _collect(itr, iteratoreltype(itr), iteratorsize(itr))
+
+_collect(itr, ::HasEltype, isz) = _collect_t(eltype(itr), itr, isz)
+
+_collect_t(T::Type, itr, ::HasLength) = copy!(Array(T,Int(length(itr)::Integer)), itr)
+_collect_t(T::Type, itr, ::HasShape) = copy!(Array(T,convert(Dims,size(itr))), itr)
+function _collect_t(T::Type, itr, ::SizeUnknown)
+    a = Array(T,0)
+    for x in itr
+        push!(a,x)
+    end
+    return a
+end
+
+_collect(itr, ::EltypeUnknown, ::HasLength) = _collect_shaped(itr, (Int(length(itr)),))
+_collect(itr, ::EltypeUnknown, ::HasShape) = _collect_shaped(itr, convert(Dims,size(itr)))
+
+_default_container(itr, elty, sz) = Array(elty, sz)
+_default_container(itr::AbstractArray, elty, sz) = similar(itr, elty, sz)
+
+function collect_to!{T}(itr, offs, st, dest::AbstractArray{T})
+    # collect to dest array, checking the type of each result. if a result does not
+    # match, widen the result type and re-dispatch.
+    i = offs
+    while !done(itr, st)
+        el, st = next(itr, st)
+        S = typeof(el)
+        if S === T || S <: T
+            @inbounds dest[i] = el::T
+            i += 1
+        else
+            R = typejoin(T, S)
+            new = similar(dest, R)
+            copy!(new,1, dest,1, i-1)
+            @inbounds new[i] = el
+            return collect_to!(itr, i+1, st, new)
+        end
+    end
+    return dest
+end
+
+function _collect_shaped(itr, sz)
+    if prod(sz) == 0
+        return _default_container(itr, Union{}, sz)
+    end
+    st = start(itr)
+    v1, st = next(itr, st)
+    dest = _default_container(itr, typeof(v1), sz)
+    dest[1] = v1
+    return collect_to!(itr, 2, st, dest)
+end
+
+function grow_to!{T}(itr, st, dest::AbstractArray{T})
+    while !done(itr, st)
+        el, st = next(itr, st)
+        S = typeof(el)
+        if S === T || S <: T
+            push!(dest, el::T)
+        else
+            R = typejoin(T, S)
+            n = length(dest)
+            new = similar(dest, R, n+1)
+            copy!(new,1, dest,1, n)
+            @inbounds new[n+1] = el
+            return grow_to!(itr, st, new)
+        end
+    end
+    return dest
+end
+
+function _collect(itr, ::EltypeUnknown, ::SizeUnknown)
+    st = start(itr)
+    if done(itr,st)
+        return _default_container(itr, Union{}, 0)
+    end
+    v1, st = next(itr, st)
+    dest = _default_container(itr, typeof(v1), 1)
+    dest[1] = v1
+    return grow_to!(itr, st, dest)
+end
 
 ## Iteration ##
 start(A::Array) = 1

base/dict.jl

@@ -4,6 +4,9 @@
 
 abstract Associative{K,V}
 
+iteratorsize{T<:Associative}(::Type{T}) = HasLength()
+iteratoreltype{T<:Associative}(::Type{T}) = HasEltype()
+
 const secret_table_token = :__c782dbf1cf4d6a2e5e3865d7e95634f2e09b5902__
 
 haskey(d::Associative, k) = in(k,keys(d))

base/generator.jl

@@ -11,11 +11,48 @@ immutable Generator{I,F}
     iter::I
 end
 
+Generator{T,I}(::Type{T}, iter::I) = Generator{I,Type{T}}(T, iter)
+
 start(g::Generator) = start(g.iter)
 done(g::Generator, s) = done(g.iter, s)
 function next(g::Generator, s)
     v, s2 = next(g.iter, s)
     g.f(v), s2
 end
 
-collect(g::Generator) = map(g.f, g.iter)
+## iterator traits
+
+abstract IteratorSize
+immutable SizeUnknown <: IteratorSize end
+immutable HasLength <: IteratorSize end
+immutable HasShape <: IteratorSize end
+
+iteratorsize(x) = iteratorsize(typeof(x))
+iteratorsize(::Type) = SizeUnknown()
+
+and_iteratorsize{T}(isz::T, ::T) = isz
+and_iteratorsize(::HasLength, ::HasShape) = HasLength()
+and_iteratorsize(::HasShape, ::HasLength) = HasLength()
+and_iteratorsize(a, b) = SizeUnknown()
+
+abstract IteratorEltype
+immutable EltypeUnknown <: IteratorEltype end
+immutable HasEltype <: IteratorEltype end
+
+iteratoreltype(x) = iteratoreltype(typeof(x))
+iteratoreltype(::Type) = EltypeUnknown()
+
+and_iteratoreltype{T}(iel::T, ::T) = iel
+and_iteratoreltype(a, b) = EltypeUnknown()
+
+iteratorsize{T<:AbstractArray}(::Type{T}) = HasShape()
+iteratorsize{T<:AbstractString}(::Type{T}) = HasLength()
+iteratorsize{T<:Tuple}(::Type{T}) = HasLength()
+iteratorsize{I,F}(::Type{Generator{I,F}}) = iteratorsize(I)
+length(g::Generator) = length(g.iter)
+size(g::Generator) = size(g.iter)
+
+iteratoreltype{T<:AbstractArray}(::Type{T}) = HasEltype()
+iteratoreltype{T<:AbstractString}(::Type{T}) = HasEltype()
+iteratoreltype{I,T}(::Type{Generator{I,Type{T}}}) = HasEltype()
+eltype{I,T}(::Type{Generator{I,Type{T}}}) = T