Creating an iterator by wrapping another one

General Usage Performance
1

I am trying to create an iterator that wraps another one. Let us assume that we have this iterator that returns l times the array a :

struct Foo
    a::Array{Float64,1}
    l::Int
end

Base.iterate(foo::Foo) = iterate(foo,0)
function Base.iterate(foo::Foo,state::Int)
    if foo.l == state
        nothing
    else
        k = state + 1
        (foo.a,k)
    end
end

My goal is to implement another iterator that wraps previous one:

struct Bar
    foo::Foo
end

function Base.iterate(bar::Bar)
    foonext = iterate(bar.foo)
    iterate(bar,foonext)
end
function Base.iterate(bar::Bar,state)
    if state == nothing
        nothing
    else
        (fooi,foostate) = state
        foonext = iterate(bar.foo,foostate)
        (fooi,foonext)
    end
end

Iterating over instances of Foo works fine (as expected):

function run(x)
    for xi in x
    end
end
foo = Foo([1.2,3.2,1.1],10000000)
@time run(foo)
@time run(foo)

  0.288515 seconds (12.11 k allocations: 750.146 KiB)
  0.274807 seconds (4 allocations: 160 bytes)

The problems is that by iterating instances of Bar I get unwanted memory allocation:

bar = Bar(foo)
@time run(bar)
@time run(bar)

  0.795432 seconds (10.02 M allocations: 306.389 MiB, 3.56% gc time)
  0.669443 seconds (10.00 M allocations: 305.176 MiB, 0.99% gc time)

I would say that the code above is type stable…what I am missing?

Thanks in advance for the help!

2

I think it is not type stable because foonext can be nothing. You may want to check with @code_warntype iterate(bar).

3

Could do something like

function Base.iterate(bar::Bar)
    foonext = iterate(bar.foo)
    foonext === nothing && return nothing
    return iterate(bar, foonext)
end

function Base.iterate(bar::Bar, state)
    fooi, foostate = state
    foonext = iterate(bar.foo, foostate)
    foonext === nothing && return nothing
    return fooi, foonext
end

julia> @time run(bar)
  0.000003 seconds (4 allocations: 160 bytes)

Basically, every time we call iterate on the inner iterator we check if it is nothing (with ===) and handle it. The compiler can then assume from thereon that the return value of the iterator is not nothing.

2 Likes
4

IterTools.@ifsomething is a macro specifically for this pattern.

2 Likes
5

A Further optimization is to check for nothing before and after the inner iterate and return fooi, foostate or nothing.

6

What do you mean before the inner innerate? Can you write out the code?

7

Thanks for the answer!

It almost fixes it. However, the suggested implementation of Bar leads to instances that iterate one iteration less that the wrapped instance of Foo.

8

I have it!

It can be done very simply by using the foostate directly instead of creating a new state.

function Base.iterate(bar::Bar)
    foonext = iterate(bar.foo)
    foonext === nothing && return nothing
    foonext
end
function Base.iterate(bar::Bar,foostate)
    foonext = iterate(bar.foo,foostate)
    foonext === nothing && return nothing
    foonext
end

julia> @time run(bar)
  0.000003 seconds (4 allocations: 160 bytes)

Thanks for all your comments!

9

This is equivalent to

Base.iterate(bar::Bar) = iterate(bar.foo)
5 Likes
10

Yes, thanks for noting!

At the end, the solution to my first question is very simple: One has to wrap the calls to iterate . That’s all.