Graph-topological iterators

treeiter.jl

using LightGraphs
using LightGraphs.SimpleGraphs: fadj
using BenchmarkTools
using Test
macro twice(body)
   quote $(esc(body));$(esc(body)) end
end

struct Done
end
struct First
end

# iteration protocol for SimpleDiGraphs

next(g::SimpleDiGraph, ::First) = one(eltype(g)), 1
function next(g::SimpleDiGraph, u)
    u >= nv(g) && return nothing
    u + 1, u + 1
end

next(g::SimpleDiGraph, u, ::First) = next(g::SimpleDiGraph, u, 1)
function next(g::SimpleDiGraph, u, i)
    i > length(fadj(g)[u]) && return Done()
    u => fadj(g)[u][i], u, i + 1
end

start(g::SimpleDiGraph, u) = u, 1


# Iteration protocol for array-like things seen as line graphs

const FastLinear = Union{Vector{T}, UnitRange{T}} where {T}

@inline function next(iter::FastLinear, ::First)
    isempty(iter) && return nothing
    @inbounds iter[1], 1
end
@inline next(iter::FastLinear, i) = @inbounds iter[i], i
@inline next(iter::FastLinear, i, ::First) = next(iter, i, false)
@inline function next(iter::FastLinear, i, done)
    done && return Done()
    i >= length(iter) && return nothing
    #@inbounds (A[i], A[i + 1]), i + 1, Done()
    @inbounds (iter[i], iter[i + 1]), i + 1, true
end
@inline function next(::FastLinear, _, ::Done)
    return Done()
end

@inline start(iter::FastLinear, i) = i, false


# any iterator can be seen as a line graph

@inline function next(iter, ::First)
    ϕ = iterate(iter)
    ϕ === nothing && return nothing
    ϕ[1], ϕ
end
@inline next(iter, ϕ) =  ϕ[1], ϕ
@inline start(iter, state) = state, false
@inline next(iter, vstate, ::First) = next(iter, vstate, false)
@inline function next(iter, vstate, done)
    #done && return Done()
    v, state = vstate
    ϕ = iterate(iter, state)
    ϕ === nothing && return nothing
    vnew, state = ϕ
    return (v, vnew), ϕ, Done()
    #return (v, vnew), ϕ, true
end
@inline function next(iter, vstate, ::Done)
    return Done()
end

usage.jl

function dosomething(g)
    ϕ = next(g, First())
    ϕ === nothing && return default
    v, u = ϕ

    # do something with the root v

    ϕ = next(g, u, First()) # next edge
    while ϕ !== nothing # while not finished
        if ϕ === Done() # done iterating through all outgoing edges of v
            ψ = next(g, u) # next vertex
            ψ === nothing && break
            v, u = ψ

            # do something with vertex v

            ϕ = next(g, u, First())
        else
            e, u, i = ϕ

            # do something with edge e

            ϕ = next(g, u, i) # next edge
        end
    end
    return result
end