diff --git a/base/iterators.jl b/base/iterators.jl index b721933a37ea7..46a7167f89c52 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -315,32 +315,52 @@ IteratorEltype(::Type{Zip{I,Z}}) where {I,Z} = and_iteratoreltype(IteratorEltype reverse(it::ZipHead) = ZipHead(reverse(it.a)) reverse(it::Zip) = Zip(reverse(it.a), reverse(it.z)) -# Iteration -@inline function iterate(it::ZipHead, state...) - isdone(it.a) === true && return nothing - a = iterate(it.a, state...) - a === nothing && return nothing - ((a[1],), a[2]) +# the states are either Tuple{} or Tuple{Any} +@inline function zip_iterate(a, z, sta, stz) + da, dz = isdone(a), isdone(z) + + # Note that stateless iterators vectorize provided there's no early return. + # In the case of stateless iterators: all branches are removed compile-time + # In the case of stateful iterators: the first condition is almost always true + if da === dz === missing || da === dz === false + ya, yz = iterate(a, sta...), iterate(z, stz...) + ya === nothing && return nothing + yz === nothing && return nothing + elseif da === true || dz === true + return nothing + elseif da === missing + # Iterate a only if z is not done and vice versa + ya = iterate(a, sta...) + ya === nothing && return nothing + yz = iterate(z, stz...) + yz === nothing && return nothing + else + yz = iterate(z, stz...) + yz === nothing && return nothing + ya = iterate(a, sta...) + ya === nothing && return nothing + end + + return ya, yz end @inline function iterate(it::Zip) - isdone(it.a) === true && return nothing - isdone(it.z) === true && return nothing - a = iterate(it.a) - z = iterate(it.z) - a === nothing && return nothing - z === nothing && return nothing - ((a[1], z[1]...), (a[2], z[2])) + y = zip_iterate(it.a, it.z, (), ()) + y === nothing && return nothing + (val_a, state_a), (val_z, state_z) = y + return (val_a, val_z...), (state_a, state_z) end - -@inline function iterate(it::Zip, state::Tuple{Any, Any}) +@inline function iterate(it::Zip, s::Tuple{Any, Any}) + y = zip_iterate(it.a, it.z, (s[1],), (s[2],)) + y === nothing && return nothing + (val_a, state_a), (val_z, state_z) = y + return (val_a, val_z...), (state_a, state_z) +end +@inline function iterate(it::ZipHead, state...) isdone(it.a) === true && return nothing - isdone(it.z) === true && return nothing - a = iterate(it.a, state[1]) - z = iterate(it.z, state[2]) + a = iterate(it.a, state...) a === nothing && return nothing - z === nothing && return nothing - ((a[1], z[1]...), (a[2], z[2])) + ((a[1],), a[2]) end # filter