Skip to content

Commit

Permalink
indices extensions for iterators & collect
Browse files Browse the repository at this point in the history
  • Loading branch information
timholy committed Aug 4, 2016
1 parent 043eb26 commit a1e5daf
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 3 deletions.
4 changes: 2 additions & 2 deletions base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -211,7 +211,7 @@ The result has the same shape and number of dimensions as `collection`.
collect{T}(::Type{T}, itr) = _collect(T, itr, iteratorsize(itr))

_collect{T}(::Type{T}, itr, isz::HasLength) = copy!(Array{T,1}(Int(length(itr)::Integer)), itr)
_collect{T}(::Type{T}, itr, isz::HasShape) = copy!(Array{T}(convert(Dims,size(itr))), itr)
_collect{T}(::Type{T}, itr, isz::HasShape) = copy!(similar(Array{T}, indices(itr)), itr)
function _collect{T}(::Type{T}, itr, isz::SizeUnknown)
a = Array{T,1}(0)
for x in itr
Expand Down Expand Up @@ -259,7 +259,7 @@ end
_default_eltype{I,T}(::Type{Generator{I,Type{T}}}) = T

_array_for(T, itr, ::HasLength) = Array{T,1}(Int(length(itr)::Integer))
_array_for(T, itr, ::HasShape) = Array{T}(convert(Dims,size(itr)))
_array_for(T, itr, ::HasShape) = similar(Array{T}, indices(itr))

function collect(itr::Generator)
isz = iteratorsize(itr.iter)
Expand Down
19 changes: 18 additions & 1 deletion base/iterator.jl
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ end
zip(a) = Zip1(a)
length(z::Zip1) = length(z.a)
size(z::Zip1) = size(z.a)
indices(z::Zip1) = indices(z.a)
eltype{I}(::Type{Zip1{I}}) = Tuple{eltype(I)}
@inline start(z::Zip1) = start(z.a)
@inline function next(z::Zip1, st)
Expand All @@ -67,6 +68,7 @@ end
zip(a, b) = Zip2(a, b)
length(z::Zip2) = _min_length(z.a, z.b, iteratorsize(z.a), iteratorsize(z.b))
size(z::Zip2) = promote_shape(size(z.a), size(z.b))
indices(z::Zip2) = promote_shape(indices(z.a), indices(z.b))
eltype{I1,I2}(::Type{Zip2{I1,I2}}) = Tuple{eltype(I1), eltype(I2)}
@inline start(z::Zip2) = (start(z.a), start(z.b))
@inline function next(z::Zip2, st)
Expand All @@ -86,6 +88,7 @@ end
zip(a, b, c...) = Zip(a, zip(b, c...))
length(z::Zip) = _min_length(z.a, z.z, iteratorsize(z.a), iteratorsize(z.z))
size(z::Zip) = promote_shape(size(z.a), size(z.z))
indices(z::Zip) = promote_shape(indices(z.a), indices(z.z))
tuple_type_cons{S}(::Type{S}, ::Type{Union{}}) = Union{}
function tuple_type_cons{S,T<:Tuple}(::Type{S}, ::Type{T})
@_pure_meta
Expand Down Expand Up @@ -308,8 +311,10 @@ iteratoreltype{O}(::Type{Repeated{O}}) = HasEltype()
abstract AbstractProdIterator

length(p::AbstractProdIterator) = prod(size(p))
_length(p::AbstractProdIterator) = prod(map(unsafe_length, indices(p)))
size(p::AbstractProdIterator) = _prod_size(p.a, p.b, iteratorsize(p.a), iteratorsize(p.b))
ndims(p::AbstractProdIterator) = length(size(p))
indices(p::AbstractProdIterator) = _prod_indices(p.a, p.b, iteratorsize(p.a), iteratorsize(p.b))
ndims(p::AbstractProdIterator) = length(indices(p))

# generic methods to handle size of Prod* types
_prod_size(a, ::HasShape) = size(a)
Expand All @@ -323,6 +328,17 @@ _prod_size(a, b, ::HasShape, ::HasShape) = (size(a)..., size(b)...)
_prod_size(a, b, A, B) =
throw(ArgumentError("Cannot construct size for objects of types $(typeof(a)) and $(typeof(b))"))

_prod_indices(a, ::HasShape) = indices(a)
_prod_indices(a, ::HasLength) = (OneTo(length(a)), )
_prod_indices(a, A) =
throw(ArgumentError("Cannot compute indices for object of type $(typeof(a))"))
_prod_indices(a, b, ::HasLength, ::HasLength) = (OneTo(length(a)), OneTo(length(b)))
_prod_indices(a, b, ::HasLength, ::HasShape) = (OneTo(length(a)), indices(b)...)
_prod_indices(a, b, ::HasShape, ::HasLength) = (indices(a)..., OneTo(length(b)))
_prod_indices(a, b, ::HasShape, ::HasShape) = (indices(a)..., indices(b)...)
_prod_indices(a, b, A, B) =
throw(ArgumentError("Cannot construct indices for objects of types $(typeof(a)) and $(typeof(b))"))

# one iterator
immutable Prod1{I} <: AbstractProdIterator
a::I
Expand All @@ -331,6 +347,7 @@ product(a) = Prod1(a)

eltype{I}(::Type{Prod1{I}}) = Tuple{eltype(I)}
size(p::Prod1) = _prod_size(p.a, iteratorsize(p.a))
indices(p::Prod1) = _prod_indices(p.a, iteratorsize(p.a))

@inline start(p::Prod1) = start(p.a)
@inline function next(p::Prod1, st)
Expand Down

0 comments on commit a1e5daf

Please sign in to comment.