Avoid creation of temporaries in non-trivial iterator

General Usage Performance
, ,
4

What an interesting little problem!

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

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.

2 Likes