Skip to content

Commit

Permalink
make map and collect more general and uniform by adding iterator …
Browse files Browse the repository at this point in the history 
…traits
  • Loading branch information
@JeffBezanson
JeffBezanson committed Mar 7, 2016
1 parent ffa7e7b commit 8bd8ea3
Show file tree
Hide file tree
Showing 9 changed files with 163 additions and 194 deletions.
150 changes: 3 additions & 147 deletions base/abstractarray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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)
Expand All @@ -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{}) = ()
Expand All @@ -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)
Expand Down
99 changes: 81 additions & 18 deletions base/array.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions base/dict.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
39 changes: 38 additions & 1 deletion base/generator.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading

0 comments on commit 8bd8ea3

Please sign in to comment.