Lean.Core.wrapAsync and MonadControl

Stream: lean4

[Eric Wieser]

I found myself wanting wrapAsync for MetaM. While it's easy to write this by hand by getting the state and context, I was wondering if using controlAt is appropriate here:

import Lean

open Lean Meta

@[inherit_doc Core.wrapAsync]
def Lean.Meta.wrapAsync (act : Unit → MetaM α) : MetaM (EIO Exception α) := do
  controlAt CoreM fun run =>
    Core.wrapAsync fun x => do
      run <| act x

[Eric Wieser]

Is this the same as:

@[inherit_doc Core.wrapAsync]
def Lean.Meta.wrapAsync2 (act : Unit → MetaM α) : MetaM (EIO Exception α) := do
  let metaCtx ← readThe Meta.Context
  let metaSt ← getThe Meta.State
  Core.wrapAsync fun _ =>
    act () |>.run' metaCtx metaSt

or is there some ST.Ref trickery here with multiple threads?

[Eric Wieser]

If I attempt to equate them with

example : @_root_.Lean.Meta.wrapAsync =  @Lean.Meta.wrapAsync2 := by
  unfold Lean.Meta.wrapAsync Lean.Meta.wrapAsync2 readThe getThe controlAt instMonadReaderOfReaderTOfMonad
    instMonadControlStateRefT' instMonadControlReaderT instMonadStateOfOfMonadLift
    instMonadControlTOfMonadControl liftWith instMonadControlTOfPure Function.comp inferInstanceAs instMonadLiftTOfMonadLift
    instMonadMetaM StateRefT'.instMonad inferInstanceAs ReaderT.instMonad ReaderT.read ReaderT.bind liftM ReaderT.instMonadLift
    StateRefT'.instMonadLift StateRefT'.lift  Core.instMonadCoreM inferInstanceAs StateRefT'.instMonad ReaderT.instMonad ReaderT.bind
    StateRefT'.instMonadStateOfOfMonadLiftTST StateRefT'.get inferInstanceAs instMonadEIO EStateM.instMonad inferInstanceAs EStateM.bind
  dsimp [Pure.pure, ReaderT.pure, EStateM.pure]
  ext α f ctx st mctx mst w
  rfl

I see that the state is

α : Type
f : Unit → MetaM α
ctx : Context
st : ST.Ref IO.RealWorld State
mctx : Core.Context
mst : ST.Ref IO.RealWorld Core.State
w : IO.RealWorld
⊢ (match Core.wrapAsync (fun x => f x ctx st) mctx mst w with
  | EStateM.Result.ok a s => EStateM.Result.ok a s
  | EStateM.Result.error e s => EStateM.Result.error e s) =
  match st.get mctx mst w with
  | EStateM.Result.ok a s => Core.wrapAsync (fun x => (f ()).run' ctx a) mctx mst s
  | EStateM.Result.error e s => EStateM.Result.error e s

[Sebastian Ullrich]

The MonadControl instance for StateRefT' is derived from ReaderT so I'm quite sure it's not suitable for multithreading