Type instability when zipping more than 7 arrays

[bramtayl]

x = zip([1], ["a"], [1.0], [1], ["a"], [1.0], [1], ["a"])
z = Iterators.filter(x -> x[1] == 1, x)
@code_warntype first(z)

I spent forever tracking down what was going on here. It seems like because states for zip are nested tuples, at a certain point it just becomes too meta for type inference?. Seems like an easy solution is a zip which is based on a tuple of arrays, not recursive zip objects? Hard solution is to get it to infer all the way down...

[bramtayl]

Improvement over 0.6 where the limit seems to be 5

[bramtayl]

Naively, this seems to work:

import Base: start, next, done, length, size, indices, eltype, iteratorsize, iteratoreltype, @propagate_inbounds

struct ManyToOne{O, M}
    one::O
    many::M
end

start(m::ManyToOne) = start(m.one), map(start, m.many)...
done(m::ManyToOne, states) = done(m.one, first(states))
length(m::ManyToOne) = length(m.one)
indices(m::ManyToOne) = indices(m.one)
size(m::ManyToOne) = size(m.one)
iteratorsize(m::ManyToOne) = iteratorsize(m.one)
iteratoreltype(m::ManyToOne) = iteratoreltype(m.one)
eltype(m::ManyToOne) = Tuple{eltype(m.one), map(eltype, m.many)...}

@propagate_inbounds function next(m::ManyToOne, state)
    mixed = map(next, (m.one, m.many...), state)
    map(first, mixed), map(last, mixed)
end

x = ManyToOne([1], (["a"], [1.0], [1], ["a"], [1.0], [1], ["a"]))
z = Iterators.filter(x -> x[1] == 1, x)
@code_warntype first(z)

[haampie]

The above is not entirely correct, as it should iterate until any iterator finishes.

Another way to escape from the long recursion is to generate the iterate method:

import Base: iterate
using Base: @nexprs, @ntuple

struct LinearZip{N,T}
    t::T
end

zip(a...) = LinearZip{length(a),typeof(a)}(a)

@generated function iterate(z::LinearZip{N}) where {N}
    quote
        @nexprs $N j -> begin
            y_j = iterate(z.t[j])
            y_j === nothing && return nothing
        end
        values = @ntuple $N j -> y_j[1]
        states = @ntuple $N j -> y_j[2]
        values, states
    end
end

@generated function iterate(z::LinearZip{N}, states) where {N}
    quote
        @nexprs $N j -> begin
            y_j = iterate(z.t[j], states[j])
            y_j === nothing && return nothing
        end
        values = @ntuple $N j -> y_j[1]
        states = @ntuple $N j -> y_j[2]
        values, states
    end
end

which infers a hundred or more zipped iterators and is at runtime probably as performant as it gets. What remains to be resolved is the issue of stateful iterators.

[bramtayl]

That would be great. I'm not sure I'm following the issue with stateful iterators. P.S. I'm pretty sure you could do the above with lispy tuple recursion instead of generated functions.

[haampie]

Indeed, you're correct :) apparently the following is easier to infer (I've verified it by zipping 300 iterators -- but compilation/execution takes a minute :p)

import Base: iterate, tail

struct LinearZip{T}
    t::T
end

zip(a, b...) = LinearZip((a, b...))

"""
    iterate_all((it1, it2, ...)) -> (val1, val2, ...), (state1, state2, ...) or nothing
"""
iterate_all(iterators::Tuple{}) = (), ()
iterate_all(iterators::Tuple{}, states::Tuple{}) = (), ()

function iterate_all(iterators::Tuple)
    curr = iterate(iterators[1])
    curr === nothing && return nothing
    others = iterate_all(tail(iterators))
    others === nothing && return nothing
    (curr[1], others[1]...), (curr[2], others[2]...)
end

function iterate_all(iterators::Tuple, states::Tuple)
    curr = iterate(iterators[1], states[1])
    curr === nothing && return nothing
    others = iterate_all(tail(iterators), tail(states))
    others === nothing && return nothing
    (curr[1], others[1]...), (curr[2], others[2]...)
end

iterate(z::LinearZip, states...) = iterate_all(z.t, states...)

[oheil]

The following code example seems to be related to this issue and to the general performance issue of the recursive implementation of zip using the Zip struct in julia/base/iterators.jl
zip(a, b, c...) = Zip(a, zip(b, c...))

When you start a new REPL:

julia> versioninfo()
Julia Version 1.1.0-DEV.69
Commit 57b3497f1b (2018-08-20 02:16 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i7-4770 CPU @ 3.40GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, haswell)

julia> a=collect(1:30);
julia> z=zip(a...);
julia> @time c=collect(z);
1339.290027 seconds (2.00 M allocations: 94.543 MiB, 0.01% gc time)