Simple backtracking search over iterators in julia

backtrack.jl

struct BacktrackSearch{V<:AbstractVector, VI<:AbstractVector, VS<:AbstractVector, O, A}
    nodes::V
    iterators::VI
    states::VS
    oracle_descend::O
    node_accept::A
end

"""
BacktrackSearch(iterators; oracle_descend, node_accept)

An object to perform a backtrack search on the tree defined by `iterators`.

Iterating over `bs::BacktrackSearch` returns nodes, i.e. vectors of elements
obtained from `iterators`.
* `oracle_descend(bs::BacktrackSearch)` controls when to descend into the children of a node,
* `node_accept` controls when a node is returned.
"""
function BacktrackSearch(iterators; oracle_descend, node_accept)
    @assert !isempty(iterators)
    @assert all(!isempty, iterators)
    nodes, states = let itr = first(iterators)
        T = eltype(itr)
        @assert all(it -> eltype(it) == T, iterators)
        S = typeof(last(iterate(itr)))
        Vector{T}(), Vector{S}()
    end

    return BacktrackSearch(nodes, iterators, states, oracle_descend, node_accept)
end

current_depth(bs::BacktrackSearch) = length(bs.states)
depth(bs::BacktrackSearch) = length(bs.iterators)

function Base.pop!(bs::BacktrackSearch)
    return pop!(bs.states), pop!(bs.nodes)
end

function Base.push!(bs::BacktrackSearch, child_state::Tuple)
    child, state = child_state
    push!(bs.nodes, child)
    push!(bs.states, state)
    return child_state
end

# the unsafe version would just return bs.nodes
# each element in interation must then be explicitly copied by the caller
node(bs::BacktrackSearch) = copy(bs.nodes)
oracle_descend(bs::BacktrackSearch) = bs.oracle_descend(bs)
node_accept(bs::BacktrackSearch) = bs.node_accept(bs)

isleaf(bs::BacktrackSearch) = current_depth(bs)==depth(bs)

struct EndsWithAnyOf{T}
    elts::T
end

(ewa::EndsWithAnyOf)(bs::BacktrackSearch) = last(bs.nodes) in ewa.elts

Base.eltype(::Type{<:BacktrackSearch{V}}) where V = V
Base.IteratorSize(::Type{<:BacktrackSearch}) = Base.SizeUnknown()

function Base.iterate(bs::BacktrackSearch)
    # initialize the backtrack object here
    resize!(bs.nodes, 0)
    resize!(bs.states, 0)
    return iterate(bs, false)
end

function Base.iterate(bs::BacktrackSearch, backtrack)
    # return (yield!) sets backtrack=true, so jump to the backtracking part immediately
    backtrack && @goto BACKTRACK

    while !backtrack
        backtrack = !oracle_descend(bs)
        backtrack && @debug "node was rejected by the oracle_descend" node(bs)
        if !backtrack && node_accept(bs)
            @debug "accepted node is returned" node(bs)
            return node(bs), true
        end
        if !backtrack
            cdepth = current_depth(bs)
            @debug "at depth=$cdepth: entering depth=$(cdepth + 1)"
            child_state = iterate(bs.iterators[cdepth + 1])
            # we're assuming that iterators are all non-empty. otherwise we'd need to
            # isnothing(child_state) && @goto BACKTRACK
            push!(bs, child_state)
            @debug "at depth=$(current_depth(bs)): going to the first child ($(first(child_state)))"
        else
            @label BACKTRACK
            while backtrack && !isempty(bs.nodes)
                cdepth = current_depth(bs)
                @debug "at depth=$cdepth: exploring next node"
                child_state = iterate(bs.iterators[cdepth], bs.states[cdepth])
                pop!(bs)
                if isnothing(child_state)
                    @debug "at depth=$cdepth: explored all childred, backtracking"
                else
                    push!(bs, child_state)
                    @debug "at depth=$cdepth: going to the next child ($(first(child_state)))"
                    backtrack = false
                end
            end
        end
    end
    return nothing
end

backtrack.log

julia> ENV["JULIA_DEBUG"]="Main"
"Main"

julia> bs = BacktrackSearch(
           ['a':'c', 'x':'z'],
           oracle_descend=bs->isempty(bs.nodes) || (last(bs.nodes) ≠ 'b' && last(bs.nodes) ≠ 'y'),
           node_accept=bs->current_depth(bs)==depth(bs)
       );

julia> collect(bs)
┌ Debug: at depth=0: entering depth=1
└ @ Main /tmp/backtrack.jl:83
┌ Debug: at depth=1: going to the first child (a)
└ @ Main /tmp/backtrack.jl:88
┌ Debug: at depth=1: entering depth=2
└ @ Main /tmp/backtrack.jl:83
┌ Debug: at depth=2: going to the first child (x)
└ @ Main /tmp/backtrack.jl:88
┌ Debug: accepted node is returned
│   node(bs) =
│    2-element Vector{Char}:
│     'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
│     'x': ASCII/Unicode U+0078 (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:78
┌ Debug: at depth=2: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=2: going to the next child (y)
└ @ Main /tmp/backtrack.jl:100
┌ Debug: node was rejected by the oracle_descend
│   node(bs) =
│    2-element Vector{Char}:
│     'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
│     'y': ASCII/Unicode U+0079 (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:76
┌ Debug: at depth=2: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=2: going to the next child (z)
└ @ Main /tmp/backtrack.jl:100
┌ Debug: accepted node is returned
│   node(bs) =
│    2-element Vector{Char}:
│     'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
│     'z': ASCII/Unicode U+007A (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:78
┌ Debug: at depth=2: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=2: explored all childred, backtracking
└ @ Main /tmp/backtrack.jl:97
┌ Debug: at depth=1: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=1: going to the next child (b)
└ @ Main /tmp/backtrack.jl:100
┌ Debug: node was rejected by the oracle_descend
│   node(bs) =
│    1-element Vector{Char}:
│     'b': ASCII/Unicode U+0062 (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:76
┌ Debug: at depth=1: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=1: going to the next child (c)
└ @ Main /tmp/backtrack.jl:100
┌ Debug: at depth=1: entering depth=2
└ @ Main /tmp/backtrack.jl:83
┌ Debug: at depth=2: going to the first child (x)
└ @ Main /tmp/backtrack.jl:88
┌ Debug: accepted node is returned
│   node(bs) =
│    2-element Vector{Char}:
│     'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
│     'x': ASCII/Unicode U+0078 (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:78
┌ Debug: at depth=2: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=2: going to the next child (y)
└ @ Main /tmp/backtrack.jl:100
┌ Debug: node was rejected by the oracle_descend
│   node(bs) =
│    2-element Vector{Char}:
│     'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
│     'y': ASCII/Unicode U+0079 (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:76
┌ Debug: at depth=2: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=2: going to the next child (z)
└ @ Main /tmp/backtrack.jl:100
┌ Debug: accepted node is returned
│   node(bs) =
│    2-element Vector{Char}:
│     'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)
│     'z': ASCII/Unicode U+007A (category Ll: Letter, lowercase)
└ @ Main /tmp/backtrack.jl:78
┌ Debug: at depth=2: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=2: explored all childred, backtracking
└ @ Main /tmp/backtrack.jl:97
┌ Debug: at depth=1: exploring next node
└ @ Main /tmp/backtrack.jl:93
┌ Debug: at depth=1: explored all childred, backtracking
└ @ Main /tmp/backtrack.jl:97
4-element Vector{Vector{Char}}:
 ['a', 'x']
 ['a', 'z']
 ['c', 'x']
 ['c', 'z']