feat: add MonadControl lemmas for ReaderT, OptionT, StateT, and ExceptT (#11591)

This PR adds missing lemmas about how `ReaderT.run`, `OptionT.run`,
`StateT.run`, and `ExceptT.run` interact with `MonadControl` operations.

This also leaves some comments noting that the lemmas may look less
general than expected; but this is because the instances are also not
very general.

Author: eric-wieser

src/Init/Control/Lawful/Instances.lean

@@ -17,6 +17,9 @@ public section
 
 open Function
 
+@[simp, grind =] theorem monadMap_refl {m : Type _ → Type _} {α} (f : ∀ {α}, m α → m α) :
+    monadMap @f = @f α := rfl
+
 /-! # ExceptT -/
 
 namespace ExceptT
@@ -57,6 +60,9 @@ theorem run_bind [Monad m] (x : ExceptT ε m α) (f : α → ExceptT ε m β)
   apply bind_congr
   intro a; cases a <;> simp [Except.map]
 
+@[simp, grind =] theorem run_monadMap [MonadFunctorT n m] (f : {β : Type u} → n β → n β) (x : ExceptT ε m α)
+    : (monadMap @f x : ExceptT ε m α).run = monadMap @f (x.run) := rfl
+
 protected theorem seq_eq {α β ε : Type u} [Monad m] (mf : ExceptT ε m (α → β)) (x : ExceptT ε m α) : mf <*> x = mf >>= fun f => f <$> x :=
   rfl
 
@@ -99,6 +105,22 @@ instance [Monad m] [LawfulMonad m] : LawfulMonad (ExceptT ε m) where
   simp only [ExceptT.instMonad, ExceptT.map, ExceptT.mk, throw, throwThe, MonadExceptOf.throw,
     pure_bind]
 
+/-! Note that the `MonadControl` instance for `ExceptT` is not monad-generic. -/
+
+@[simp] theorem run_restoreM [Monad m] (x : stM m (ExceptT ε m) α) :
+    ExceptT.run (restoreM x) = pure x := rfl
+
+@[simp] theorem run_liftWith [Monad m] (f : ({β : Type u} → ExceptT ε m β → m (stM m (ExceptT ε m) β)) → m α) :
+    ExceptT.run (liftWith f) = Except.ok <$> (f fun x => x.run) :=
+  rfl
+
+@[simp] theorem run_controlAt [Monad m] [LawfulMonad m] (f : ({β : Type u} → ExceptT ε m β → m (stM m (ExceptT ε m) β)) → m (stM m (ExceptT ε m) α)) :
+    ExceptT.run (controlAt m f) = f fun x => x.run := by
+  simp [controlAt, run_bind, bind_map_left]
+
+@[simp] theorem run_control [Monad m] [LawfulMonad m] (f : ({β : Type u} → ExceptT ε m β → m (stM m (ExceptT ε m) β)) → m (stM m (ExceptT ε m) α)) :
+    ExceptT.run (control f) = f fun x => x.run := run_controlAt f
+
 end ExceptT
 
 /-! # Except -/
@@ -152,6 +174,9 @@ namespace OptionT
   apply bind_congr
   intro a; cases a <;> simp [OptionT.pure, OptionT.mk]
 
+@[simp, grind =] theorem run_monadMap [MonadFunctorT n m] (f : {β : Type u} → n β → n β) (x : OptionT m α)
+    : (monadMap @f x : OptionT m α).run = monadMap @f (x.run) := rfl
+
 protected theorem seq_eq {α β : Type u} [Monad m] (mf : OptionT m (α → β)) (x : OptionT m α) : mf <*> x = mf >>= fun f => f <$> x :=
   rfl
 
@@ -213,6 +238,24 @@ instance [Monad m] [LawfulMonad m] : LawfulMonad (OptionT m) where
     (x <|> y).run = Option.elimM x.run y.run (fun x => pure (some x)) :=
   bind_congr fun | some _ => by rfl | none => by rfl
 
+/-! Note that the `MonadControl` instance for `OptionT` is not monad-generic. -/
+
+@[simp] theorem run_restoreM [Monad m] (x : stM m (OptionT m) α) :
+    OptionT.run (restoreM x) = pure x := rfl
+
+@[simp] theorem run_liftWith [Monad m] [LawfulMonad m] (f : ({β : Type u} → OptionT m β → m (stM m (OptionT m) β)) → m α) :
+    OptionT.run (liftWith f) = Option.some <$> (f fun x => x.run) := by
+  dsimp [liftWith]
+  rw [← bind_pure_comp]
+  rfl
+
+@[simp] theorem run_controlAt [Monad m] [LawfulMonad m] (f : ({β : Type u} → OptionT m β → m (stM m (OptionT m) β)) → m (stM m (OptionT m) α)) :
+    OptionT.run (controlAt m f) = f fun x => x.run := by
+  simp [controlAt, Option.elimM, Option.elim]
+
+@[simp] theorem run_control [Monad m] [LawfulMonad m] (f : ({β : Type u} → OptionT m β → m (stM m (OptionT m) β)) → m (stM m (OptionT m) α)) :
+    OptionT.run (control f) = f fun x => x.run := run_controlAt f
+
 end OptionT
 
 /-! # Option -/
@@ -284,6 +327,22 @@ instance [Monad m] [LawfulMonad m] : LawfulMonad (ReaderT ρ m) where
   pure_bind      := by intros; apply ext; intros; simp
   bind_assoc     := by intros; apply ext; intros; simp
 
+/-! Note that the `MonadControl` instance for `ReaderT` is not monad-generic. -/
+
+@[simp] theorem run_restoreM [Monad m] (x : stM m (ReaderT ρ m) α) (ctx : ρ) :
+    ReaderT.run (restoreM x) ctx = pure x := rfl
+
+@[simp] theorem run_liftWith [Monad m] (f : ({β : Type u} → ReaderT ρ m β → m (stM m (ReaderT ρ m) β)) → m α) (ctx : ρ) :
+    ReaderT.run (liftWith f) ctx = (f fun x => x.run ctx) :=
+  rfl
+
+@[simp] theorem run_controlAt [Monad m] [LawfulMonad m] (f : ({β : Type u} → ReaderT ρ m β → m (stM m (ReaderT ρ m) β)) → m (stM m (ReaderT ρ m) α)) (ctx : ρ) :
+    ReaderT.run (controlAt m f) ctx = f fun x => x.run ctx := by
+  simp [controlAt]
+
+@[simp] theorem run_control [Monad m] [LawfulMonad m] (f : ({β : Type u} → ReaderT ρ m β → m (stM m (ReaderT ρ m) β)) → m (stM m (ReaderT ρ m) α)) (ctx : ρ) :
+    ReaderT.run (control f) ctx = f fun x => x.run ctx := run_controlAt f ctx
+
 end ReaderT
 
 /-! # StateRefT -/
@@ -307,11 +366,11 @@ namespace StateT
 @[simp, grind =] theorem run_pure [Monad m] (a : α) (s : σ) : (pure a : StateT σ m α).run s = pure (a, s) := rfl
 
 @[simp, grind =] theorem run_bind [Monad m] (x : StateT σ m α) (f : α → StateT σ m β) (s : σ)
-    : (x >>= f).run s = x.run s >>= λ p => (f p.1).run p.2 := by
-  simp [bind, StateT.bind, run]
+    : (x >>= f).run s = x.run s >>= λ p => (f p.1).run p.2 := rfl
 
 @[simp, grind =] theorem run_map {α β σ : Type u} [Monad m] [LawfulMonad m] (f : α → β) (x : StateT σ m α) (s : σ) : (f <$> x).run s = (fun (p : α × σ) => (f p.1, p.2)) <$> x.run s := by
-  simp [Functor.map, StateT.map, run, ←bind_pure_comp]
+  rw [← bind_pure_comp (m := m)]
+  rfl
 
 @[simp, grind =] theorem run_get [Monad m] (s : σ)    : (get : StateT σ m σ).run s = pure (s, s) := rfl
 
@@ -320,13 +379,13 @@ namespace StateT
 @[simp, grind =] theorem run_modify [Monad m] (f : σ → σ) (s : σ) : (modify f : StateT σ m PUnit).run s = pure (⟨⟩, f s) := rfl
 
 @[simp, grind =] theorem run_modifyGet [Monad m] (f : σ → α × σ) (s : σ) : (modifyGet f : StateT σ m α).run s = pure ((f s).1, (f s).2) := by
-  simp [modifyGet, MonadStateOf.modifyGet, StateT.modifyGet, run]
+  rfl
 
 @[simp, grind =] theorem run_lift {α σ : Type u} [Monad m] (x : m α) (s : σ) : (StateT.lift x : StateT σ m α).run s = x >>= fun a => pure (a, s) := rfl
 
 @[grind =]
 theorem run_bind_lift {α σ : Type u} [Monad m] [LawfulMonad m] (x : m α) (f : α → StateT σ m β) (s : σ) : (StateT.lift x >>= f).run s = x >>= fun a => (f a).run s := by
-  simp [StateT.lift, StateT.run, bind, StateT.bind]
+  simp
 
 @[simp, grind =] theorem run_monadLift {α σ : Type u} [Monad m] [MonadLiftT n m] (x : n α) (s : σ) : (monadLift x : StateT σ m α).run s = (monadLift x : m α) >>= fun a => pure (a, s) := rfl
 
@@ -366,6 +425,24 @@ instance [Monad m] [LawfulMonad m] : LawfulMonad (StateT σ m) where
   pure_bind      := by intros; apply ext; intros; simp
   bind_assoc     := by intros; apply ext; intros; simp
 
+/-! Note that the `MonadControl` instance for `StateT` is not monad-generic. -/
+
+@[simp] theorem run_restoreM [Monad m] [LawfulMonad m] (x : stM m (StateT σ m) α) (s : σ) :
+    StateT.run (restoreM x) s = pure x := by
+  simp [restoreM, MonadControl.restoreM]
+  rfl
+
+@[simp] theorem run_liftWith [Monad m] [LawfulMonad m] (f : ({β : Type u} → StateT σ m β → m (stM m (StateT σ m) β)) → m α) (s : σ) :
+    StateT.run (liftWith f) s = ((·, s) <$> f fun x => x.run s) := by
+  simp [liftWith, MonadControl.liftWith, Function.comp_def]
+
+@[simp] theorem run_controlAt [Monad m] [LawfulMonad m] (f : ({β : Type u} → StateT σ m β → m (stM m (StateT σ m) β)) → m (stM m (StateT σ m) α)) (s : σ) :
+    StateT.run (controlAt m f) s = f fun x => x.run s := by
+  simp [controlAt]
+
+@[simp] theorem run_control [Monad m] [LawfulMonad m] (f : ({β : Type u} → StateT σ m β → m (stM m (StateT σ m) β)) → m (stM m (StateT σ m) α)) (s : σ) :
+    StateT.run (control f) s = f fun x => x.run s := run_controlAt f s
+
 end StateT
 
 /-! # EStateM -/