# Avoid creation of temporaries in non-trivial iterator

**URL:** <https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498>\
**Category:** Performance\
**Tags:** memory-allocation, type-stability, iterators\
**Created:** [July 12, 2021, 11:03am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498 "2021-07-12T11:03:12Z")\
**Posts on this page:** 17\
**Page:** 1

<div class="post-metadata">

**Author:** ![mwallerb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mwallerb/32/52642_2.png) [@mwallerb](https://discourse.julialang.org/u/mwallerb)\
**Post date:** [July 12, 2021, 11:03am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/1 "2021-07-12T11:03:12Z")

</div>

Consider the following Julia “compound” iterator: it merges two iterators, `a` and `b` , each of which are assumed to be sorted according to `order`, to a single ordered sequence:

```julia
struct MergeSorted{T,A,B,O}
    a::A
    b::B
    order::O
    MergeSorted(a::A, b::B, order::O=Base.Order.Forward) where {A,B,O} =
        new{promote_type(eltype(A),eltype(B)),A,B,O}(a, b, order)
end

Base.eltype(::Type{MergeSorted{T,A,B,O}}) where {T,A,B,O} = T

function Base.iterate(self::MergeSorted{T}, 
                      state=(iterate(self.a), iterate(self.b))) where T
    a_result, b_result = state
    if b_result === nothing
        a_result === nothing && return nothing
        a_curr, a_state = a_result
        return T(a_curr), (iterate(self.a, a_state), b_result)
    end
    b_curr, b_state = b_result
    if a_result !== nothing
        a_curr, a_state = a_result
        Base.Order.lt(self.order, a_curr, b_curr) &&
            return T(a_curr), (iterate(self.a, a_state), b_result)
    end
    return T(b_curr), (a_result, iterate(self.b, b_state))
end

x = MergeSorted([1,4,5,9,32,44], [0,7,9,24,134])
sum(x); print(@allocated sum(x))

```

This code works, but creates a temporary in each iteration step. Running `iterate(::MergeSorted)` through `@code_warntype` reveals that the return type is `Union{Nothing, Tuple{Int, Any}}`, which means that Julia gives up on specializing on the return type.

I first posted this question on [stackoverflow](https://stackoverflow.com/questions/68271829/how-to-avoid-memory-allocations-in-custom-julia-iterators), where user Bogumił Kamiński suggested a rather ugly workaround, which introduces a special state type to “force” julia to specialize. We were wondering whether there is a better possiblity, e.g., by forcing Julia to specialize on more complicated types?

Or maybe there is a smart way to rewrite this such that julia can figure this out on its own? I got encouraged by the standard library’s `Base.Iterators` package: it has a `flatten()` routine which does pretty much the same thing as `MergeSorted()` except that it does not intersperse the elements, yet it does not create temporaries. Does anyone with more iterator-foo transfer this to this case?

---

<div class="post-metadata">

**Author:** ![josuagrw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josuagrw/32/1015_2.png) [@josuagrw](https://discourse.julialang.org/u/josuagrw)\
**Post date:** [July 12, 2021, 11:43am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/2 "2021-07-12T11:43:15Z")

</div>

This seems like a case for multiple dispatch—implement four methods for the four different possibilities (A + B, A + nothing, nothing + B, nothing + nothing) and then let the type system figure out the rest

---

<div class="post-metadata">

**Author:** ![mwallerb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mwallerb/32/52642_2.png) [@mwallerb](https://discourse.julialang.org/u/mwallerb)\
**Post date:** [July 12, 2021, 11:57am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/3 "2021-07-12T11:57:47Z")

</div>

@josuagrw : Unfortunately, no. Here is the code for multiple dispatch, which does not help avoiding temporaries:

```julia
struct MergeSorted{T,A,B,O,S}
    a::A
    b::B
    order::O
    init::S

    function MergeSorted(a::A, b::B, order::O=Base.Order.Forward) where {A,B,O}
        T = promote_type(eltype(A), eltype(B))
        init = _make_state(iterate(a), iterate(b))
        new{T,A,B,O,typeof(init)}(a,b,order,init)
    end
end

Base.eltype(::Type{MergeSorted{T,A,B,O,S}}) where {T,A,B,O,S} = T

Base.iterate(self::MergeSorted) = iterate(self, self.init)

Base.iterate(self::MergeSorted{T}, ::Val{:done}) where T = nothing

function Base.iterate(self::MergeSorted{T}, state::Tuple{Val{:a_only},A}) where {T,A}
    _, (a_curr, a_state) = state
    return T(a_curr), _make_state(nothing, iterate(self.a, a_state))
end
    
function Base.iterate(self::MergeSorted{T}, state::Tuple{Val{:b_only},B}) where {T,B}
    _, (b_curr, b_state) = state
    return T(b_curr), _make_state(nothing, iterate(self.b, b_state))
end

function Base.iterate(self::MergeSorted{T}, state::Tuple{Val{:both},A,B}) where {T,A,B}
    _, (a_curr, a_state), (b_curr, b_state) = state
    if Base.Order.lt(self.order, a_curr, b_curr)
        return T(a_curr), _make_state(iterate(self.a, a_state), (b_curr, b_state))
    else
        return T(b_curr), _make_state((a_curr, a_state), iterate(self.b, b_state))
    end        
end

_make_state(a_result::Nothing, b_result::Nothing) = Val(:done)
_make_state(a_result::Nothing, b_result) = Val(:b_only), b_result
_make_state(a_result, b_result::Nothing) = Val(:a_only), a_result
_make_state(a_result, b_result) = Val(:both), a_result, b_result

```

---

<div class="post-metadata">

**Author:** ![pablosanjose](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pablosanjose/32/7006_2.png) [@pablosanjose](https://discourse.julialang.org/u/pablosanjose)\
**Post date:** [July 13, 2021, 12:37am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/4 "2021-07-13T00:37:12Z")

</div>

What an interesting little problem!

Here’s a multiple dispatch + look-ahead approach that is inferable and free of allocations in v1.6+

```julia
struct MergeSorted{T,A,B,O}
    a::A
    b::B
    order::O
    MergeSorted(a::A, b::B, order::O=Base.Order.Forward) where {A,B,O} =
        new{promote_type(eltype(A),eltype(B)),A,B,O}(a, b, order)
end

Base.eltype(::Type{MergeSorted{T,A,B,O}}) where {T,A,B,O} = T

Base.iterate(self::MergeSorted, (sa, sb) = ((), ())) =
    _iterate(self, sa, sb, iterate(self.a, sa...), iterate(self.b, sb...))

_iterate(self::MergeSorted{T}, sa, sb, ::Nothing, ::Nothing) where {T} = nothing
_iterate(self::MergeSorted{T}, sa, sb, (ra, sa´), ::Nothing) where {T} = T(ra), (sa´, sb)
_iterate(self::MergeSorted{T}, sa, sb, ::Nothing, (rb, sb´)) where {T} = T(rb), (sa, sb´)
_iterate(self::MergeSorted{T}, sa, sb, (ra, sa´), (rb, sb´)) where {T} =
    Base.Order.lt(self.order, ra, rb) ? (T(ra), (sa´, sb)) : (T(rb), (sa, sb´))

x = MergeSorted([1,4,5,9,32,44], [0,7,9,24,134])
sum(x); print(@allocated sum(x))

```

EDIT: it may still need some small tweak to avoid allocations also in the case of empty collections

EDIT2: Forget the above, one just needs to properly cast to eltype `T`. Fixed.

---

<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:** [July 13, 2021, 2:07am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/5 "2021-07-13T02:07:09Z")

</div>

Interesting solution, but does it require run-time dispatch to advance the iterator?

---

<div class="post-metadata">

**Author:** ![pablosanjose](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pablosanjose/32/7006_2.png) [@pablosanjose](https://discourse.julialang.org/u/pablosanjose)\
**Post date:** [July 13, 2021, 9:23am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/6 "2021-07-13T09:23:09Z")

</div>

No, it just avoids type-unstable states `(sa, sb)`. In the `x` example here it is always `Tuple{Int,Int}`.

---

<div class="post-metadata">

**Author:** ![mwallerb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mwallerb/32/52642_2.png) [@mwallerb](https://discourse.julialang.org/u/mwallerb)\
**Post date:** [July 20, 2021, 4:00pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/7 "2021-07-20T16:00:35Z")

</div>

@pablosanjose Thanks for neat and tidy solution!

Is there any way to understand why the added lookup is required here to avoid temporaries?

---

<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:** [July 20, 2021, 7:49pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/8 "2021-07-20T19:49:52Z")

</div>

The pre-1.0 iteration protocol of `start`/`next`/`done` was type-stable, so your problem would have had a straightforward type-stable non-allocating solution in the old protocol. The newer `iterate` protocol is predicated on the assumption that the compiler is able to efficiently implement certain type-unstable operations, in particular, containers of objects of type `Union{T,Nothing}`. So solving your problem with efficient code in the `iterate` protocol requires reliance on aggressive compiler optimization, which is apparently not available yet. At least, I don’t see a solution.

---

<div class="post-metadata">

**Author:** ![jzr](https://avatars.discourse-cdn.com/v4/letter/j/eb9ed0/32.png) [@jzr](https://discourse.julialang.org/u/jzr)\
**Post date:** [July 20, 2021, 8:15pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/9 "2021-07-20T20:15:45Z")

</div>

Tangentially, there is an interesting blog post about the state of the iterator protocol [https://mikeinnes.github.io/2020/06/04/iterate.html](https://mikeinnes.github.io/2020/06/04/iterate.html)

---

<div class="post-metadata">

**Author:** ![pablosanjose](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pablosanjose/32/7006_2.png) [@pablosanjose](https://discourse.julialang.org/u/pablosanjose)\
**Post date:** [July 20, 2021, 8:24pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/10 "2021-07-20T20:24:59Z")

</div>

What do you mean? The above solution is infered fine using union splitting. You just need to remain within the bounds of what unions can be split by the optimizer

---

<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:** [July 20, 2021, 8:48pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/11 "2021-07-20T20:48:27Z")

</div>

In your proposed solution, on each iteration, one of the four versions of `_iterate` is called. The compiler does not know in advance which version is invoked since the choice is data dependent. Therefore, the running code must choose at run-time on each iteration which one to call, correct? This leads to some loss of performance.

---

<div class="post-metadata">

**Author:** ![pablosanjose](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pablosanjose/32/7006_2.png) [@pablosanjose](https://discourse.julialang.org/u/pablosanjose)\
**Post date:** [July 20, 2021, 9:39pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/12 "2021-07-20T21:39:32Z")

</div>

Ah, but the key is that the return type of all `_iterate` functions is the same, regardless of which one is run. Well, to be precise, it is `Union{Nothing, Tuple{Int,Int}}`, so there are just two possible return types, but that is precisely what union splitting is designed to optimize (the two different type branches are precompiled, runtime dispatch is avoided, and you have essentially no runtime penalty).

The difference with the OP’s approach is that there, `iterate` returned a much more complex `Union` that shows up in red in `@code_warntype`, instead of yellow, signaling that you’re paying the price of runtime dispatch in that case. And the underlying reason is that the state of `iterate` is itself a `Union`.

So what I learned from this nice exercise is that to keep the compiler happy when using the post-1.0 iteration protocol is to construct your iterations so that `iterate` returns a `Union{Nothing,T}` for some concrete state type `T`, nothing much more complex than that, i.e. make your iteration states type-stable.

---

<div class="post-metadata">

**Author:** ![pablosanjose](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pablosanjose/32/7006_2.png) [@pablosanjose](https://discourse.julialang.org/u/pablosanjose)\
**Post date:** [July 20, 2021, 9:44pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/13 "2021-07-20T21:44:53Z")

</div>

> Is there any way to understand why the added lookup is required here to avoid temporaries?

Hi @mwallerb, maybe see the post above. The idea is “make `iterate` always return the same concrete type `T` unless it has finished its job, in which case return `nothing`”. So the proposed solution just returns the a tuple with the states for the `iterate` of `a` and `b`. When either of them is finished (no state is returned), instead of returning a tuple with a `nothing`, you return the last state. That way your state remains always `Tuple{Int,Int}`, instead of `Tuple{Union{Nothing,Int}, Union{Nothing,Int}}`, which is too complex for the optimized to handle at this point.

---

<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:** [July 20, 2021, 9:52pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/14 "2021-07-20T21:52:10Z")

</div>

I see, this is interesting! So the compiled code has all four versions of `_iterate` in-lined at the call site? Is there any way for the user to confirm this, short of poring over the machine code?

---

<div class="post-metadata">

**Author:** ![pablosanjose](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pablosanjose/32/7006_2.png) [@pablosanjose](https://discourse.julialang.org/u/pablosanjose)\
**Post date:** [July 20, 2021, 9:55pm UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/15 "2021-07-20T21:55:41Z")

</div>

I like that `@code_warntype` does this kind of interpretation for you with its color coding (yellow = union splitting, red = runtime dispatch), but perhaps there are other lower-level ways to see it directly, I’d also be interested to know!

---

<div class="post-metadata">

**Author:** ![mwallerb](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mwallerb/32/52642_2.png) [@mwallerb](https://discourse.julialang.org/u/mwallerb)\
**Post date:** [October 6, 2021, 8:50am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/16 "2021-10-06T08:50:33Z")

</div>

Sorry for resurrecting this thread now, but I just found out that @pablosanjose 's excellent solution in fact _does_ create temporaries for Julia \<= 1.5. It is only from Julia 1.6 that it does not (probably because of a more aggressively optimizing compiler).

All in all, iteration in julia seems **increadibly** brittle, because it is inherently type-unstable. This means you always have to hope and pray that your solution is simple enough for the compiler to be able to optimize out your instabilities.

I fail to understand the value of this. Why not, instead of using `iterate()`, have three methods, say:

```julia
IteratorElType(::I) = T
IteratorState(::I) = S
start(iter::I)::S = ... #Returns start state
stop(iter::I)::S = ... #Returns state at stop
next(iter::I, s::S)::Tuple{T,S} = ... #Returns next

```

and have `for x in iter; block; end` translate to:

```julia
state = start(iter)
while state != stop(iter)
    block
    state = next(state)
end

```

This is type-stable, super-easy to understand, and super-easy to compose upon. Now, if there is no good “stop” state, you can still use a `Union{Nothing,S}`. However, the vast majority of indexing in Julia is over scalar indices in an array, where this is total overkill. Any thoughts?

---

<div class="post-metadata">

**Author:** ![kristoffer.carlsson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/kristoffer.carlsson/32/22_2.png) [@kristoffer.carlsson](https://discourse.julialang.org/u/kristoffer.carlsson)\
**Post date:** [October 6, 2021, 9:23am UTC](https://discourse.julialang.org/t/avoid-creation-of-temporaries-in-non-trivial-iterator/64498/17 "2021-10-06T09:23:57Z")

</div>

> [@mwallerb](#):
>
> Any thoughts?

See for example [iterators that don't know they're done until they've tried · Issue #8149 · JuliaLang/julia · GitHub](https://github.com/JuliaLang/julia/issues/8149) and issues linked within.