metagraphoptimization.jl/src/trie.jl

# helper struct for NodeTrie
mutable struct NodeIdTrie
    value::Vector{Node}
    children::Dict{UUID, NodeIdTrie}
end

# Trie data structure for node reduction, inserts nodes by children
# Assumes that given nodes have ordered vectors of children (see sort_node)
# First level is the task type and thus does not have a value
# Should be constructed with all Types that will be used
mutable struct NodeTrie
    children::Dict{DataType, NodeIdTrie}
end

function NodeTrie()
    return NodeTrie(Dict{DataType, NodeIdTrie}())
end

function NodeIdTrie()
    return NodeIdTrie(Vector{Node}(), Dict{UUID, NodeIdTrie}())
end

function insert_helper!(trie::NodeIdTrie, node::Node, depth::Int)
    if (length(node.children) == depth)
        push!(trie.value, node)
        return nothing
    end

    depth = depth + 1
    id = node.children[depth].id

    if (!haskey(trie.children, id))
        trie.children[id] = NodeIdTrie()
    end
    return insert_helper!(trie.children[id], node, depth)
end

function insert!(trie::NodeTrie, node::Node)
    t = typeof(node.task)
    if (!haskey(trie.children, t))
        trie.children[t] = NodeIdTrie()
    end
    return insert_helper!(trie.children[typeof(node.task)], node, 0)
end

function collect_helper(trie::NodeIdTrie, acc::Set{Vector{Node}})
    if (length(trie.value) >= 2)
        push!(acc, trie.value)
    end

    for (id, child) in trie.children
        collect_helper(child, acc)
    end
    return nothing
end

# returns all sets of multiple nodes that have accumulated in leaves
function collect(trie::NodeTrie)
    acc = Set{Vector{Node}}()
    for (t, child) in trie.children
        collect_helper(child, acc)
    end
    return acc
end