SetupIterator.jl

Sometimes, you have a data structure where iterate is complicated to express for that structure, but Base.Iterators has already gone through the trouble of making an appropriate iterator for you.

For example, consider

struct VectorOfVectors{T}
   v::Vector{Vector{T}}
end
Base.IteratorEltype(::VectorOfVectors{T}) where {T} = T
Base.IteratorSize(::VectorOfVectors) = Base.HasLength()
Base.length(vov::VectorOfVectors) = isempty(vov.v) ? 0 : sum(length, vov.v)

and supose we want to treat this thing semantically like a flat vector, i.e. concatenate their contents. Writing a method for Base.iterate on this type is actually somewhat difficult, since you'll need to juggle the state of all the constituent vectors and check which sub-vector the current iteration state corresponds to.

Wouldn't it be great if we could just do something akin to

# This is a wish, not something that actually works
Base.iterator(vov::VectorOfVectors) = Iterators.flatten(vov.v) 

and be done with it?

Enter SetupIterator.jl:

using SetupIterator
SetupIterator.iterator(vov::VectorOfVectors) = Iterators.flatten(vov.v)
@setup_iterate VectorOfVectors

julia> let vov = VectorOfVectors([[1, 2, 3], [4, 5]])
           for x ∈ vov
               @show x
           end
       end
x = 1
x = 2
x = 3
x = 4
x = 5

Normally, when you write

for x ∈ xs
    #do stuff
end

this gets lowered to the code

next = iterate(xs)
while next !== nothing
  x, state = next
  #do stuff
  next = iterate(xs, state)
end

but when we use SetupIterator.jl, it becomes more equivalent to

itr = SetupIterator.iterator(xs)
next = iterate(itr)
while next !== nothing
  x, state = next
  #do stuff
  next = iterate(itr, state)
end

This way, there are no remaining references to xs inside the loop, and if we wish to implement something like a destructive iterator to reclaim memory from xs while the loop runs, then that is now possible.

Inspiration for this comes from Iterate on it by Mike Innes, and from this issue

Credit to Jameson Nash for pointing out that this can be done in this issue.