Make zip more complex & correct again

Author: haampie

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