zip

_zip.jl

module ZipStuff

import Base: iterate
using Test, Base.Iterators

abstract type AbstractZip end

struct ZipBase{H} <: AbstractZip
    head::H
end

struct Zip{H,T<:AbstractZip} <: AbstractZip
    head::H
    tail::T
end

@inline function iterate(z::ZipBase, s...)
    head = iterate(z.head, s...)
    head === nothing && return nothing
    value, state = head
    return (value,), state
end

@inline function iterate(z::Zip)
    head = iterate(z.head)
    tail = iterate(z.tail)
    head === nothing && return nothing
    tail === nothing && return nothing
    headvalue, headstate = head
    tailvalue, tailstate = tail
    return (headvalue, tailvalue...), (headstate, tailstate)
end

@inline function iterate(z::Zip, state)
    head = iterate(z.head, state[1])
    tail = iterate(z.tail, state[2])
    head === nothing && return nothing
    tail === nothing && return nothing
    headvalue, headstate = head
    tailvalue, tailstate = tail
    return (headvalue, tailvalue...), (headstate, tailstate)
end

zip(head) = ZipBase(head)
zip(head, tail...) = Zip(head, zip(tail...))

"""
LLVM should be able to recognize this as having a closed form; should be O(1) complexity
"""
function optimize_away_loop_with_closed_form(n)
    s = 0

    for (a, (b, c), d, e) in zip(n:2n, zip(2:n+1, 3:n+3), 4:n+5, 5:n+8)
        s += (a + b) * c + d * e
    end
    s
end


"""
Should vectorize and be on par with enumerate(v) and a handwritten loop
"""
function vectorize_loop(v)
    s = 0
    for (i, w) = zip(LinearIndices(v), v)
        s += i * w
    end
    s
end

"""
Should be able to infer value type with a bunch of iterators when nested in another iterator
"""
function inference_with_lots_of_iterators(N = 14)
    candidates = [[1], (1, 2), 1:10, [1.0], rand(UInt, 2,2), "ab"]
    iterators = ntuple(x -> rand(candidates), N)
    it = Iterators.filter(x -> x !== 10, zip(iterators...))
    @inferred first(it)
end

end

results.jl

using BenchmarkTools

# The following shows that things get simplified to a closed-form expression

@btime ZipStuff.optimize_away_loop_with_closed_form(10)
#  1.888 ns (0 allocations: 0 bytes)
#2630

@btime ZipStuff.optimize_away_loop_with_closed_form(10000)
#  1.887 ns (0 allocations: 0 bytes)
#1500950180000

#############################################

# Vectorization:

@code_llvm ZipStuff.vectorize_loop(rand(Int, 100))

# ...
#  %bin.rdx = add <4 x i64> %29, %28
#  %rdx.shuf = shufflevector <4 x i64> %bin.rdx, <4 x i64> undef, <4 x i32> <i32 2, i32 3, i32 undef, i32 undef>
#  %bin.rdx94 = add <4 x i64> %bin.rdx, %rdx.shuf
#  %rdx.shuf95 = shufflevector <4 x i64> %bin.rdx94, <4 x i64> undef, <4 x i32> <i32 1, i32 undef, i32 undef, i32 undef>
#  %bin.rdx96 = add <4 x i64> %bin.rdx94, %rdx.shuf95
#  %31 = extractelement <4 x i64> %bin.rdx96, i32 0
#  %cmp.n = icmp eq i64 %n.vec, %20
# ...

# However, when we define iterate as follows:

@inline function iterate(z::Zip, state)
    head = iterate(z.head, state[1])
    head === nothing && return nothing
    tail = iterate(z.tail, state[2])
    tail === nothing && return nothing
    headvalue, headstate = head
    tailvalue, tailstate = tail
    return (headvalue, tailvalue...), (headstate, tailstate)
end

# then suddenly vectorization is gone!

#############################################

# Inference:

ZipStuff.inference_with_lots_of_iterators(14)
# (1.0, 'a', 1, 1, 0x5c24f0ca2d5419c3, 0x5c24f0ca2d5419c3, 'a', 1.0, 1, 1, 'a', 'a', 1, 1)

ZipStuff.inference_with_lots_of_iterators(15)
# ERROR: return type Tuple{Int64,UInt64,UInt64,Int64,Int64,Char,UInt64,UInt64,Int64,Char,UInt64,Float64,UInt64,Int64,Char} does not match inferred return type Tuple{Int64,UInt64,UInt64,Int64,Int64,Char,UInt64,UInt64,Int64,Char,UInt64,Float64,UInt64,Vararg{Union{Char, Int64},N} where N}