Documentation

Mathlib.Data.MLList.BestFirst

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:

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 lazyness 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 RBMap, 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)].)

Instances For
    def BestFirstNode.estimate {ω : Type u_1} {α : 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
    Instances For
      instance instOrdBestFirstNode {ω : Type u_1} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [Ord ω] [Ord α] :
      Ord (BestFirstNode prio ε)
      Equations
      def BestFirstQueue {ω : Type u_1} {α : Type} (prio : αThunk ω) (ε : αType) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] (m : TypeType) (β : Type) [Ord ω] [Ord α] (maxSize : Option ) :

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

      Equations
      Instances For
        def BestFirstQueue.insertAndEject {ω : Type u_1} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] {m : TypeType} {β : 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.
        Instances For
          partial def BestFirstQueue.ensureFirstIsBest {ω : Type u_1} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] {m : TypeType} [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 u_1} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] {m : TypeType} [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} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] {m : TypeType} [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_liftq.popWithBound match __do_lift with | none => pure none | some (a, (e, b), q') => pure (some (((a, EstimatorData.bound (prio a) e), b), q'))
          Instances For
            def BestFirstQueue.pop {ω : Type u_1} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] {m : TypeType} [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_liftq.popWithBound match __do_lift with | none => pure none | some (a, (fst, b), q') => pure (some ((a, b), q'))
            Instances For
              partial def BestFirstQueue.toMLListWithPriority {ω : Type} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] {m : TypeType} [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} {α : Type} {prio : αThunk ω} {ε : αType} [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] {m : TypeType} [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
              Instances For
                def impl {ω : Type u_1} {α : Type} (prio : αThunk ω) (ε : αType) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : TypeType} [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
                Instances For
                  def impl.go {ω : Type u_1} {α : Type} (prio : αThunk ω) (ε : αType) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : TypeType} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (maxSize : Option ) (f : αMLList m α) :
                  BestFirstQueue prio ε m α maxSizem (α × 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_lifts.pop match __do_lift with | none => failure | some ((fst, b), s') => pure (b, s'.insertAndEject { key := b, estimator := } (f b))
                  Instances For
                    def implMaxDepth {ω : Type u_1} {α : Type} (prio : αThunk ω) (ε : αType) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : TypeType} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (maxSize : Option ) (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.
                    Instances For
                      def bestFirstSearchCore {ω : Type u_1} {α : Type} (prio : αThunk ω) (ε : αType) [LinearOrder ω] [(a : α) → Estimator (prio a) (ε a)] [I : ∀ (a : α), WellFoundedGT (Set.range (EstimatorData.bound (prio a)))] [Ord ω] [Ord α] {m : TypeType} [Monad m] [Alternative m] [(a : α) → Bot (ε a)] (f : αMLList m α) (a : α) (β : Type) [Ord β] (removeDuplicatesBy? : optParam (Option (αβ)) none) (maxQueued : optParam (Option ) none) (maxDepth : optParam (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.
                      Instances For
                        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
                        Instances For
                          def bestFirstSearch {m : TypeType} {α : Type} [Monad m] [Alternative m] [LinearOrder α] (f : αMLList m α) (a : α) (maxQueued : optParam (Option ) none) (maxDepth : optParam (Option ) none) (removeDuplicates : optParam 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.
                          Instances For