A type stable product iterator

productiterator.jl

import Base: start, next, done, size, length, eltype, iteratorsize, iteratoreltype, first, last
using Base: tail, IteratorEltype, HasEltype, EltypeUnknown, IteratorSize, HasShape, HasLength, IsInfinite, SizeUnknown
immutable ProductIterator{T<:Tuple}
    iterators::T
end

ProductIterator(iterators...) = ProductIterator(iterators)

iteratorsize(P::ProductIterator) = _itsize(P.iterators)
@inline _itsize(::Tuple{}) = HasShape()
@inline _itsize(t::Tuple) = isa(iteratorsize(t[1]), Union{IsInfinite,SizeUnknown}) ?
    iteratorsize(t[1]) : _itsize(tail(t))

size(P::ProductIterator) = _size(P.iterators)
_size(::Tuple{}) = ()
_size(t::Tuple) = tuple(size(t[1])...,_size(tail(t))...)

length(P::ProductIterator) = prod(size(P))

iteratoreltype(P::ProductIterator) = _iteltype(P.iterators)
@inline _iteltype(t::Tuple{}) = HasEltype()
@inline _iteltype(t::Tuple) = isa(iteratoreltype(t[1]), EltypeUnknown) ?
    EltypeUnknown() : _iteltype(tail(t))

eltype(P::ProductIterator) = _eltype(P.iterators)
@inline _eltype(::Tuple{}) = Tuple{}
@inline _eltype(t::Tuple) = Base.tuple_type_cons(eltype(t[1]),_eltype(tail(t)))

@inline function start(P::ProductIterator)
    iterators = P.iterators
    states = map(start, iterators)
    d = any(map(done, iterators, states))
    return (d, states)
end
@inline next(P::ProductIterator, state) = nextval(P, state), nextstate(P, state)
@inline done(P::ProductIterator, state) = state[1]

@inline nextval(iter, s) = next(iter, s)[1]
@inline nextstate(iter, s) = next(iter, s)[2]

@inline nextval(P::ProductIterator, state) = map(nextval, P.iterators, state[2])
@inline nextstate(P::ProductIterator{Tuple{}}, state) = (true, ())
@inline nextstate(P::ProductIterator, state) = carry(P.iterators, inc(P.iterators, state[2]))

# increment & carry
@inline inc(iterators::Tuple, states::Tuple) = (nextstate(iterators[1], states[1]), tail(states)...)
@inline carry(iterators::Tuple{Any}, states::Tuple{Any}) = done(iterators[1], states[1]), states
@inline function carry(iterators::Tuple, states::Tuple)
    !done(iterators[1],states[1]) && return (false, states)
    itertail = tail(iterators)
    statetail = tail(states)
    d, statetail = carry(itertail, inc(itertail, statetail))
    return d, tuple(start(iterators[1]), statetail...)
end

first(P::ProductIterator) = map(first, P.iterators)
last(P::ProductIterator) = map(last, P.iterators)