feat: LawfulMonad and WPMonad instances for Option and OptionT

sgraf812 · sgraf812 · commit 1c79aacaeb24 · 2025-09-29T16:06:38.000+02:00
diff --git a/src/Init/Control/Lawful/Instances.lean b/src/Init/Control/Lawful/Instances.lean
@@ -9,6 +9,8 @@ prelude
 public import Init.Control.Lawful.Basic
 public import Init.Control.Except
 import all Init.Control.Except
+public import Init.Control.Option
+import all Init.Control.Option
 public import Init.Control.State
 import all Init.Control.State
 public import Init.Control.StateRef
@@ -110,6 +112,118 @@ instance : LawfulMonad (Except ε) := LawfulMonad.mk'
 instance : LawfulApplicative (Except ε) := inferInstance
 instance : LawfulFunctor (Except ε) := inferInstance
 
+/-! # OptionT -/
+
+namespace OptionT
+
+@[ext] theorem ext {x y : OptionT m α} (h : x.run = y.run) : x = y := by
+  simp [run] at h
+  assumption
+
+@[simp, grind =] theorem run_pure [Monad m] (x : α) : run (pure x : OptionT m α) = pure (some x) := by
+  simp [run, pure, OptionT.pure, OptionT.mk]
+
+@[simp, grind =] theorem run_lift  [Monad.{u, v} m] (x : m α) : run (OptionT.lift x : OptionT m α) = (return some (← x) : m (Option α)) := by
+  simp [run, OptionT.lift, OptionT.mk]
+
+@[simp, grind =] theorem run_throw [Monad m] : run (throw e : OptionT m β) = pure none := by
+  simp [run, throw, throwThe, MonadExceptOf.throw, OptionT.fail, OptionT.mk]
+
+@[simp, grind =] theorem run_bind_lift [Monad m] [LawfulMonad m] (x : m α) (f : α → OptionT m β) : run (OptionT.lift x >>= f : OptionT m β) = x >>= fun a => run (f a) := by
+  simp [OptionT.run, OptionT.lift, bind, OptionT.bind, OptionT.mk]
+
+@[simp, grind =] theorem bind_throw [Monad m] [LawfulMonad m] (f : α → OptionT m β) : (throw e >>= f) = throw e := by
+  simp [throw, throwThe, MonadExceptOf.throw, bind, OptionT.bind, OptionT.mk, OptionT.fail]
+
+@[simp, grind =] theorem run_bind (f : α → OptionT m β) [Monad m] :
+    (x >>= f).run = Option.elimM x.run (pure none) (fun x => (f x).run) := by
+  change x.run >>= _ = _
+  simp [Option.elimM]
+  exact bind_congr fun |some _ => rfl | none => rfl
+
+@[simp, grind =] theorem lift_pure [Monad m] [LawfulMonad m] {α : Type u} (a : α) : OptionT.lift (pure a : m α) = pure a := by
+  simp only [OptionT.lift, OptionT.mk, bind_pure_comp, map_pure, pure, OptionT.pure]
+
+@[simp, grind =] theorem run_map [Monad m] [LawfulMonad m] (f : α → β) (x : OptionT m α)
+    : (f <$> x).run = Option.map f <$> x.run := by
+  simp [Functor.map, Option.map, ←bind_pure_comp]
+  apply bind_congr
+  intro a; cases a <;> simp [OptionT.pure, OptionT.mk]
+
+protected theorem seq_eq {α β : Type u} [Monad m] (mf : OptionT m (α → β)) (x : OptionT m α) : mf <*> x = mf >>= fun f => f <$> x :=
+  rfl
+
+protected theorem bind_pure_comp [Monad m] (f : α → β) (x : OptionT m α) : x >>= pure ∘ f = f <$> x := by
+  intros; rfl
+
+protected theorem seqLeft_eq {α β : Type u} {m : Type u → Type v} [Monad m] [LawfulMonad m] (x : OptionT m α) (y : OptionT m β) : x <* y = const β <$> x <*> y := by
+  change (x >>= fun a => y >>= fun _ => pure a) = (const (α := α) β <$> x) >>= fun f => f <$> y
+  rw [← OptionT.bind_pure_comp]
+  apply ext
+  simp [Option.elimM, Option.elim]
+  apply bind_congr
+  intro
+  | none => simp
+  | some _ =>
+    simp [←bind_pure_comp]; apply bind_congr; intro b;
+    cases b <;> simp [const]
+
+protected theorem seqRight_eq [Monad m] [LawfulMonad m] (x : OptionT m α) (y : OptionT m β) : x *> y = const α id <$> x <*> y := by
+  change (x >>= fun _ => y) = (const α id <$> x) >>= fun f => f <$> y
+  rw [← OptionT.bind_pure_comp]
+  apply ext
+  simp [Option.elimM, Option.elim]
+  apply bind_congr
+  intro a; cases a <;> simp
+
+instance [Monad m] [LawfulMonad m] : LawfulMonad (OptionT m) where
+  id_map         := by intros; apply ext; simp
+  map_const      := by intros; rfl
+  seqLeft_eq     := OptionT.seqLeft_eq
+  seqRight_eq    := OptionT.seqRight_eq
+  pure_seq       := by intros; apply ext; simp [OptionT.seq_eq, Option.elimM, Option.elim]
+  bind_pure_comp := OptionT.bind_pure_comp
+  bind_map       := by intros; rfl
+  pure_bind      := by intros; apply ext; simp [Option.elimM, Option.elim]
+  bind_assoc     := by intros; apply ext; simp [Option.elimM, Option.elim]; apply bind_congr; intro a; cases a <;> simp
+
+@[simp] theorem run_seq [Monad m] [LawfulMonad m] (f : OptionT m (α → β)) (x : OptionT m α) :
+    (f <*> x).run = Option.elimM f.run (pure none) (fun f => Option.map f <$> x.run) := by
+  simp [seq_eq_bind, Option.elimM, Option.elim]
+
+@[simp] theorem run_seqLeft [Monad m] [LawfulMonad m] (x : OptionT m α) (y : OptionT m β) :
+    (x <* y).run = Option.elimM x.run (pure none)
+      (fun x => Option.map (Function.const β x) <$> y.run) := by
+  simp [seqLeft_eq, seq_eq_bind, Option.elimM, OptionT.run_bind]
+
+@[simp] theorem run_seqRight [Monad m] [LawfulMonad m] (x : OptionT m α) (y : OptionT m β) :
+    (x *> y).run = Option.elimM x.run (pure none) (Function.const α y.run) := by
+  simp only [seqRight_eq, run_seq, Option.elimM, run_map, Option.elim, bind_map_left]
+  refine bind_congr (fun | some _ => by simp | none => by simp)
+
+@[simp, grind =] theorem run_failure [Monad m] : (failure : OptionT m α).run = pure none := by rfl
+
+@[simp] theorem map_failure [Monad m] [LawfulMonad m] {α β : Type _} (f : α → β) :
+    f <$> (failure : OptionT m α) = (failure : OptionT m β) := by
+  simp [OptionT.mk, Functor.map, Alternative.failure, OptionT.fail, OptionT.bind]
+
+@[simp] theorem run_orElse [Monad m] (x : OptionT m α) (y : OptionT m α) :
+    (x <|> y).run = Option.elimM x.run y.run (fun x => pure (some x)) :=
+  bind_congr fun | some _ => by rfl | none => by rfl
+
+end OptionT
+
+/-! # Option -/
+
+instance : LawfulMonad Option := LawfulMonad.mk'
+  (id_map := fun x => by cases x <;> rfl)
+  (pure_bind := fun _ _ => by rfl)
+  (bind_assoc := fun a _ _ => by cases a <;> rfl)
+  (bind_pure_comp := bind_pure_comp)
+
+instance : LawfulApplicative Option := inferInstance
+instance : LawfulFunctor Option := inferInstance
+
 /-! # ReaderT -/
 
 namespace ReaderT
diff --git a/src/Init/Control/Lawful/MonadLift/Instances.lean b/src/Init/Control/Lawful/MonadLift/Instances.lean
@@ -64,10 +64,6 @@ namespace OptionT
 
 variable [Monad m] [LawfulMonad m]
 
-@[simp]
-theorem lift_pure {α : Type u} (a : α) : OptionT.lift (pure a : m α) = pure a := by
-  simp only [OptionT.lift, OptionT.mk, bind_pure_comp, map_pure, pure, OptionT.pure]
-
 @[simp]
 theorem lift_bind {α β : Type u} (ma : m α) (f : α → m β) :
     OptionT.lift (ma >>= f) = OptionT.lift ma >>= (fun a => OptionT.lift (f a)) := by
diff --git a/src/Init/Control/Option.lean b/src/Init/Control/Option.lean
@@ -39,13 +39,14 @@ variable {m : Type u → Type v} [Monad m] {α β : Type u}
 Converts an action that returns an `Option` into one that might fail, with `none` indicating
 failure.
 -/
+@[always_inline, inline, expose]
 protected def mk (x : m (Option α)) : OptionT m α :=
   x
 
 /--
 Sequences two potentially-failing actions. The second action is run only if the first succeeds.
 -/
-@[always_inline, inline]
+@[always_inline, inline, expose]
 protected def bind (x : OptionT m α) (f : α → OptionT m β) : OptionT m β := OptionT.mk do
   match (← x) with
   | some a => f a
@@ -54,7 +55,7 @@ protected def bind (x : OptionT m α) (f : α → OptionT m β) : OptionT m β :
 /--
 Succeeds with the provided value.
 -/
-@[always_inline, inline]
+@[always_inline, inline, expose]
 protected def pure (a : α) : OptionT m α := OptionT.mk do
   pure (some a)
 
diff --git a/src/Std/Do/PredTrans.lean b/src/Std/Do/PredTrans.lean
@@ -137,28 +137,37 @@ instance instLawfulMonad : LawfulMonad (PredTrans ps) :=
 
 -- The interpretation of `StateT σ (PredTrans ps) α` into `PredTrans (.arg σ ps) α`.
 -- Think: modifyGetM
-def pushArg {σ : Type u} (x : StateT σ (PredTrans ps) α) : PredTrans (.arg σ ps) α :=
-  { apply := fun Q s => (x s).apply (fun (a, s) => Q.1 a s, Q.2),
-    conjunctive := by
-      intro Q₁ Q₂
-      apply SPred.bientails.of_eq
-      ext s
-      dsimp only [SPred.and_cons, ExceptConds.and]
-      rw [← ((x s).conjunctive _ _).to_eq]
-  }
+def pushArg {σ : Type u} (x : StateT σ (PredTrans ps) α) : PredTrans (.arg σ ps) α where
+  apply := fun Q s => (x s).apply (fun (a, s) => Q.1 a s, Q.2)
+  conjunctive := by
+    intro Q₁ Q₂
+    apply SPred.bientails.of_eq
+    ext s
+    dsimp only [SPred.and_cons, ExceptConds.and]
+    rw [← ((x s).conjunctive _ _).to_eq]
 
 -- The interpretation of `ExceptT ε (PredTrans ps) α` into `PredTrans (.except ε ps) α`
-def pushExcept {ps : PostShape} {α ε} (x : ExceptT ε (PredTrans ps) α) : PredTrans (.except ε ps) α :=
-  { apply Q := x.apply (fun | .ok a => Q.1 a | .error e => Q.2.1 e, Q.2.2),
-    conjunctive := by
-      intro Q₁ Q₂
-      apply SPred.bientails.of_eq
-      dsimp
-      rw[← (x.conjunctive _ _).to_eq]
-      congr
-      ext x
-      cases x <;> simp
-  }
+def pushExcept {ps : PostShape} {α ε} (x : ExceptT ε (PredTrans ps) α) : PredTrans (.except ε ps) α where
+  apply Q := x.apply (fun | .ok a => Q.1 a | .error e => Q.2.1 e, Q.2.2)
+  conjunctive := by
+    intro Q₁ Q₂
+    apply SPred.bientails.of_eq
+    dsimp
+    rw[← (x.conjunctive _ _).to_eq]
+    congr
+    ext x
+    cases x <;> simp
+
+def pushOption {ps : PostShape} {α} (x : OptionT (PredTrans ps) α) : PredTrans (.except PUnit ps) α where
+  apply Q := x.apply (fun | .some a => Q.1 a | .none => Q.2.1 ⟨⟩, Q.2.2)
+  conjunctive := by
+    intro Q₁ Q₂
+    apply SPred.bientails.of_eq
+    dsimp
+    rw[← (x.conjunctive _ _).to_eq]
+    congr
+    ext x
+    cases x <;> simp
 
 @[simp]
 def pushArg_apply {ps} {α σ : Type u} {Q : PostCond α (.arg σ ps)} (f : σ → PredTrans ps (α × σ)) :
@@ -168,6 +177,10 @@ def pushArg_apply {ps} {α σ : Type u} {Q : PostCond α (.arg σ ps)} (f : σ 
 def pushExcept_apply {ps} {α ε : Type u} {Q : PostCond α (.except ε ps)} (x : PredTrans ps (Except ε α)) :
   (pushExcept x).apply Q = x.apply (fun | .ok a => Q.1 a | .error e => Q.2.1 e, Q.2.2) := rfl
 
+@[simp]
+def pushOption_apply {ps} {α : Type u} {Q : PostCond α (.except PUnit ps)} (x : PredTrans ps (Option α)) :
+  (pushOption x).apply Q = x.apply (fun | .some a => Q.1 a | .none => Q.2.1 ⟨⟩, Q.2.2) := rfl
+
 def dite_apply {ps} {Q : PostCond α ps} (c : Prop) [Decidable c] (t : c → PredTrans ps α) (e : ¬ c → PredTrans ps α) :
   (if h : c then t h else e h).apply Q = if h : c then (t h).apply Q else (e h).apply Q := by split <;> rfl
 
diff --git a/src/Std/Do/WP/Basic.lean b/src/Std/Do/WP/Basic.lean
@@ -67,6 +67,7 @@ protected meta def unexpandWP : Lean.PrettyPrinter.Unexpander
     | `(wp $e) => `(wp⟦$e⟧)
     | _ => throw ()
   | _ => throw ()
+
 instance Id.instWP : WP Id .pure where
   wp x := PredTrans.pure x.run
 
@@ -79,6 +80,9 @@ instance ReaderT.instWP [WP m ps] : WP (ReaderT ρ m) (.arg ρ ps) where
 instance ExceptT.instWP [WP m ps] : WP (ExceptT ε m) (.except ε ps) where
   wp x := PredTrans.pushExcept (wp x)
 
+instance OptionT.instWP [WP m ps] : WP (OptionT m) (.except PUnit ps) where
+  wp x := PredTrans.pushOption (wp x)
+
 instance EStateM.instWP : WP (EStateM ε σ) (.except ε (.arg σ .pure)) where
   wp x := -- Could define as PredTrans.mkExcept (PredTrans.modifyGetM (fun s => pure (EStateM.Result.toExceptState (x s))))
     { apply := fun Q s => match x s with
@@ -98,6 +102,8 @@ instance Reader.instWP : WP (ReaderM ρ) (.arg ρ .pure) :=
   inferInstanceAs (WP (ReaderT ρ Id) (.arg ρ .pure))
 instance Except.instWP : WP (Except ε) (.except ε .pure) :=
   inferInstanceAs (WP (ExceptT ε Id) (.except ε .pure))
+instance Option.instWP : WP Option (.except PUnit .pure) :=
+  inferInstanceAs (WP (OptionT Id) (.except PUnit .pure))
 
 /--
 Adequacy lemma for `Id.run`.
diff --git a/src/Std/Do/WP/Monad.lean b/src/Std/Do/WP/Monad.lean
@@ -72,6 +72,18 @@ instance ExceptT.instWPMonad [Monad m] [WPMonad m ps] : WPMonad (ExceptT ε m) (
     case error a => simp [wp_pure]
     case ok a => rfl
 
+instance OptionT.instWPMonad [Monad m] [WPMonad m ps] : WPMonad (OptionT m) (.except PUnit ps) where
+  wp_pure a := by ext; simp only [wp, pure, OptionT.pure, OptionT.mk, WPMonad.wp_pure,
+    PredTrans.pure, PredTrans.pushOption_apply]
+  wp_bind x f := by
+    ext Q
+    simp only [wp, bind, OptionT.bind, OptionT.mk, WPMonad.wp_bind, PredTrans.bind, PredTrans.pushOption_apply]
+    congr
+    ext b
+    cases b
+    case none => simp [wp_pure]
+    case some a => rfl
+
 instance EStateM.instWPMonad : WPMonad (EStateM ε σ) (.except ε (.arg σ .pure)) where
   wp_pure a := by simp only [wp, pure, EStateM.pure, PredTrans.pure]
   wp_bind x f := by
@@ -84,6 +96,10 @@ instance Except.instWPMonad : WPMonad (Except ε) (.except ε .pure) where
   wp_pure a := rfl
   wp_bind x f := by cases x <;> rfl
 
+instance Option.instWPMonad : WPMonad Option (.except PUnit .pure) where
+  wp_pure a := rfl
+  wp_bind x f := by cases x <;> rfl
+
 instance State.instWPMonad : WPMonad (StateM σ) (.arg σ .pure) :=
   inferInstanceAs (WPMonad (StateT σ Id) (.arg σ .pure))
 instance Reader.instWPMonad : WPMonad (ReaderM ρ) (.arg ρ .pure) :=
