feat: run hint tactics in parallel

Mathlib.lean

@@ -1662,6 +1662,7 @@ import Mathlib.Data.MLList.BestFirst
 import Mathlib.Data.MLList.Dedup
 import Mathlib.Data.MLList.DepthFirst
 import Mathlib.Data.MLList.IO
+import Mathlib.Data.MLList.Meta
 import Mathlib.Data.MLList.Split
 import Mathlib.Data.Matrix.Auto
 import Mathlib.Data.Matrix.Basic

Mathlib/Data/MLList/Meta.lean

@@ -0,0 +1,224 @@
+/-
+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.MLList.IO
+import Mathlib.Data.List.Defs
+
+/-!
+# Parallelization in Lean's tactic monads.
+
+`MetaM.runGreedily` will run a list of `MetaM α` in parallel, returning
+* a cancellation hook and
+* an `MLList MetaM α` that returns the results (and the relevant `MetaM` state)
+  in the order that they finish.
+
+After calling the cancellation hook the behaviour of the monadic lazy list should not be relied on.
+It may continue returning values, or fail.
+Recommended usage is to take a prefix of the list
+(e.g. with `MLList.takeUpToFirst` followed by `MLList.force`, or `MLList.takeAsList`)
+and then call the cancellation hook to request cancellation of later unwanted tasks.
+
+(Similarly also for `CoreM`, `TermElabM`, and `TacticM`.)
+
+## Implementation notes:
+Calling `IO.cancel` on `t.map f` does not cancel `t`,
+so we have to be careful throughout this file
+to construct cancellation hooks connected to the underlying task,
+rather than the various maps of it that we construct to pass state.
+
+Thomas Murrills has a suggestion to significantly refactor this code,
+reducing duplication using `MonadControl`, but it will require a core change.
+See https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/.60StateRefT'.60.20.60MonadControl.60.20instance.3F
+-/
+
+set_option autoImplicit true
+
+namespace IO
+
+/--
+Given a list of values in `IO α`, executes them all in parallel as tasks, and returns
+* a cancellation hook and
+* a monadic lazy list which returns the values in the order they complete.
+
+Note that the cancellation hook merely requests cooperative cancellation:
+the tasks must call `IO.checkCanceled` themselves.
+
+After calling the cancellation hook the behaviour of the monadic lazy list should not be relied on.
+It may continue returning values, or fail.
+Recommended usage is to take a prefix of the list
+(e.g. with `MLList.takeUpToFirst` followed by `MLList.force`, or `MLList.takeAsList`)
+and then call the cancellation hook to request cancellation of later unwanted tasks.
+-/
+def runGreedily (tasks : List (IO α)) : IO (BaseIO Unit × MLList IO α) := do
+  let t ← show BaseIO _ from (tasks.map IO.asTask).traverse id
+  return (t.forM cancel, MLList.ofTaskList t |>.liftM |>.mapM fun
+  | .ok a => pure a
+  | .error e => throw (IO.userError s!"{e}"))
+
+/--
+Variant of `IO.runGreedily` without a cancellation hook.
+-/
+def runGreedily' (tasks : List (IO α)) : MLList IO α :=
+  .squash fun _ => (·.2) <$> runGreedily tasks
+
+end IO
+
+namespace Lean.Core.CoreM
+
+/--
+Given a monadic value in `CoreM`, creates a task that runs it in the current state,
+returning
+* a cancellation hook and
+* a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask (t : CoreM α) : CoreM (BaseIO Unit × Task (CoreM α)) := do
+  let task ← (t.toIO (← read) (← get)).asTask
+  return (IO.cancel task, task.map (prio := .max) fun
+  | .ok (a, s) => do set s; pure a
+  | .error e => throwError m!"Task failed:\n{e}")
+
+/--
+Given a monadic value in `CoreM`, creates a task that runs it in the current state,
+returning a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask' (t : CoreM α) : CoreM (Task (CoreM α)) := (·.2) <$> asTask t
+
+/--
+Given a list of monadic values in `CoreM`, runs them all as tasks,
+and returns
+* a cancellation hook and
+* the monadic lazy list which returns the values in the order they complete.
+
+See the doc-string for `IO.runGreedily` for details about the cancellation hook behaviour.
+-/
+def runGreedily (jobs : List (CoreM α)) : CoreM (BaseIO Unit × MLList CoreM α) := do
+  let (cancels, tasks) := (← jobs.mapM asTask).unzip
+  return (cancels.forM id, .squash fun _ => return MLList.ofTaskList tasks |>.liftM.mapM id)
+
+/--
+Variant of `CoreM.runGreedily` without a cancellation hook.
+-/
+def runGreedily' (jobs : List (CoreM α)) : MLList CoreM α :=
+  .squash fun _ => (·.2) <$> runGreedily jobs
+
+end Lean.Core.CoreM
+
+namespace Lean.Meta.MetaM
+
+/--
+Given a monadic value in `MetaM`, creates a task that runs it in the current state,
+returning
+* a cancellation hook and
+* a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask (t : MetaM α) : MetaM (BaseIO Unit × Task (MetaM α)) := do
+  let (cancel, task) ← (t.run (← read) (← get)).asTask
+  return (cancel, task.map (prio := .max)
+    fun c : CoreM (α × Meta.State) => do let (a, s) ← c; set s; pure a)
+
+/--
+Given a monadic value in `MetaM`, creates a task that runs it in the current state,
+returning a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask' (t : MetaM α) : MetaM (Task (MetaM α)) := (·.2) <$> asTask t
+
+/--
+Given a list of monadic values in `MetaM`, runs them all as tasks,
+and returns
+* a cancellation hook and
+* the monadic lazy list which returns the values in the order they complete.
+
+See the doc-string for `IO.runGreedily` for details about the cancellation hook behaviour.
+-/
+def runGreedily (jobs : List (MetaM α)) : MetaM (BaseIO Unit × MLList MetaM α) := do
+  let (cancels, tasks) := (← jobs.mapM asTask).unzip
+  return (cancels.forM id, .squash fun _ => return MLList.ofTaskList tasks |>.liftM.mapM id)
+
+/--
+Variant of `MetaM.runGreedily` without a cancellation hook.
+-/
+def runGreedily' (jobs : List (MetaM α)) : MLList MetaM α :=
+  .squash fun _ => (·.2) <$> runGreedily jobs
+
+
+end Lean.Meta.MetaM
+
+namespace Lean.Elab.Term.TermElabM
+
+/--
+Given a monadic value in `TermElabM`, creates a task that runs it in the current state,
+returning
+* a cancellation hook and
+* a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask (t : TermElabM α) : TermElabM (BaseIO Unit × Task (TermElabM α)) := do
+  let (cancel, task) ← (t.run (← read) (← get)).asTask
+  return (cancel, task.map (prio := .max)
+    fun c : MetaM (α × Term.State) => do let (a, s) ← c; set s; pure a)
+
+/--
+Given a monadic value in `TermElabM`, creates a task that runs it in the current state,
+returning a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask' (t : TermElabM α) : TermElabM (Task (TermElabM α)) := (·.2) <$> asTask t
+
+/--
+Given a list of monadic values in `TermElabM`, runs them all as tasks,
+and returns
+* a cancellation hook and
+* the monadic lazy list which returns the values in the order they complete.
+
+See the doc-string for `IO.runGreedily` for details about the cancellation hook behaviour.
+-/
+def runGreedily (jobs : List (TermElabM α)) : TermElabM (BaseIO Unit × MLList TermElabM α) := do
+  let (cancels, tasks) := (← jobs.mapM asTask).unzip
+  return (cancels.forM id, .squash fun _ => return MLList.ofTaskList tasks |>.liftM.mapM id)
+
+/--
+Variant of `TermElabM.runGreedily` without a cancellation hook.
+-/
+def runGreedily' (jobs : List (TermElabM α)) : MLList TermElabM α :=
+  .squash fun _ => (·.2) <$> runGreedily jobs
+
+end Lean.Elab.Term.TermElabM
+
+namespace Lean.Elab.Tactic.TacticM
+
+/--
+Given a monadic value in `TacticM`, creates a task that runs it in the current state,
+returning
+* a cancellation hook and
+* a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask (t : TacticM α) : TacticM (BaseIO Unit × Task (TacticM α)) := do
+  let (cancel, task) ← (t (← read) |>.run (← get)).asTask
+  return (cancel, task.map (prio := .max)
+    fun c : TermElabM (α × Tactic.State) => do let (a, s) ← c; set s; pure a)
+
+/--
+Given a monadic value in `TacticM`, creates a task that runs it in the current state,
+returning a monadic value with the cached result (and subsequent state as it was after running).
+-/
+def asTask' (t : TacticM α) : TacticM (Task (TacticM α)) := (·.2) <$> asTask t
+
+/--
+Given a list of monadic values in `TacticM`, runs them all as tasks,
+and returns
+* a cancellation hook and
+* the monadic lazy list which returns the values in the order they complete.
+
+See the doc-string for `IO.runGreedily` for details about the cancellation hook behaviour.
+-/
+def runGreedily (jobs : List (TacticM α)) : TacticM (BaseIO Unit × MLList TacticM α) := do
+  let (cancels, tasks) := (← jobs.mapM asTask).unzip
+  return (cancels.forM id, .squash fun _ => return MLList.ofTaskList tasks |>.liftM.mapM id)
+
+/--
+Variant of `TacticM.runGreedily` without a cancellation hook.
+-/
+def runGreedily' (jobs : List (TacticM α)) : MLList TacticM α :=
+  .squash fun _ => (·.2) <$> runGreedily jobs
+
+end Lean.Elab.Tactic.TacticM

Mathlib/Tactic/Hint.lean

@@ -5,9 +5,7 @@ Authors: Scott Morrison
 -/
 import Std.Tactic.TryThis
 import Std.Linter.UnreachableTactic
-import Mathlib.Data.Nondet.Basic
-import Mathlib.Tactic.FailIfNoProgress
-import Mathlib.Mathport.Rename
+import Mathlib.Data.MLList.Meta
 
 /-!
 # The `hint` tactic.
@@ -16,8 +14,6 @@ The `hint` tactic tries the kitchen sink:
 it runs every tactic registered via the `register_hint tac` command
 on the current goal, and reports which ones succeed.
 
-## Future work
-It would be nice to run the tactics in parallel.
 -/
 
 open Lean Elab Tactic
@@ -93,28 +89,40 @@ def withoutInfoTrees (t : TacticM Unit) : TacticM Unit := do
   modifyInfoState fun s => { s with trees }
 
 /--
-Run all tactics registered using `register_hint`.
-Print a "Try these:" suggestion for each of the successful tactics.
+Run all tactics registered using `register_hint`,
+returning a list of triples consisting of
+* a `Suggestion` (for "Try this"),
+* the `List MVarId` of remaining goals, and
+* the `Tactic.SavedState` after running the tactic.
 
-If one tactic succeeds and closes the goal, we don't look at subsequent tactics.
+If one tactic succeeds and closes the goal, we cancel subsequent tactics.
 -/
--- TODO We could run the tactics in parallel.
--- TODO With widget support, could we run the tactics in parallel
---      and do live updates of the widget as results come in?
-def hint (stx : Syntax) : TacticM Unit := do
-  let tacs := Nondet.ofList (← getHints)
-  let results := tacs.filterMapM fun t : TSyntax `tactic => do
+def hintCore : TacticM (Array (Suggestion × List MVarId × SavedState)) := do
+  let tacs ← getHints
+  let tasks := tacs.map fun t : TSyntax `tactic => do
     if let some msgs ← observing? (withMessageLog (withoutInfoTrees (evalTactic t))) then
-      return some (← getGoals, ← suggestion t msgs)
+      return some (← suggestion t msgs, ← getGoals, ← saveState)
     else
       return none
-  let results ← (results.toMLList.takeUpToFirst fun r => r.1.1.isEmpty).asArray
-  let results := results.qsort (·.1.1.length < ·.1.1.length)
-  addSuggestions stx (results.map (·.1.2))
-  match results.find? (·.1.1.isEmpty) with
-  | some r =>
+  let (cancel, results) ← TacticM.runGreedily tasks
+  let results := results.filterMap id
+  let results ← (results.takeUpToFirst fun (_, mvars, _) => mvars.isEmpty).asArray
+  cancel -- Cancel any remaining tasks, in case one closed the goal early.
+  return results.qsort fun (_, mvars₁, _) (_, mvars₂, _) => mvars₁.length < mvars₂.length
+
+/--
+Run all tactics registered using `register_hint`.
+Print a "Try these:" suggestion for each of the successful tactics.
+
+If one tactic succeeds and closes the goal, we cancel subsequent tactics.
+-/
+def hint (stx : Syntax) : TacticM Unit := do
+  let results ← hintCore
+  addSuggestions stx (results.map (·.1))
+  match results.find? fun (_, mvars, _) => mvars.isEmpty with
+  | some (_, _, s) =>
     -- We don't restore the entire state, as that would delete the suggestion messages.
-    setMCtx r.2.term.meta.meta.mctx
+    setMCtx s.term.meta.meta.mctx
   | none => admitGoal (← getMainGoal)
 
 /--

test/MLList.lean

@@ -6,6 +6,8 @@ Authors: Scott Morrison
 import Std.Data.MLList.Basic
 import Mathlib.Control.Basic
 import Mathlib.Tactic.RunCmd
+import Std.Tactic.GuardMsgs
+import Mathlib.Data.MLList.Meta
 
 @[reducible] def S (α : Type) := StateT (List Nat) Option α
 def append (x : Nat) : S Unit :=
@@ -67,3 +69,60 @@ run_cmd Lean.Elab.Command.liftTermElabM do
     guard (n = 5)
     pure n
   guard $ n = 5
+
+/-!
+Tests for `runGreedily`, which converts a `List (MetaM α)` into a `MLList MetaM α`,
+streaming results as they become available.
+-/
+
+def busy_wait (millis : Nat) : IO Unit := do
+  let start ← IO.monoMsNow
+  while !(← IO.checkCanceled) && (← IO.monoMsNow) < start + millis do pure ()
+  if (← IO.checkCanceled) then throw <| IO.userError "cancelled"
+
+open Lean.Meta.MetaM
+
+/--
+info: 0
+-/
+#guard_msgs in
+#eval show MetaM _ from do
+  -- We put an `IO.sleep 0` in the long calculation to prevent Lean from optimizing away the `do`.
+  let t : List (MetaM Nat) := [do busy_wait 1000; pure 1000, do busy_wait 0; pure 0]
+  let r := runGreedily' t
+  r.head
+
+/--
+info: 1
+-/
+#guard_msgs in
+#eval show MetaM _ from do
+  let t : List (MetaM Nat) := [do busy_wait 0; pure 1, do busy_wait 5; pure 0]
+  let r := runGreedily' t
+  r.head
+
+/--
+info: Without cancellation:
+Result: 0
+Results after waiting: [10, 0]
+With cancellation:
+Result: 0
+Results after waiting: [0]
+-/
+#guard_msgs in
+#eval show MetaM _ from do
+  let ref : IO.Ref (List Nat) ← IO.mkRef []
+  let store : Nat → IO Nat := fun n => do ref.modify fun l => n :: l; pure n
+  let t : List (MetaM Nat) := [do busy_wait 10; store 10, do busy_wait 0; store 0]
+  let r := runGreedily' t
+  IO.println "Without cancellation:"
+  IO.println s!"Result: {← r.head}"
+  IO.sleep 20
+  IO.println s!"Results after waiting: {← ref.get}"
+  IO.println "With cancellation:"
+  ref.set []
+  let (c, r) ← runGreedily t
+  IO.println s!"Result: {← r.head}"
+  _ ← c
+  IO.sleep 20
+  IO.println s!"Results after waiting: {← ref.get}"