Iterator tool similar to zip (mesh?)

General Usage
,
8

Here’s a quick attempt:

interleave(iters...) = Interleaved(iters)
struct Interleaved{T<:Tuple}
    iters::T
end
Base.length(x::Interleaved) = sum(length, x.iters)
function Base.IteratorSize(x::Interleaved)
    traits = map(Base.IteratorSize, x.iters)
    if all(t -> t isa Union{Base.HasLength, Base.HasShape}, traits)
        Base.HasLength()
    elseif any(t -> t isa Base.IsInfinite, traits)
        Base.IsInfinite()
    else
        Base.SizeUnknown()
    end
end

struct _Done end
struct _Init end
function Base.iterate(x::Interleaved, (id, states) = (1, map(_ -> _Init(), x.iters)))
    nextid = mod1(id+1, length(x.iters))

    all(s -> s isa _Done, states) && return nothing
    states[id] isa _Done && return iterate(x, (nextid, states))

    it = if states[id] isa _Init
        iterate(x.iters[id])
    else
        iterate(x.iters[id], states[id])
    end
    if it isa Nothing
        donestate = ntuple(i -> i==id ? _Done() : states[i], length(x.iters))
        return iterate(x, (nextid, donestate))
    else
        y, s = it
        newstate = ntuple(i -> i==id ? s : states[i], length(x.iters))
        return y, (nextid, newstate)
    end
end

interleave(1:5, 'a':'c', (), "AB") |> collect

first(interleave(1:5, 'a':'c', Iterators.cycle([10.0, 100])), 20)

Edit: a much shorter version, with the same struct, by cycling them around:

function Base.iterate(x::Interleaved, (this, rest...) = map(tuple, x.iters))
    it = iterate(this...)
    isnothing(it) && return iterate(x, rest)
    val, st = it
    return val, (rest..., (this[1], st))
end
Base.iterate(x::Interleaved, state::Tuple{}) = nothing

(No idea if these perform well, not carefully checked, etc.)

Edit’: added IteratorSize method to allow infinite iterators.

3 Likes