Iterative construction of lazy Iterators leads to excessive compilation time

General Usage Performance
, ,
21

I’m adding another version using the Deque data structure. For this specific example it might be slightly slower, because the test graph is linear and the size of the ‘frontier’ is small. But on more expanding graphs, the popfirst! becomes alogrithmically and practically faster with a Deque.

Additionally, using some default parameters, the whole iterator is one function which is nicer:

using DataStructures

struct ImplicitGraph{F<:Function}
    adj::F
end

struct TraversalBFS{T,F}
    graph::ImplicitGraph{F}
    start::T
end

@inline function Base.iterate(bfs::TraversalBFS{T,F}, 
  state=(push!(Deque{T}(),bfs.start), Set{T}())) where {T,F}
    Q, explored = state
    isempty(Q) && return nothing
    v = popfirst!(Q)
    for w in bfs.graph.adj(v)
        w ∈ explored && continue
        push!(Q, w)
    end
    push!(explored, v)
    return (v, state)
end

and the benchmark:

julia> G = ImplicitGraph(v -> (v - 1, v + 1))
ImplicitGraph{var"#3#4"}(var"#3#4"())

julia> t = TraversalBFS(G, 0)
TraversalBFS{Int64, var"#3#4"}(ImplicitGraph{var"#3#4"}(var"#3#4"()), 0)

julia> function test_iterator(t, n)
           s = 0
           i = 0
           for v in t
               (i = i + 1) > n && break
               s += v
           end
           s
       end
test_iterator (generic function with 1 method)

julia> test_iterator(t, 100)
-50

julia> @btime test_iterator(t, 100);
  1.648 μs (13 allocations: 11.86 KiB)

To make an example for the performance difference, let’s run on a 2D grid graph:

julia> G2 = ImplicitGraph(((x,y),) -> ((x+1,y),(x,y+1),(x-1,y),(x,y-1)))
ImplicitGraph{var"#7#8"}(var"#7#8"())

julia> t2 = TraversalBFS(G2, (0,0))
TraversalBFS{Tuple{Int64, Int64}, var"#7#8"}(ImplicitGraph{var"#7#8"}(var"#7#8"()), (0, 0))

Some tweaks needed to test_ routines. Eventually:

julia> @btime test_direct(t2, 10000);
  564.711 μs (30 allocations: 407.39 KiB)

julia> @btime test_iterator(t2, 10000);
  282.938 μs (55 allocations: 347.31 KiB)

And the advantage to the Deque based iterator increases with node count.

2 Likes
22

Sorry for the late reply. Thanks for the suggestions! :slight_smile: