Skip to content

Commit

Permalink
Merge 3eb00e0 into fc10c8b
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jutho
Jutho authored Jul 28, 2017
2 parents fc10c8b + 3eb00e0 commit 7d371c7
Showing 1 changed file with 77 additions and 99 deletions.
176 changes: 77 additions & 99 deletions base/iterators.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ module Iterators

import Base: start, done, next, isempty, length, size, eltype, iteratorsize, iteratoreltype, indices, ndims

using Base: tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, @propagate_inbounds
using Base: tail, tuple_type_head, tuple_type_tail, tuple_type_cons, SizeUnknown, HasLength, HasShape, IsInfinite, EltypeUnknown, HasEltype, OneTo, @propagate_inbounds

export enumerate, zip, rest, countfrom, take, drop, cycle, repeated, product, flatten, partition

Expand Down Expand Up @@ -575,63 +575,18 @@ iteratoreltype(::Type{<:Repeated}) = HasEltype()


# Product -- cartesian product of iterators

abstract type AbstractProdIterator end

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))
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)
_prod_size(a, ::HasLength) = (length(a),)
_prod_size(a, A) =
throw(ArgumentError("Cannot compute size for object of type $(typeof(a))"))
_prod_size(a, b, ::HasLength, ::HasLength) = (length(a), length(b))
_prod_size(a, b, ::HasLength, ::HasShape) = (length(a), size(b)...)
_prod_size(a, b, ::HasShape, ::HasLength) = (size(a)..., length(b))
_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
struct Prod1{I} <: AbstractProdIterator
a::I
struct ProductIterator{T<:Tuple}
iterators::T
end
product(a) = Prod1(a)

eltype(::Type{Prod1{I}}) where {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)
n, st = next(p.a, st)
(n,), st
end
@inline done(p::Prod1, st) = done(p.a, st)

iteratoreltype(::Type{Prod1{I}}) where {I} = iteratoreltype(I)
iteratorsize(::Type{Prod1{I}}) where {I} = iteratorsize(I)

# two iterators
struct Prod2{I1, I2} <: AbstractProdIterator
a::I1
b::I2
struct ProductIteratorState{T,S}
done::Bool
states::T
values::S
ProductIteratorState{T,S}(done, states::T, values::S) where {T,S} = new(done, states, values)
ProductIteratorState{T,S}(done, states::T) where {T,S} = new(done, states)
end
ProductIteratorState(done, states::T, values::S) where {T,S} = ProductIteratorState{T,S}(done, states, values)

"""
product(iters...)
Expand All @@ -648,62 +603,85 @@ julia> collect(Iterators.product(1:2,3:5))
(2, 3) (2, 4) (2, 5)
```
"""
product(a, b) = Prod2(a, b)
product(iters...) = ProductIterator(iters)

eltype(::Type{Prod2{I1,I2}}) where {I1,I2} = Tuple{eltype(I1), eltype(I2)}
iteratorsize(::Type{ProductIterator{Tuple{}}}) = HasShape()
iteratorsize(::Type{ProductIterator{Tuple{I}}}) where {I} = iteratorsize(I)
iteratorsize(::Type{ProductIterator{T}}) where {T<:Tuple} =
prod_iteratorsize( iteratorsize(tuple_type_head(T)), iteratorsize(ProductIterator{tuple_type_tail(T)}) )

iteratoreltype(::Type{Prod2{I1,I2}}) where {I1,I2} = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2))
iteratorsize(::Type{Prod2{I1,I2}}) where {I1,I2} = prod_iteratorsize(iteratorsize(I1),iteratorsize(I2))
prod_iteratorsize(::Union{HasLength,HasShape}, ::Union{HasLength,HasShape}) = HasShape()
# products can have an infinite iterator
prod_iteratorsize(::IsInfinite, ::IsInfinite) = IsInfinite()
prod_iteratorsize(a, ::IsInfinite) = IsInfinite()
prod_iteratorsize(::IsInfinite, b) = IsInfinite()
prod_iteratorsize(a, b) = SizeUnknown()

function start(p::AbstractProdIterator)
s1, s2 = start(p.a), start(p.b)
s1, s2, Nullable{eltype(p.b)}(), (done(p.a,s1) || done(p.b,s2))
end
size(P::ProductIterator) = _prod_size(P.iterators)
_prod_size(::Tuple{}) = ()
_prod_size(t::Tuple) = tuple(_prod_size1(t[1], iteratorsize(t[1]))..., _prod_size(tail(t))...)
_prod_size1(a, ::HasShape) = size(a)
_prod_size1(a, ::HasLength) = (length(a),)
_prod_size1(a, A) =
throw(ArgumentError("Cannot compute size for object of type $(typeof(a))"))

@inline function prod_next(p, st)
s1, s2 = st[1], st[2]
v1, s1 = next(p.a, s1)
indices(P::ProductIterator) = _prod_indices(P.iterators)
_prod_indices(::Tuple{}) = ()
_prod_indices(t::Tuple) = tuple(_prod_indices1(t[1], iteratorsize(t[1]))..., _prod_indices(tail(t))...)
_prod_indices1(a, ::HasShape) = indices(a)
_prod_indices1(a, ::HasLength) = (OneTo(length(a)),)
_prod_indices1(a, A) =
throw(ArgumentError("Cannot compute indices for object of type $(typeof(a))"))

nv2 = st[3]
if isnull(nv2)
v2, s2 = next(p.b, s2)
else
v2 = nv2.value
end
ndims(p::ProductIterator) = length(indices(p))
length(P::ProductIterator) = prod(size(P))
_length(p::ProductIterator) = prod(map(unsafe_length, indices(p)))

if done(p.a, s1)
return (v1,v2), (start(p.a), s2, Nullable{eltype(nv2)}(), done(p.b,s2))
end
return (v1,v2), (s1, s2, Nullable(v2), false)
iteratoreltype(::Type{ProductIterator{Tuple{}}}) = HasEltype()
iteratoreltype(::Type{ProductIterator{Tuple{I}}}) where {I} = iteratoreltype(I)
function iteratoreltype(::Type{ProductIterator{T}}) where {T<:Tuple}
I = tuple_type_head(T)
P = ProductIterator{tuple_type_tail(T)}
iteratoreltype(I) == EltypeUnknown() ? EltypeUnknown() : iteratoreltype(P)
end

@inline next(p::Prod2, st) = prod_next(p, st)
@inline done(p::AbstractProdIterator, st) = st[4]
Base.eltype(P::ProductIterator) = _prod_eltype(P.iterators)
_prod_eltype(::Tuple{}) = Tuple{}
_prod_eltype(t::Tuple) = Base.tuple_type_cons(eltype(t[1]),_prod_eltype(tail(t)))

# n iterators
struct Prod{I1, I2<:AbstractProdIterator} <: AbstractProdIterator
a::I1
b::I2
function Base.start(P::ProductIterator)
iterators = P.iterators
states = map(start, iterators)
if any(map(done, iterators, states))
return ProductIteratorState{typeof(states),eltype(P)}(true, states)
else
return ProductIteratorState{typeof(states),eltype(P)}(false, _prod_start(iterators, states)...)
end
end
product(a, b, c...) = Prod(a, product(b, c...))

eltype(::Type{Prod{I1,I2}}) where {I1,I2} = tuple_type_cons(eltype(I1), eltype(I2))

iteratoreltype(::Type{Prod{I1,I2}}) where {I1,I2} = and_iteratoreltype(iteratoreltype(I1),iteratoreltype(I2))
iteratorsize(::Type{Prod{I1,I2}}) where {I1,I2} = prod_iteratorsize(iteratorsize(I1),iteratorsize(I2))
Base.next(P::ProductIterator, state) = state.values, ProductIteratorState(_prod_next(P.iterators, state.states, state.values)...)
Base.done(P::ProductIterator, state) = state.done

@inline function next(p::Prod, st)
x = prod_next(p, st)
((x[1][1],x[1][2]...), x[2])
_prod_start(iterators::Tuple{}, states::Tuple{}) = (), ()
function _prod_start(iterators, states)
val, state = next(iterators[1], states[1])
tailstates, tailvals = _prod_start(tail(iterators), tail(states))
return (state, tailstates...), (val, tailvals...)
end
_prod_next(iterators::Tuple{}, states::Tuple{}, values::Tuple{}) = true, (), ()
function _prod_next(iterators, states, values)
if !done(iterators[1], states[1])
nextval, nextstate = next(iterators[1], states[1])
return false, tuple(nextstate, tail(states)...), tuple(nextval, tail(values)...)
else
d, nextstates, nextvals = _prod_next(tail(iterators), tail(states), tail(values))
if d # all iterators are done
return true, states, values
else
nextval, nextstate = next(iterators[1], start(iterators[1]))
return false, tuple(nextstate, nextstates...), tuple(nextval, nextvals...)
end
end
end

prod_iteratorsize(::Union{HasLength,HasShape}, ::Union{HasLength,HasShape}) = HasShape()
# products can have an infinite iterator
prod_iteratorsize(::IsInfinite, ::IsInfinite) = IsInfinite()
prod_iteratorsize(a, ::IsInfinite) = IsInfinite()
prod_iteratorsize(::IsInfinite, b) = IsInfinite()
prod_iteratorsize(a, b) = SizeUnknown()


# flatten an iterator of iterators

Expand Down

0 comments on commit 7d371c7

Please sign in to comment.