feat: depth first and best first search using ListM (#3221)

Implementations of depth first search, best first search, and beam search, for graphs described by a neighbours function `α → ListM m α`. There are also wrappers for using `α → List α`.

This is only intended for use in meta code.



Co-authored-by: Scott Morrison <scott.morrison@gmail.com>

Author: kim-em

Mathlib.lean

@@ -861,7 +861,9 @@ import Mathlib.Data.List.Sublists
 import Mathlib.Data.List.TFAE
 import Mathlib.Data.List.ToFinsupp
 import Mathlib.Data.List.Zip
-import Mathlib.Data.ListM
+import Mathlib.Data.ListM.Basic
+import Mathlib.Data.ListM.BestFirst
+import Mathlib.Data.ListM.DepthFirst
 import Mathlib.Data.ListM.Heartbeats
 import Mathlib.Data.Matrix.Basic
 import Mathlib.Data.Matrix.Basis

Mathlib/Data/ListM.lean renamed to Mathlib/Data/ListM/Basic.lean

@@ -85,7 +85,7 @@ unsafe def fixl [Alternative m] (f : α → m (α × List β)) (s : α) : ListM
 
 /-- Deconstruct a `ListM`, returning inside the monad an optional pair `α × ListM m α`
 representing the head and tail of the list. -/
-unsafe def uncons {α : Type u} : ListM m α → m (Option (α × ListM m α))
+unsafe def uncons : ListM m α → m (Option (α × ListM m α))
   | nil => pure none
   | cons l => do
     let (x, xs) ← l
@@ -94,27 +94,27 @@ unsafe def uncons {α : Type u} : ListM m α → m (Option (α × ListM m α))
 #align tactic.mllist.uncons ListM.uncons
 
 /-- Compute, inside the monad, whether a `ListM` is empty. -/
-unsafe def isEmpty {α : Type u} (xs : ListM m α) : m (ULift Bool) :=
+unsafe def isEmpty (xs : ListM m α) : m (ULift Bool) :=
   (ULift.up ∘ Option.isSome) <$> uncons xs
 #align tactic.mllist.empty ListM.isEmpty
 
 /-- Convert a `List` to a `ListM`. -/
-unsafe def ofList {α : Type u} : List α → ListM m α
+unsafe def ofList : List α → ListM m α
   | [] => nil
   | h :: t => cons (pure (h, ofList t))
 #align tactic.mllist.of_list ListM.ofList
 
 /-- The empty `ListM`. -/
-unsafe def empty {α : Type u} : ListM m α := ofList []
+unsafe def empty : ListM m α := ofList []
 
 /-- Convert a `List` of values inside the monad into a `ListM`. -/
-unsafe def ofListM {α : Type u} : List (m α) → ListM m α
+unsafe def ofListM : List (m α) → ListM m α
   | [] => nil
   | h :: t => cons ((fun x => (x, ofListM t)) <$> some <$> h)
 #align tactic.mllist.m_of_list ListM.ofListM
 
 /-- Extract a list inside the monad from a `ListM`. -/
-unsafe def force {α} (L : ListM m α) : m (List α) := do
+unsafe def force (L : ListM m α) : m (List α) := do
   match ← uncons L with
   | none => pure []
   | some (x, xs) => (x :: ·) <$> force xs
@@ -142,7 +142,7 @@ unsafe def folds (f : β → α → β) (init : β) (L : ListM m α) : ListM m 
   L.foldsM (fun b a => pure (f b a)) init
 
 /-- Take the first `n` elements, as a list inside the monad. -/
-unsafe def takeAsList {α} : ListM m α → Nat → m (List α)
+unsafe def takeAsList : ListM m α → Nat → m (List α)
   | _,  0 => pure []
   | xs, n+1 => do
     match ← uncons xs with
@@ -167,33 +167,33 @@ unsafe def drop : ListM m α → Nat → ListM m α
     | none => return (none, empty)
 
 /-- Apply a function which returns values in the monad to every element of a `ListM`. -/
-unsafe def mapM {α β : Type u} (f : α → m β) (L : ListM m α) : ListM m β :=
+unsafe def mapM (f : α → m β) (L : ListM m α) : ListM m β :=
   cons do match ← uncons L with
   | some (x, xs) => return (← f x, mapM f xs)
   | none => return (none, empty)
 #align tactic.mllist.mmap ListM.mapM
 
 /-- Apply a function to every element of a `ListM`. -/
-unsafe def map {α β : Type u} (f : α → β) (L : ListM m α) : ListM m β :=
+unsafe def map (f : α → β) (L : ListM m α) : ListM m β :=
   L.mapM fun a => pure (f a)
 #align tactic.mllist.map ListM.map
 
 /-- Filter a `ListM` using a monadic function. -/
-unsafe def filterM {α : Type u} (p : α → m (ULift Bool)) (L : ListM m α) : ListM m α :=
+unsafe def filterM (p : α → m (ULift Bool)) (L : ListM m α) : ListM m α :=
   cons do match ← uncons L with
   | some (x, xs) => return (if (← p x).down then some x else none, filterM p xs)
   | none => return (none, empty)
 #align tactic.mllist.mfilter ListM.filterM
 
 /-- Filter a `ListM`. -/
-unsafe def filter {α : Type u} (p : α → Bool) (L : ListM m α) : ListM m α :=
+unsafe def filter (p : α → Bool) (L : ListM m α) : ListM m α :=
   L.filterM fun a => pure <| .up (p a)
 #align tactic.mllist.filter ListM.filter
 
 /-- Filter and transform a `ListM` using a function that returns values inside the monad. -/
 -- Note that the type signature has changed since Lean 3, when we allowed `f` to fail.
 -- Use `try?` from `Mathlib.Control.Basic` to lift a possibly failing function to `Option`.
-unsafe def filterMapM {α β : Type u} (f : α → m (Option β)) (L : ListM m α) : ListM m β :=
+unsafe def filterMapM (f : α → m (Option β)) (L : ListM m α) : ListM m β :=
   cons do match ← uncons L with
   | none => return (none, empty)
   | some (x, xs) => match ← f x with
@@ -202,7 +202,7 @@ unsafe def filterMapM {α β : Type u} (f : α → m (Option β)) (L : ListM m 
 #align tactic.mllist.mfilter_map ListM.filterMapM
 
 /-- Filter and transform a `ListM` using an `Option` valued function. -/
-unsafe def filterMap {α β : Type u} (f : α → Option β) : ListM m α → ListM m β :=
+unsafe def filterMap (f : α → Option β) : ListM m α → ListM m β :=
   filterMapM fun a => do pure (f a)
 #align tactic.mllist.filter_map ListM.filterMap
 
@@ -217,14 +217,14 @@ unsafe def takeWhile (f : α → Bool) : ListM m α → ListM m α :=
   takeWhileM fun a => pure (.up (f a))
 
 /-- Concatenate two monadic lazy lists. -/
-unsafe def append {α : Type u} (L M : ListM m α) : ListM m α :=
+unsafe def append (L M : ListM m α) : ListM m α :=
   cons do match ← uncons L with
   | none => return (none, M)
   | some (x, xs) => return (some x, append xs M)
 #align tactic.mllist.append ListM.append
 
 /-- Join a monadic lazy list of monadic lazy lists into a single monadic lazy list. -/
-unsafe def join {α : Type u} (L : ListM m (ListM m α)) : ListM m α :=
+unsafe def join (L : ListM m (ListM m α)) : ListM m α :=
   cons do match ← uncons L with
   | none => return (none, empty)
   | some (x, xs) => match ← uncons x with
@@ -238,14 +238,14 @@ unsafe def squash (t : m (ListM m α)) : ListM m α :=
 #align tactic.mllist.squash ListM.squash
 
 /-- Enumerate the elements of a monadic lazy list, starting at a specified offset. -/
-unsafe def enum_from {α : Type u} (n : Nat) (L : ListM m α) : ListM m (Nat × α) :=
+unsafe def enum_from (n : Nat) (L : ListM m α) : ListM m (Nat × α) :=
   cons do match ← uncons L with
   | none => return (none, empty)
   | some (x, xs) => return (some (n, x), xs.enum_from (n+1))
 #align tactic.mllist.enum_from ListM.enum_from
 
 /-- Enumerate the elements of a monadic lazy list. -/
-unsafe def enum {α : Type u} : ListM m α → ListM m (Nat × α) :=
+unsafe def enum : ListM m α → ListM m (Nat × α) :=
   enum_from 0
 #align tactic.mllist.enum ListM.enum
 
@@ -255,7 +255,7 @@ unsafe def range {m : Type → Type} [Alternative m] : ListM m Nat :=
 #align tactic.mllist.range ListM.range
 
 /-- Add one element to the end of a monadic lazy list. -/
-unsafe def concat {α : Type u} : ListM m α → α → ListM m α
+unsafe def concat : ListM m α → α → ListM m α
   | L, a => (ListM.ofList [L, ListM.ofList [a]]).join
 #align tactic.mllist.concat ListM.concat
 
@@ -267,7 +267,7 @@ unsafe def zip (L : ListM m α) (M : ListM m β) : ListM m (α × β) :=
 
 /-- Apply a function returning a monadic lazy list to each element of a monadic lazy list,
 joining the results. -/
-unsafe def bind {α β : Type u} (L : ListM m α) (f : α → ListM m β) : ListM m β :=
+unsafe def bind (L : ListM m α) (f : α → ListM m β) : ListM m β :=
   cons do match ← uncons L with
   | none => return (none, empty)
   | some (x, xs) => match ← uncons (f x) with
@@ -287,13 +287,13 @@ unsafe def liftM [Monad n] [MonadLift m n] (L : ListM m α) : ListM n α :=
     | some (a, L') => pure <| cons do pure (a, L'.liftM)
 
 /-- Given a lazy list in a state monad, run it on some initial state, recording the states. -/
-unsafe def runState {σ α : Type u} (L : ListM (StateT.{u} σ m) α) (s : σ) : ListM m (α × σ) :=
+unsafe def runState (L : ListM (StateT.{u} σ m) α) (s : σ) : ListM m (α × σ) :=
   squash do match ← StateT.run (uncons L) s with
     | (none, _) => pure empty
     | (some (a, L'), s') => pure <| cons do pure (some (a, s'), L'.runState s')
 
 /-- Given a lazy list in a state monad, run it on some initial state. -/
-unsafe def runState' {σ α : Type u} (L : ListM (StateT.{u} σ m) α) (s : σ) : ListM m α :=
+unsafe def runState' (L : ListM (StateT.{u} σ m) α) (s : σ) : ListM m α :=
   L.runState s |>.map (·.1)
 
 /-- Return the head of a monadic lazy list if it exists, as an `Option` in the monad. -/

Mathlib/Data/ListM/BestFirst.lean

@@ -0,0 +1,93 @@
+/-
+Copyright (c) 2023 Scott Morrison. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Scott Morrison
+-/
+import Mathlib.Data.ListM.Basic
+
+/-!
+# 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 `α → ListM 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
+* `removeDuplicates` maintains an `RBSet` of previously visited nodes;
+  otherwise if the graph is not a tree nodes may be visited multiple times.
+-/
+
+
+variable {α : Type u} [Monad m] [Alternative m] [Ord α]
+
+open Std ListM
+
+/--
+Auxiliary function for `bestFirstSearch`, that updates the internal state,
+consisting of a priority queue of triples `α × Nat × ListM m α`.
+We remove the next element from the list contained in the best triple
+(discarding the triple if there is no next element),
+enqueue it and return it.
+-/
+-- The return type has `× List α` rather than just `× Option α` so `bestFirstSearch` can use `fixl`.
+unsafe def bestFirstSearchAux
+    (f : Nat → α → ListM m α) (maxQueued : Option Nat := none) :
+    RBMap α (Nat × ListM m α) compare → m (RBMap α (Nat × ListM m α) compare × List α) :=
+fun s => do
+  match s.min with
+  | none => failure
+  | some (a, (n, L)) =>
+    match ← uncons L with
+    | none => pure (s.erase a, [])
+    | some (b, L') => do
+      let s' := s.insert a (n, L') |>.insert b (n + 1, f (n+1) b)
+      let s' := match maxQueued with
+      | some q => if s'.size > q then
+          match s'.max with | some x => s'.erase x.1 | none => unreachable!
+        else
+          s'
+      | none => s'
+      pure (s', [b])
+
+/--
+A lazy list recording the best first search of a graph generated by a function
+`f : α → ListM 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.
+-/
+unsafe def bestFirstSearch (f : α → ListM m α) (a : α)
+    (maxDepth : Option Nat := none) (maxQueued : Option Nat := none) (removeDuplicates := true) :
+    ListM m α :=
+let f := match maxDepth with
+| none => fun _ a => f a
+| some d => fun n a => if d < n then empty else f a
+if removeDuplicates then
+  let f' : Nat → α → ListM (StateT.{u} (RBSet α compare) m) α := fun n a =>
+    (f n a).liftM >>= fun b => do
+      let s ← get
+      if s.contains b then failure
+      set <| s.insert b
+      pure b
+  cons (do pure (some a, fixl (bestFirstSearchAux f' maxQueued) (RBMap.single a (0, f' 0 a))))
+    |>.runState' (RBSet.empty.insert a)
+else
+  cons do pure (some a, fixl (bestFirstSearchAux f maxQueued) (RBMap.single a (0, f 0 a)))

Mathlib/Data/ListM/DepthFirst.lean

@@ -0,0 +1,83 @@
+/-
+Copyright (c) 2023 Scott Morrison. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Scott Morrison
+-/
+import Mathlib.Data.ListM.Basic
+import Mathlib.Control.Traversable.Basic
+
+/-!
+# Depth first search
+
+We perform depth first search of a tree or graph,
+where the neighbours of a vertex are provided either by list `α → List α`
+or a lazy list `α → ListM MetaM α`.
+
+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.
+-/
+
+section
+variable [Monad m] [Alternative m]
+
+/-- A generalisation of `depthFirst`, which allows the generation function to know the current
+depth, and to count the depth starting from a specified value. -/
+partial def depthFirst' (f : Nat → α → m α) (n : Nat) (a : α) : m α :=
+pure a <|> joinM ((f n a) <&> (depthFirst' f (n+1)))
+
+/--
+Depth first search of a graph generated by a function
+`f : α → m α`.
+
+Here `m` must be an `Alternative` `Monad`,
+and perhaps the only sensible values are `List` and `ListM MetaM`.
+
+The option `maxDepth` limits the search depth.
+
+Note that if the graph is not a tree then elements will be visited multiple times.
+(See `depthFirstRemovingDuplicates`)
+-/
+def depthFirst (f : α → m α) (a : α) (maxDepth : Option Nat := none) : m α :=
+match maxDepth with
+| some d => depthFirst' (fun n a => if n ≤ d then f a else failure) 0 a
+| none => depthFirst' (fun _ a => f a) 0 a
+
+end
+
+variable [Monad m]
+
+open Lean in
+/--
+Variant of `depthFirst`,
+using an internal `HashSet` to record and avoid already visited nodes.
+
+This version describes the graph using `α → ListM m α`,
+and returns the monadic lazy list of nodes visited in order.
+
+This is potentially very expensive.
+If you want to do efficient enumerations from a generation function,
+avoiding duplication up to equality or isomorphism,
+use Brendan McKay's method of "generation by canonical construction path".
+-/
+-- TODO can you make this work in `List` and `ListM m` simultaneously, by being tricky with monads?
+unsafe def depthFirstRemovingDuplicates {α : Type u} [BEq α] [Hashable α]
+    (f : α → ListM m α) (a : α) (maxDepth : Option Nat := none) : ListM m α :=
+let f' : α → ListM (StateT.{u} (HashSet α) m) α := fun a =>
+  (f a).liftM >>= fun b => do
+    let s ← get
+    if s.contains b then failure
+    set <| s.insert b
+    pure b
+(depthFirst f' a maxDepth).runState' (HashSet.empty.insert a)
+
+/--
+Variant of `depthFirst`,
+using an internal `HashSet` to record and avoid already visited nodes.
+
+This version describes the graph using `α → List α`, and returns the list of nodes visited in order.
+-/
+def depthFirstRemovingDuplicates' [BEq α] [Hashable α]
+    (f : α → List α) (a : α) (maxDepth : Option Nat := none) : List α :=
+unsafe depthFirstRemovingDuplicates
+  (fun a => (.ofList (f a) : ListM Option α)) a maxDepth |>.force |>.get!

Mathlib/Data/ListM/Heartbeats.lean

@@ -3,7 +3,7 @@ Copyright (c) 2023 Scott Morrison. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
-import Mathlib.Data.ListM
+import Mathlib.Data.ListM.Basic
 import Mathlib.Lean.CoreM
 
 /-!

test/ListM.lean

@@ -3,7 +3,7 @@ Copyright (c) 2019 Scott Morrison. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
-import Mathlib.Data.ListM
+import Mathlib.Data.ListM.Basic
 import Mathlib.Control.Basic
 
 @[reducible] def S (α : Type) := StateT (List Nat) Option α