Best first search

We perform best first search of a tree or graph, where the neighbours of a vertex are provided by a lazy list α → MLList m α.

We maintain a priority queue of visited-but-not-exhausted nodes, and at each step take the next child of the highest priority node in the queue.

This is useful in meta code for searching for solutions in the presence of alternatives. It can be nice to represent the choices via a lazy list, so the later choices don't need to be evaluated while we do depth first search on earlier choices.

Options:

• maxDepth allows bounding the search depth
• maxQueued implements "beam" search, by discarding elements from the priority queue when it grows too large
• removeDuplicatesBy? maintains an RBSet of previously visited nodes; otherwise if the graph is not a tree nodes may be visited multiple times.

We begin by defining a best-first queue of MLLists. This is a somewhat baroque data structure designed for the application in this file (and in particularly for the implementation of rewrite_search). If someone would like to generalize appropriately that would be great.

We want to maintain a priority queue of MLList m β, each indexed by some a : α with a priority. (One could simplify matters here by simply flattening this out to a priority queue of pairs α × β, with the priority determined by the α factor. However the laziness of MLList is essential to performance here: we will extract elements from these lists one at a time, and only when they at the head of the queue. If another item arrives at the head of the queue, we may not need to continue calculate the previous head's elements.)

To complicate matters, the priorities might be expensive to calculate, so we instead keep track of a lower bound (where less is better) for each such a : α. The priority queue maintains the MLList m β in order of the current best lower bound for the corresponding a : α. When we insert a new α × MLList m β into the queue, we have to provide a lower bound, and we just insert it at a position depending on the estimate. When it is time to pop a β off the queue, we iteratively improve the lower bound for the front element of the queue, until we decide that either it must be the least element, or we can exchange it with the second element of the queue and continue.

A BestFirstQueue prio ε m β maxSize consists of an TreeMap, where the keys are in BestFirstNode prio ε and the values are MLList m β.

A BestFirstNode prio ε consists of a key : α and an estimator ε : key. Here ε provides the current best lower bound for prio key : Thunk ω. (The actual priority is hidden behind a Thunk to avoid evaluating it, in case it is expensive.)

We ask for the type classes LinearOrder ω and ∀ a : α, Estimator (prio a) (ε a). This later typeclass ensures that we can always produce progressively better estimates for a priority. We also need a WellFounded instance to ensure that improving estimates terminates.

This whole structure is designed around the problem of searching rewrite graphs, prioritising according to edit distances (either between sides of an equation, or from a goal to a target). Edit distance computations are particularly suited to this design because the usual algorithm for computing them produces improving lower bounds at each step.

With this design, it is okay if we visit nodes with very large edit distances: while these would be expensive to compute, we never actually finish the computation except in cases where the node arrives at the front of the queue.

structure BestFirstNode {α : Sort u_1} {ω : Type u_2} (prio : α → Thunk ω) (ε : α → Type) : Sort (max 1 u_1)

A node in a BestFirstQueue.

• key : α

  The data to store at a node, from which we can calculate a priority using prio.

• estimator : ε self.key

  An estimator for the priority of the key. (We will assume we have [∀ a : α, Estimator (prio a) (ε a)].)

def BestFirstNode.estimate {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] (n : BestFirstNode prio ε) : ω

Calculate the current best lower bound for the priority of a node.

Equations

• n.estimate = EstimatorData.bound (prio n.key) n.estimator

instance instOrdBestFirstNode {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [Ord ω] [Ord α] : Ord (BestFirstNode prio ε)

Equations

• instOrdBestFirstNode = { compare := compareLex (compareOn BestFirstNode.estimate) (compareOn BestFirstNode.key) }

def BestFirstQueue {ω α : Type} (prio : α → Thunk ω) (ε : α → Type) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] (m : Type → Type) (β : Type) [Ord ω] [Ord α] (maxSize : Option ℕ) : Type

A queue of MLList m βs, lazily prioritized by lower bounds.

Equations

• BestFirstQueue prio ε m β maxSize = Std.TreeMap (BestFirstNode prio ε) (MLList m β) compare

def BestFirstQueue.insertAndEject {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] {m : Type → Type} {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) (n : BestFirstNode prio ε) (l : MLList m β) : BestFirstQueue prio ε m β maxSize

Add a new MLList m β to the BestFirstQueue, and if this takes the size above maxSize, eject a MLList from the tail of the queue.

Equations

• One or more equations did not get rendered due to their size.
• q_2.insertAndEject n l = Std.TreeMap.insert q_2 n l

partial def BestFirstQueue.ensureFirstIsBest {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] {m : Type → Type} [Monad m] {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) : m (BestFirstQueue prio ε m β maxSize)

By improving priority estimates as needed, and permuting elements, ensure that the first element of the queue has the greatest priority.

partial def BestFirstQueue.popWithBound {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] {m : Type → Type} [Monad m] {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) : m (Option (((a : α) × ε a × β) × BestFirstQueue prio ε m β maxSize))

Pop a β off the MLList m β with lowest priority, also returning the index in α and the current best lower bound for its priority. This may require improving estimates of priorities and shuffling the queue.

def BestFirstQueue.popWithPriority {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] {m : Type → Type} [Monad m] {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) : m (Option (((α × ω) × β) × BestFirstQueue prio ε m β maxSize))

Pop a β off the MLList m β with lowest priority, also returning the index in α and the value of the current best lower bound for its priority.

Equations

• q.popWithPriority = do let __do_lift ← q.popWithBound match __do_lift with | none => pure none | some (⟨a, (e, b)⟩, q') => pure (some (((a, EstimatorData.bound (prio a) e), b), q'))

def BestFirstQueue.pop {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] {m : Type → Type} [Monad m] {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) : m (Option ((α × β) × BestFirstQueue prio ε m β maxSize))

Pop a β off the MLList m β with lowest priority.

Equations

• q.pop = do let __do_lift ← q.popWithBound match __do_lift with | none => pure none | some (⟨a, (fst, b)⟩, q') => pure (some ((a, b), q'))

partial def BestFirstQueue.toMLListWithPriority {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] {m : Type → Type} [Monad m] {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) : MLList m ((α × ω) × β)

Convert a BestFirstQueue to a MLList ((α × ω) × β), by popping off all elements, recording also the values in ω of the best current lower bounds.

def BestFirstQueue.toMLList {ω α : Type} {prio : α → Thunk ω} {ε : α → Type} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] {m : Type → Type} [Monad m] {β : Type} [Ord ω] [Ord α] {maxSize : Option ℕ} (q : BestFirstQueue prio ε m β maxSize) : MLList m (α × β)

Convert a BestFirstQueue to a MLList (α × β), by popping off all elements.

Equations

• q.toMLList = MLList.map (fun (t : (α × ω) × β) => (t.1.1, t.2)) q.toMLListWithPriority

def impl {ω α : Type} (prio : α → Thunk ω) (ε : α → Type) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : Type → Type} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (maxSize : Option ℕ) (f : α → MLList m α) (a : α) : MLList m α

Core implementation of bestFirstSearch, that works by iteratively updating an internal state, consisting of a priority queue of MLList m α.

At each step we pop an element off the queue, compute its children (lazily) and put these back on the queue.

Equations

• impl prio ε maxSize f a = MLList.cons a ((MLList.iterate (impl.go prio ε maxSize f)).runState' (Std.TreeMap.empty.insert { key := a, estimator := ⊥ } (f a)))

def impl.go {ω α : Type} (prio : α → Thunk ω) (ε : α → Type) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : Type → Type} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (maxSize : Option ℕ) (f : α → MLList m α) : BestFirstQueue prio ε m α maxSize → m (α × BestFirstQueue prio ε m α maxSize)

A single step of the best first search. Pop an element, and insert its children back into the queue, with a trivial estimator for their priority.

Equations

• impl.go prio ε maxSize f s = do let __do_lift ← s.pop match __do_lift with | none => failure | some ((fst, b), s') => pure (b, s'.insertAndEject { key := b, estimator := ⊥ } (f b))

def implMaxDepth {ω α : Type} (prio : α → Thunk ω) (ε : α → Type) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : Type → Type} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (maxSize maxDepth : Option ℕ) (f : α → MLList m α) (a : α) : MLList m α

Wrapper for impl implementing the maxDepth option.

Equations

• One or more equations did not get rendered due to their size.
• implMaxDepth prio ε maxSize none f a = impl prio ε maxSize f a

def bestFirstSearchCore {ω α : Type} (prio : α → Thunk ω) (ε : α → Type) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT ↑(Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : Type → Type} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (f : α → MLList m α) (a : α) (β : Type) [Ord β] (removeDuplicatesBy? : Option (α → β) := none) (maxQueued maxDepth : Option ℕ := none) : MLList m α

A lazy list recording the best first search of a graph generated by a function f : α → MLList m α.

We maintain a priority queue of visited-but-not-exhausted nodes, and at each step take the next child of the highest priority node in the queue.

The option maxDepth limits the search depth.

The option maxQueued bounds the size of the priority queue, discarding the lowest priority nodes as needed. This implements a "beam" search, which may be incomplete but uses bounded memory.

The option removeDuplicates keeps an RBSet of previously visited nodes. Otherwise, if the graph is not a tree then nodes will be visited multiple times.

This version allows specifying a custom priority function prio : α → Thunk ω along with estimators ε : α → Type equipped with [∀ a, Estimator (prio a) (ε a)] that control the behaviour of the priority queue. This function returns values a : α that have the lowest possible prio a amongst unvisited neighbours of visited nodes, but lazily estimates these priorities to avoid unnecessary computations.

Equations

• One or more equations did not get rendered due to their size.
• bestFirstSearchCore prio ε f a β none maxQueued maxDepth = implMaxDepth prio ε maxQueued maxDepth f a

def instOrderBotEq {α : Type} [LinearOrder α] {a : α} : OrderBot { x : α // x = a }

A local instance that enables using "the actual value" as a priority estimator, for simple use cases.

Equations

• instOrderBotEq = { bot := ⟨a, ⋯⟩, bot_le := ⋯ }

def bestFirstSearch {m : Type → Type} {α : Type} [Monad m] [Alternative m] [LinearOrder α] (f : α → MLList m α) (a : α) (maxQueued maxDepth : Option ℕ := none) (removeDuplicates : Bool := true) : MLList m α

A lazy list recording the best first search of a graph generated by a function f : α → MLList m α.

We maintain a priority queue of visited-but-not-exhausted nodes, and at each step take the next child of the highest priority node in the queue.

The option maxDepth limits the search depth.

The option maxQueued bounds the size of the priority queue, discarding the lowest priority nodes as needed. This implements a "beam" search, which may be incomplete but uses bounded memory.

The option removeDuplicates keeps an RBSet of previously visited nodes. Otherwise, if the graph is not a tree then nodes will be visited multiple times.

This function returns values a : α that are least in the [LinearOrder α] amongst unvisited neighbours of visited nodes.

Equations

• One or more equations did not get rendered due to their size.