Switch to heap-based priority queue in RPG search.

src/heuristics/pgraph.jl

@@ -310,16 +310,19 @@ function relaxed_pgraph_search(domain::Domain, state::State, spec::Specification
 
     # Set up initial facts and priority queue
     init_idxs = pgraph_init_idxs(graph, domain, state)
-    dists[init_idxs] .= 0
-    costs[init_idxs] .= 0
-    queue = PriorityQueue{Int,Float32}(i => 0 for i in findall(init_idxs))
+    dists[init_idxs] .= 0.0f0
+    costs[init_idxs] .= 0.0f0
+    queue = FastPriorityQueue{Int,Float32}()
+    append!(queue, (i => 0.0f0 for i in findall(init_idxs)))
 
     # Perform Djikstra / uniform-cost search until goals are reached
     goal_idx, goal_cost = nothing, Inf32
     last_nongoal_idx = n_actions - graph.n_goals
     while !isempty(queue) && isnothing(goal_idx)
         # Dequeue nearest fact/condition
-        cond_idx = dequeue!(queue)
+        cond_idx, cond_dist = dequeue_pair!(queue)
+        # Skip if distance greater than stored value
+        cond_dist > dists[cond_idx] && continue
         # Iterate over child actions
         for (act_idx, precond_idx) in graph.cond_children[cond_idx]
             # Check if goal action is reached
@@ -373,10 +376,8 @@ function relaxed_pgraph_search(domain::Domain, state::State, spec::Specification
                     costs[c_idx] = next_cost
                     achievers[c_idx] = act_idx
                 end
-                if !(c_idx in keys(queue)) # Enqueue new conditions
+                if less_dist # Adjust distances and place on queue
                     enqueue!(queue, c_idx, next_dist)
-                elseif less_dist # Adjust distances
-                    queue[c_idx] = next_dist
                     dists[c_idx] = next_dist
                 end
             end

src/utils.jl

@@ -1,5 +1,7 @@
 # Miscellaneous utilities
 
+using DataStructures: DataStructures, BinaryHeap
+
 """
     lazy_collect(collection)
     lazy_collect(element_type, collection)
@@ -177,3 +179,67 @@ function show_struct(
     end
     nothing
 end
+
+"""
+    FastPriorityQueue{K, V}([ord])
+
+Fast heap-based priority queue implementation which supports enqueuing and
+dequeueing, but not priority adjustment.
+"""
+struct FastPriorityQueue{K, V, O <: Base.Ordering} <: AbstractDict{K, V}
+    heap::BinaryHeap{Pair{K, V}, Base.Order.By{typeof(last), O}}
+end 
+
+function FastPriorityQueue{K, V}(
+    ord::Base.Ordering = DataStructures.FasterForward()
+) where {K, V}
+    heap = BinaryHeap{Pair{K, V}}(Base.By(last, ord))
+    return FastPriorityQueue{K, V, typeof(ord)}(heap)
+end
+
+function FastPriorityQueue{K, V}(
+    ord::Base.Ordering, pair::Union{Pair{K, V}, Tuple{K, V}},
+    pairs::Union{Pair{K, V}, Tuple{K, V}}...
+) where {K, V}
+    pq = FastPriorityQueue{K, V}(ord)
+    enqueue!(pq, pair)
+    for kv in pairs
+        enqueue!(pq, kv)
+    end
+    return pq
+end
+
+function FastPriorityQueue{K, V}(ord::Base.Ordering, pairs) where {K, V}
+    pairs = pairs isa AbstractVector ? pairs : collect(pairs)
+    heap = BinaryHeap{Pair{K, V}}(Base.By(last, ord), pairs)
+    return FastPriorityQueue{K, V, typeof(ord)}(heap)
+end
+
+FastPriorityQueue{K, V}(pairs) where {K, V} =
+    FastPriorityQueue{K, V}(DataStructures.FasterForward(), pairs)
+
+FastPriorityQueue(pairs::AbstractArray{Pair{K, V}}) where {K, V} =
+    FastPriorityQueue{K, V}(DataStructures.FasterForward(), pairs)
+
+Base.length(pq::FastPriorityQueue) = length(pq.heap)
+Base.isempty(pq::FastPriorityQueue) = isempty(pq.heap)
+Base.first(pq::FastPriorityQueue) = first(pq.heap)
+Base.peek(pq::FastPriorityQueue) = peek(pq.heap)
+
+Base.empty!(pq::FastPriorityQueue) =  (empty!(pq.heap.valtree); pq)
+
+function Base.append!(pq::FastPriorityQueue, pairs)
+    empty!(pq.heap.valtree)
+    append!(pq.heap.valtree, pairs)
+    DataStructures.heapify!(pq.heap.valtree, pq.heap.ordering)
+    return pq
+end
+
+DataStructures.enqueue!(pq::FastPriorityQueue{K, V}, kv::Pair{K, V}) where {K, V} =
+    (push!(pq.heap, kv); return pq)
+DataStructures.enqueue!(pq::FastPriorityQueue{K, V}, k::K, v::V) where {K, V} =
+    (push!(pq.heap, k => v); return pq)
+DataStructures.dequeue!(pq::FastPriorityQueue) =
+    first(pop!(pq.heap))
+DataStructures.dequeue_pair!(pq::FastPriorityQueue) =
+    pop!(pq.heap)