# Defining a custom iterator that wraps multiple iterators

**URL:** <https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790>\
**Category:** General Usage\
**Created:** [September 10, 2018, 11:10pm UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790 "2018-09-10T23:10:36Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![schmrlng](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/schmrlng/32/2325_2.png) [@schmrlng](https://discourse.julialang.org/u/schmrlng)\
**Post date:** [September 10, 2018, 11:10pm UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790/1 "2018-09-10T23:10:36Z")

</div>

I’m having a tough time writing performant iterators that wrap multiple other iterators; as an illustrative example this is how I would currently write a simple merge sort iterator:

```julia
struct MergeSortIterator{I1,I2}
    a::I1
    b::I2
end

Base.length(itr::MergeSortIterator) = length(itr.a) + length(itr.b)
Base.eltype(::Type{MergeSortIterator{I1,I2}}) where {I1,I2} = promote_type(eltype(I1), eltype(I2))

function Base.iterate(itr::MergeSortIterator, state=(iterate(itr.a), iterate(itr.b)))
    a_iter_result, b_iter_result = state
    a_iter_result === nothing && b_iter_result === nothing && return nothing
    a_iter_result === nothing && return (b_iter_result[1], (nothing, iterate(itr.b, b_iter_result[2])))
    b_iter_result === nothing && return (a_iter_result[1], (iterate(itr.a, a_iter_result[2]), nothing))
    item_a, state_a = a_iter_result
    item_b, state_b = b_iter_result
    item_a < item_b ? (item_a, (iterate(itr.a, state_a), b_iter_result)) : (item_b, (a_iter_result, iterate(itr.b, state_b)))
end

```

This works, e.g.,

```julia
v1 = sort!(rand(10))
v2 = sort!(rand(5))
collect(MergeSortIterator(v1, v2))
# [0.0135208, 0.125779, 0.221346, 0.386235, 0.530192, 0.614787, 0.629101, 0.702266, 0.708002, 0.85657, 0.86487, 0.886334, 0.895294, 0.912969, 0.947071]

```

but it’s pretty slow because the iterator state type can be a combinatorial explosion of component iterator `iterate` return types (`Nothing` or `Tuple{Float64,Int64}}` in this case):

```julia
julia> @code_warntype iterate(MergeSortIterator(v1, v2))
Body::Union{Nothing, Tuple{Float64,Any}}
2 1 ─ %1 = (Base.getfield)(itr, :a)::Array{Float64,1} │╻ getproperty
  │ %2 = (Base.arraylen)(%1)::Int64 ││╻╷ iterate
  │ %3 = (Base.sle_int)(0, %2)::Bool │││╻╷ <
  │ %4 = (Base.bitcast)(UInt64, %2)::UInt64 ││││╻ unsigned
  │ %5 = (Base.ult_int)(0x0000000000000000, %4)::Bool ││││╻ <
  │ %6 = (Base.and_int)(%3, %5)::Bool ││││╻ &
  └── goto #3 if not %6 │││  
  2 ─ %8 = (Base.arrayref)(false, %1, 1)::Float64 │││╻ getindex
  │ %9 = (Core.tuple)(%8, 2)::Tuple{Float64,Int64} │││  
  └── goto #4 │││  
  3 ─ %11 = Base.nothing::Nothing │││  
  └── goto #4 │││  
  4 ┄ %13 = φ (#2 => %9, #3 => %11)::Union{Nothing, Tuple{Float64,Int64}} ││   
  └── goto #5 ││   
  5 ─ %15 = (Base.getfield)(itr, :b)::Array{Float64,1} │╻ getproperty
  │ %16 = (Base.arraylen)(%15)::Int64 ││╻╷ iterate
  │ %17 = (Base.sle_int)(0, %16)::Bool │││╻╷ <
  │ %18 = (Base.bitcast)(UInt64, %16)::UInt64 ││││╻ unsigned
  │ %19 = (Base.ult_int)(0x0000000000000000, %18)::Bool ││││╻ <
  │ %20 = (Base.and_int)(%17, %19)::Bool ││││╻ &
  └── goto #7 if not %20 │││  
  6 ─ %22 = (Base.arrayref)(false, %15, 1)::Float64 │││╻ getindex
  │ %23 = (Core.tuple)(%22, 2)::Tuple{Float64,Int64} │││  
  └── goto #8 │││  
  7 ─ %25 = Base.nothing::Nothing │││  
  └── goto #8 │││  
  8 ┄ %27 = φ (#6 => %23, #7 => %25)::Union{Nothing, Tuple{Float64,Int64}} ││   
  └── goto #9 ││   
  9 ─ %29 = (Core.tuple)(%13, %27)::Tuple{Union{Nothing, Tuple{Float64,Int64}},Union{Nothing, Tuple{Float64,Int64}}}
  │ %30 = (#self#)(itr, %29)::Union{Nothing, Tuple{Float64,Any}} │    
  └── return %30 │    

```

Does anyone know of a better way to write iterators that need to maintain the states of sub-iterators?

* * *

I’ve taken a look at the implementation of `Base.zip` (in particular [`Zip2`](https://github.com/JuliaLang/julia/blob/master/base/iterators.jl#L284)). It doesn’t suffer from this type inference problem because it truncates on the first `nothing` it sees from any of the wrapped iterators (and thus the return type of `iterate` is just a union of `Nothing` and `Tuple{[Tuple of value types], [Tuple of state types]}`:

```julia
julia> @code_warntype iterate(zip(v1, v2))
Body::Union{Nothing, Tuple{Tuple{Float64,Float64},Tuple{Int64,Int64}}}
320 1 ── %1 = (Base.getfield)(z, :a)::Array{Float64,1} │╻ getproperty
    │ %2 = (Base.getfield)(z, :b)::Array{Float64,1} ││    
    └─── goto #8 if not true │╻ zip_iterate
...

```

In order to achieve this sort of semi-type-stability, I can keep around an extra sentinel as in this pretty hacky implementation (this works, and is performant), but I’m hoping there’s a better way (in particular, `_dummy_iterate_result` feels problematic):

```julia
_dummy_iterate_result(itr::Vector{T}) where {T} = (zero(T), 0)

function _start(itr::MergeSortIterator)
    a_iter_result = iterate(itr.a)
    b_iter_result = iterate(itr.b)
    if (a_done = a_iter_result === nothing)
        a_iter_result = _dummy_iterate_result(itr.a)
    end
    if (b_done = b_iter_result === nothing)
        b_iter_result = _dummy_iterate_result(itr.b)
    end
    (a_done, b_done, a_iter_result, b_iter_result)
end

function Base.iterate(itr::MergeSortIterator, state=_start(itr))
    a_done, b_done, a_iter_result, b_iter_result = state
    a_done && b_done && return nothing
    if b_done || (!a_done && a_iter_result[1] < b_iter_result[1])
        item_a, state_a = a_iter_result
        a_iter_result = iterate(itr.a, state_a)
        a_done = a_iter_result === nothing
        return (item_a, (a_done, b_done, a_done ? _dummy_iterate_result(itr.a) : a_iter_result, b_iter_result))
    else
        item_b, state_b = b_iter_result
        b_iter_result = iterate(itr.b, state_b)
        b_done = b_iter_result === nothing
        return (item_b, (a_done, b_done, a_iter_result, b_done ? _dummy_iterate_result(itr.b) : b_iter_result))
    end
end

```

---

<div class="post-metadata">

**Author:** ![Stephen\_Vavasis](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stephen_vavasis/32/3389_2.png) [@Stephen\_Vavasis](https://discourse.julialang.org/u/Stephen_Vavasis)\
**Post date:** [September 14, 2018, 3:40am UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790/2 "2018-09-14T03:40:34Z")

</div>

This would have been relatively easy under the old (0.6) iteration protocol. I think the following more complicated solution may work for the new iteration protocol. First, define a `mutable struct` to hold the vector of iteration states; each field is parameterized. If a component container is initially empty, then the corresponding field is populated with a `nothing` instead of an iteration state. This mutable struct also contains a Bool for each component iterator to indicate whether the particular container is exhausted. This is an object whose type cannot be determined at compile time, but once it is initialized, its type is stable. Therefore, after it is initialized, you can use a “function barrier” to send it to the actual worker routine.

---

<div class="post-metadata">

**Author:** ![Stephen\_Vavasis](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stephen_vavasis/32/3389_2.png) [@Stephen\_Vavasis](https://discourse.julialang.org/u/Stephen_Vavasis)\
**Post date:** [September 14, 2018, 12:46pm UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790/3 "2018-09-14T12:46:48Z")

</div>

Upon further thought, I decided that the solution proposed in my previous posting would have significant overhead in recompiling the worker function each time an iterator is created, so that it would be non-performant unless the containers are very large to compensate for the overhead.

However, I think the proposed solution can be fixed by defining some helper routines like this (not tested):

```julia
   iterator_state_type(::Any) = Any
   iterator_state_type(::Vector{T}) where {T<:Number} = Tuple{Int,T}
   # etc for all the common types over which you might want to call MergeSort
   default_iterator_state(::Any) = nothing
   default_iterator_state(::Vector{T}) where {T<:Number} = (0,zero(T))
   # etc for all the common types

```

Then in the setup routine described in my previous message, use these helper functions to create the iteration state for mergesort. In the common cases, the compiler will be able to use the helper function to infer the type of the struct, so that the worker arguments have a defined type signature at compile time. By the way, in my previous message, I said that the iteration state should be a mutable struct, but on further reflection, I think an immutable struct would be better.

---

<div class="post-metadata">

**Author:** ![schmrlng](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/schmrlng/32/2325_2.png) [@schmrlng](https://discourse.julialang.org/u/schmrlng)\
**Post date:** [September 14, 2018, 4:23pm UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790/4 "2018-09-14T16:23:18Z")

</div>

Thanks for your suggestions – unless I’m mistaken it seems the solution you’re proposing is similar in spirit to the “hacky implementation” at the bottom of the OP (there, a tuple is used as an anonymous version of the immutable struct you’re proposing). Over the past few days I’ve actually come to have a much more positive view of this style of approach. I can confirm that it’s zero overhead (i.e., `iterate(::MergeSortIterator)` doesn’t allocate as long as the iterators it wraps don’t) and in my code variants of it actually outperform the indexing-based `for` loops (applicable only to `AbstractVector`s, not arbitrary iterators) they are replacing.

My only lingering concern is that `_dummy_iterate_result` (or in your case, `default_iterator_state` with the order of element/state flipped) needs to instantiate a default value of the eltype `T` in general (I’m not thrilled with, but ultimately okay with defining `_dummy_iterate_result` for each iterator type I might use). We both use `zero(T)` but I can imagine cases where `zero(T)` doesn’t make any sense. It would be nice to have a language feature to define placeholder instances of types where the contract for usage is that the programmer must verify that those instances won’t actually be used but, notably, the burden of proof is taken off of the compiler. Maybe a better version of something like this (which actually seems to work okay for `isbitstype`s):

```julia
julia> struct A end

julia> undef(A)
ERROR: MethodError: objects of type UndefInitializer are not callable
Stacktrace:
 [1] top-level scope at none:0

julia> (::UndefInitializer)(::Type{T}) where {T} = Ref{T}()[]

julia> undef(Int)
139950704410448

julia> undef(A)
A()

julia> using StaticArrays

julia> undef(SVector{3,Float64})
3-element SArray{Tuple{3},Float64,1,3}:
 6.9529415980266e-310 
 1.0e-323             
 6.91450508901983e-310

julia> undef(typeof(+))
+ (generic function with 171 methods)

```

which would allow a definition like

```julia
_dummy_iterate_result(itr::Vector{T}) where {T} = (undef(T), 0)

```

that I would be more happy with.

---

<div class="post-metadata">

**Author:** ![Stephen\_Vavasis](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stephen_vavasis/32/3389_2.png) [@Stephen\_Vavasis](https://discourse.julialang.org/u/Stephen_Vavasis)\
**Post date:** [September 14, 2018, 7:22pm UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790/5 "2018-09-14T19:22:58Z")

</div>

You are right-- my proposed solution is quite similar to your OP; sorry I didn’t notice this. So your solution seems fine to me, except possibly (IIUC) there may be the following small improvements:

- The Ref{T}() trick (which I’ve never seen before!) doesn’t always work (see trace below). So there may need to be a more generic fallback if you want to write a routine that won’t fail.

- Even when it does work, it appears to me that it might allocate, i.e., it seems that `Ref{Int}()[]` must put something on the heap, So, assuming this is the case, you might want to change your `Base.iterate` so that it does not call `_dummy_iterate_result()` but instead uses the last known good iteration state as the placeholder, and you confine the invocation of `_dummy_iterate_result` to `_start`. Not sure about this.

```julia
julia> x = Ref{Vector{Int}}()[]
ERROR: UndefRefError: access to undefined reference
Stacktrace:
 [1] getproperty at .\sysimg.jl:18 [inlined]
 [2] getindex(::Base.RefValue{Array{Int64,1}}) at .\refvalue.jl:32
 [3] top-level scope at none:0

```

---

<div class="post-metadata">

**Author:** ![schmrlng](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/schmrlng/32/2325_2.png) [@schmrlng](https://discourse.julialang.org/u/schmrlng)\
**Post date:** [September 15, 2018, 2:01am UTC](https://discourse.julialang.org/t/defining-a-custom-iterator-that-wraps-multiple-iterators/14790/6 "2018-09-15T02:01:04Z")

</div>

Yeah, I think `undef(::Type{T})` would have to be a language feature/compiler hint (along the lines of `@simd` or `@inbounds`, as opposed to being a function that returns an actually accessible value) to be most useful/general.

I agree that there’s no reason to call this in any case outside of the initial `Nothing`/non-`Nothing` ambiguity at the first `iterate` call.