feat: sugar for SatisfiesM

Batteries/Classes/SatisfiesM.lean

@@ -1,8 +1,10 @@
 /-
 Copyright (c) 2022 Mario Carneiro. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Mario Carneiro
+Authors: Mario Carneiro, Kim Morrison
 -/
+import Batteries.Lean.EStateM
+import Batteries.Lean.Except
 
 /-!
 ## SatisfiesM
@@ -12,6 +14,13 @@ and enables Hoare-like reasoning over monadic expressions. For example, given a
 function `f : α → m β`, to say that the return value of `f` satisfies `Q` whenever
 the input satisfies `P`, we write `∀ a, P a → SatisfiesM Q (f a)`.
 
+For any monad equipped with `Satisfying m`
+one can lift `SatisfiesM` to a monadic value in `Subtype`,
+using `satisfying x h : m {a // p a}`, where `x : m α` and `h : SatisfiesM p x`.
+This includes `Option`, `ReaderT`, `StateT`, and `ExceptT`, and the Lean monad stack.
+(Although it is not entirely clear one should treat the Lean monad stack as lawful,
+even though Lean accepts this.)
+
 ## Notes
 
 `SatisfiesM` is not yet a satisfactory solution for verifying the behaviour of large scale monadic
@@ -23,7 +32,7 @@ presumably requiring more syntactic support (and smarter `do` blocks) from Lean.
 Or it may be that such a solution will look different!
 This is an open research program, and for now one should not be overly ambitious using `SatisfiesM`.
 
-In particular lemmas about pure operations on data structures in `batteries` except for `HashMap`
+In particular lemmas about pure operations on data structures in `Batteries` except for `HashMap`
 should avoid `SatisfiesM` for now, so that it is easy to migrate to other approaches in future.
 -/
 
@@ -158,6 +167,23 @@ end SatisfiesM
   ⟨by revert x; intro | .ok _, ⟨.ok ⟨_, h⟩, rfl⟩, _, rfl => exact h,
    fun h => match x with | .ok a => ⟨.ok ⟨a, h _ rfl⟩, rfl⟩ | .error e => ⟨.error e, rfl⟩⟩
 
+@[simp] theorem SatisfiesM_EStateM_eq :
+    SatisfiesM (m := EStateM ε σ) p x ↔ ∀ s a s', x.run s = .ok a s' → p a := by
+  constructor
+  · rintro ⟨x, rfl⟩ s a s' h
+    match w : x.run s with
+    | .ok a s' => simp at h; exact h.1
+    | .error e s' => simp [w] at h
+  · intro w
+    refine ⟨?_, ?_⟩
+    · intro s
+      match q : x.run s with
+      | .ok a s' => exact .ok ⟨a, w s a s' q⟩ s'
+      | .error e s' => exact .error e s'
+    · ext s
+      rw [EStateM.run_map, EStateM.run]
+      split <;> simp_all
+
 @[simp] theorem SatisfiesM_ReaderT_eq [Monad m] :
     SatisfiesM (m := ReaderT ρ m) p x ↔ ∀ s, SatisfiesM p (x s) :=
   (exists_congr fun a => by exact ⟨fun eq _ => eq ▸ rfl, funext⟩).trans Classical.skolem.symm
@@ -174,9 +200,93 @@ theorem SatisfiesM_StateRefT_eq [Monad m] :
     show _ >>= _ = _; simp [← h]
 
 @[simp] theorem SatisfiesM_ExceptT_eq [Monad m] [LawfulMonad m] :
-    SatisfiesM (m := ExceptT ρ m) (α := α) p x ↔ SatisfiesM (m := m) (∀ a, · = .ok a → p a) x := by
+    SatisfiesM (m := ExceptT ρ m) (α := α) p x ↔ SatisfiesM (m := m) (∀ a, · = .ok a → p a) x.run := by
+  change _ ↔ SatisfiesM (m := m) (∀ a, · = .ok a → p a) x
   refine ⟨fun ⟨f, eq⟩ => eq ▸ ?_, fun ⟨f, eq⟩ => eq ▸ ?_⟩
   · exists (fun | .ok ⟨a, h⟩ => ⟨.ok a, fun | _, rfl => h⟩ | .error e => ⟨.error e, nofun⟩) <$> f
     show _ = _ >>= _; rw [← comp_map, map_eq_pure_bind]; congr; funext a; cases a <;> rfl
   · exists ((fun | ⟨.ok a, h⟩ => .ok ⟨a, h _ rfl⟩ | ⟨.error e, _⟩ => .error e) <$> f : m _)
     show _ >>= _ = _; simp [← comp_map, ← bind_pure_comp]; congr; funext ⟨a, h⟩; cases a <;> rfl
+
+/--
+If a monad has `Satisfying m`, then we can lift a `h : SatisfiesM (m := m) p x` predicate
+to monadic value `satisfying x p : m { x // p x }`.
+
+Reader, state, and exception monads have `Satisfying` instances if the base monad does.
+-/
+class Satisfying (m : Type u → Type v) [Functor m] where
+  /-- Lift a `SatisfiesM` predicate to a monadic value. -/
+  satisfying {p : α → Prop} (x : m α) (h : SatisfiesM (m := m) p x) :
+    { y : m {a // p a} // Subtype.val <$> y = x }
+
+export Satisfying (satisfying)
+
+instance : Satisfying Id where
+  satisfying x h := ⟨⟨x, by obtain ⟨⟨_, h⟩, rfl⟩ := h; exact h⟩, rfl⟩
+
+instance : Satisfying Option where
+  satisfying x? h :=
+  have h' := SatisfiesM_Option_eq.mp h
+  { val := match x? with
+    | none => none
+    | some x => some ⟨x, h' x rfl⟩
+    property := by cases x? <;> simp }
+
+instance : Satisfying (Except ε) where
+  satisfying x? h :=
+  have h' := SatisfiesM_Except_eq.mp h
+  { val := match x? with
+    | .ok x => .ok ⟨x, h' x rfl⟩
+    | .error e => .error e
+    property := by cases x? <;> simp }
+
+-- This will be redundant after nightly-2024-11-08.
+attribute [ext] ReaderT.ext
+
+instance [Monad m] [Satisfying m] : Satisfying (ReaderT ρ m) where
+  satisfying x h :=
+  have h' := SatisfiesM_ReaderT_eq.mp h
+  { val := fun r => (satisfying (x.run r) (h' r)).val
+    property := by
+      ext r
+      let t := satisfying (x.run r) (h' r)
+      rw [ReaderT.run_map, ← t.2]
+      rfl }
+
+instance [Monad m] [Satisfying m] : Satisfying (StateRefT' ω σ m) :=
+  inferInstanceAs <| Satisfying (ReaderT _ _)
+
+-- This will be redundant after nightly-2024-11-08.
+attribute [ext] StateT.ext
+
+instance [Monad m] [LawfulMonad m] [Satisfying m] : Satisfying (StateT ρ m) where
+  satisfying x h :=
+  have h' := SatisfiesM_StateT_eq.mp h
+  { val := fun r => (fun ⟨⟨a, r'⟩, h⟩ => ⟨⟨a, h⟩, r'⟩) <$> (satisfying (x.run r) (h' r)).1
+    property := by
+      ext r
+      rw [← (satisfying (x.run r) (h' r)).2, StateT.run_map]
+      simp [StateT.run] }
+
+instance [Monad m] [LawfulMonad m] [Satisfying m] : Satisfying (ExceptT ε m) where
+  satisfying x h :=
+  have h' := SatisfiesM_ExceptT_eq.mp h
+  have x' := satisfying x.run h'
+  { val := ExceptT.mk ((fun ⟨y, w⟩ => y.pmap fun a h => ⟨a, w _ h⟩) <$> x'.1)
+    property := by
+      ext
+      simp [← x'.2] }
+
+instance : Satisfying (EStateM ε σ) where
+  satisfying x h :=
+  have h' := SatisfiesM_EStateM_eq.mp h
+  { val := fun s => match w : x.run s with
+    | .ok a s' => .ok ⟨a, h' s a s' w⟩ s'
+    | .error e s' => .error e s'
+    property := by
+      ext s
+      rw [EStateM.run_map, EStateM.run]
+      split <;> simp_all }
+
+instance : Satisfying (EIO ε) := inferInstanceAs <| Satisfying (EStateM _ _)
+instance : Satisfying IO := inferInstanceAs <| Satisfying (EIO _)

Batteries/Lean/EStateM.lean

@@ -0,0 +1,40 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+
+namespace EStateM
+
+namespace Result
+
+/-- Map a function over an `EStateM.Result`, preserving states and errors. -/
+def map {ε σ α β : Type u} (f : α → β) (x : Result ε σ α) : Result ε σ β :=
+  match x with
+  | .ok a s' => .ok (f a) s'
+  | .error e s' => .error e s'
+
+@[simp] theorem map_ok {ε σ α β : Type u} (f : α → β) (a : α) (s : σ) :
+    (Result.ok a s : Result ε σ α).map f = .ok (f a) s := rfl
+
+@[simp] theorem map_error {ε σ α β : Type u} (f : α → β) (e : ε) (s : σ) :
+    (Result.error e s : Result ε σ α).map f = .error e s := rfl
+
+@[simp] theorem map_eq_ok {ε σ α β : Type u} (f : α → β) (x : Result ε σ α) (b : β) (s : σ) :
+    x.map f = .ok b s ↔ ∃ a, x = .ok a s ∧ b = f a := by
+  cases x <;> simp [and_assoc, and_comm, eq_comm]
+
+@[simp] theorem map_eq_error {ε σ α β : Type u} (f : α → β) (x : Result ε σ α) (e : ε) (s : σ) :
+    x.map f = .error e s ↔ x = .error e s := by
+  cases x <;> simp [eq_comm]
+
+end Result
+
+@[simp] theorem run_map (f : α → β) (x : EStateM ε σ α) :
+    (f <$> x).run s = (x.run s).map f := rfl
+
+@[ext] theorem ext {ε σ α : Type u} (x y : EStateM ε σ α) (h : ∀ s, x.run s = y.run s) : x = y := by
+  funext s
+  exact h s
+
+end EStateM

Batteries/Lean/Except.lean

@@ -7,7 +7,56 @@ import Lean.Util.Trace
 
 open Lean
 
+namespace Except
+
 /-- Visualize an `Except` using a checkmark or a cross. -/
-def Except.emoji : Except ε α → String
+def emoji : Except ε α → String
   | .error _ => crossEmoji
   | .ok _ => checkEmoji
+
+@[simp] theorem map_error {ε : Type u} (f : α → β) (e : ε) :
+    f <$> (.error e : Except ε α) = .error e := rfl
+
+@[simp] theorem map_ok {ε : Type u} (f : α → β) (x : α) :
+    f <$> (.ok x : Except ε α) = .ok (f x) := rfl
+
+/-- Map a function over an `Except` value, using a proof that the value is `.ok`. -/
+def pmap {ε : Type u} {α β : Type v} (x : Except ε α) (f : (a : α) → x = .ok a → β) : Except ε β :=
+  match x with
+  | .error e => .error e
+  | .ok a => .ok (f a rfl)
+
+@[simp] theorem pmap_error {ε : Type u} {α β : Type v} (e : ε)
+    (f : (a : α) → Except.error e = Except.ok a → β) :
+    Except.pmap (.error e) f = .error e := rfl
+
+@[simp] theorem pmap_ok {ε : Type u} {α β : Type v} (a : α)
+    (f : (a' : α) → (.ok a : Except ε α) = .ok a' → β) :
+    Except.pmap (.ok a) f = .ok (f a rfl) := rfl
+
+@[simp] theorem pmap_id {ε : Type u} {α : Type v} (x : Except ε α) :
+    x.pmap (fun a _ => a) = x := by cases x <;> simp
+
+@[simp] theorem map_pmap (g : β → γ) (x : Except ε α) (f : (a : α) → x = .ok a → β) :
+    g <$> x.pmap f = x.pmap fun a h => g (f a h) := by
+  cases x <;> simp
+
+end Except
+
+namespace ExceptT
+
+-- This will be redundant after nightly-2024-11-08.
+attribute [ext] ExceptT.ext
+
+@[simp] theorem run_mk {m : Type u → Type v} (x : m (Except ε α)) : (ExceptT.mk x).run = x := rfl
+@[simp] theorem mk_run (x : ExceptT ε m α) : ExceptT.mk (ExceptT.run x) = x := rfl
+
+@[simp]
+theorem map_mk [Monad m] [LawfulMonad m] (f : α → β) (x : m (Except ε α)) :
+    f <$> ExceptT.mk x = ExceptT.mk ((f <$> · ) <$> x) := by
+  simp only [Functor.map, Except.map, ExceptT.map, map_eq_pure_bind]
+  congr
+  funext a
+  split <;> simp
+
+end ExceptT

Batteries/Lean/SatisfiesM.lean

@@ -0,0 +1,28 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+import Batteries.Classes.SatisfiesM
+import Lean.Elab.Command
+
+/-!
+# Construct `Satisfying` instances for the Lean monad stack.
+-/
+
+open Lean Elab Term Tactic Command
+
+instance : Satisfying CoreM :=
+  inferInstanceAs <| Satisfying (ReaderT _ <| StateRefT' _ _ (EIO _))
+
+instance : Satisfying MetaM :=
+  inferInstanceAs <| Satisfying (ReaderT _ <| StateRefT' _ _ CoreM)
+
+instance : Satisfying TermElabM :=
+  inferInstanceAs <| Satisfying (ReaderT _ <| StateRefT' _ _ MetaM)
+
+instance : Satisfying TacticM :=
+  inferInstanceAs <| Satisfying (ReaderT _ $ StateRefT' _ _ TermElabM)
+
+instance : Satisfying CommandElabM :=
+  inferInstanceAs <| Satisfying (ReaderT _ $ StateRefT' _ _ (EIO _))

BatteriesTest/satisfying.lean

@@ -0,0 +1,21 @@
+import Batteries.Lean.SatisfiesM
+import Batteries.Data.Array.Monadic
+
+open Lean Meta Array Elab Term Tactic Command
+
+-- For now these live in the test file, as it's not really clear we want people relying on them!
+instance : LawfulMonad (EIO ε) := inferInstanceAs <| LawfulMonad (EStateM _ _)
+instance : LawfulMonad CoreM :=
+  inferInstanceAs <| LawfulMonad (ReaderT _ <| StateRefT' _ _ (EIO Exception))
+instance : LawfulMonad MetaM :=
+  inferInstanceAs <| LawfulMonad (ReaderT _ <| StateRefT' _ _ CoreM)
+instance : LawfulMonad TermElabM :=
+  inferInstanceAs <| LawfulMonad (ReaderT _ <| StateRefT' _ _ MetaM)
+instance : LawfulMonad TacticM :=
+  inferInstanceAs <| LawfulMonad (ReaderT _ $ StateRefT' _ _ $ TermElabM)
+instance : LawfulMonad CommandElabM :=
+  inferInstanceAs <| LawfulMonad (ReaderT _ $ StateRefT' _ _ $ EIO _)
+
+example (xs : Array Expr) : MetaM { ts : Array Expr // ts.size = xs.size } := do
+  let r ← satisfying (xs.mapM inferType) (size_mapM _ _)
+  return r