|
| 1 | +/- |
| 2 | +Copyright (c) 2026 Quang Dao. All rights reserved. |
| 3 | +Released under Apache 2.0 license as described in the file LICENSE. |
| 4 | +Authors: Quang Dao |
| 5 | +-/ |
| 6 | + |
| 7 | +module |
| 8 | + |
| 9 | +public import Cslib.Init |
| 10 | +public import Mathlib.Data.PFunctor.Univariate.Basic |
| 11 | + |
| 12 | +/-! |
| 13 | +# Free Monad of a Polynomial Functor |
| 14 | +
|
| 15 | +We define the free monad on a **polynomial functor** (`PFunctor`), and prove some basic properties. |
| 16 | +
|
| 17 | +The free monad `PFunctor.FreeM P` extends the W-type construction with an extra `pure` |
| 18 | +constructor, yielding a monad that is free over the polynomial functor `P`. |
| 19 | +
|
| 20 | +This construction is ported from the [VCV-io](https://github.com/Verified-zkEVM/VCV-io) library. |
| 21 | +
|
| 22 | +## Main Definitions |
| 23 | +
|
| 24 | +- `PFunctor.FreeM`: The free monad on a polynomial functor. |
| 25 | +- `PFunctor.FreeM.lift`: Lift an object of the base polynomial functor into the free monad. |
| 26 | +- `PFunctor.FreeM.liftA`: Lift a position of the base polynomial functor into the free monad. |
| 27 | +- `PFunctor.FreeM.mapM`: Canonical mapping of `FreeM P` into any other monad. |
| 28 | +- `PFunctor.FreeM.inductionOn`: Induction principle for `FreeM P`. |
| 29 | +- `PFunctor.FreeM.construct`: Dependent eliminator for `FreeM P`. |
| 30 | +-/ |
| 31 | + |
| 32 | +@[expose] public section |
| 33 | + |
| 34 | +universe u v uA uB |
| 35 | + |
| 36 | +namespace PFunctor |
| 37 | + |
| 38 | +/-- The free monad on a polynomial functor. |
| 39 | +This extends the `W`-type construction with an extra `pure` constructor. -/ |
| 40 | +inductive FreeM (P : PFunctor.{uA, uB}) : Type v → Type (max uA uB v) |
| 41 | + /-- A leaf node wrapping a pure value. -/ |
| 42 | + | pure {α} (a : α) : FreeM P α |
| 43 | + /-- A node with shape `a : P.A` and subtrees given by the continuation |
| 44 | + `cont : P.B a → FreeM P α`. -/ |
| 45 | + | roll {α} (a : P.A) (cont : P.B a → FreeM P α) : FreeM P α |
| 46 | +deriving Inhabited |
| 47 | + |
| 48 | +namespace FreeM |
| 49 | + |
| 50 | +variable {P : PFunctor.{uA, uB}} {α β γ : Type v} |
| 51 | + |
| 52 | +/-- Lift an object of the base polynomial functor into the free monad. -/ |
| 53 | +@[always_inline, inline, reducible] |
| 54 | +def lift (x : P.Obj α) : FreeM P α := FreeM.roll x.1 (fun y ↦ FreeM.pure (x.2 y)) |
| 55 | + |
| 56 | +/-- Lift a position of the base polynomial functor into the free monad. -/ |
| 57 | +@[always_inline, inline, reducible] |
| 58 | +def liftA (a : P.A) : FreeM P (P.B a) := lift ⟨a, id⟩ |
| 59 | + |
| 60 | +instance : MonadLift P (FreeM P) where |
| 61 | + monadLift x := FreeM.lift x |
| 62 | + |
| 63 | +@[simp] lemma lift_ne_pure (x : P α) (y : α) : |
| 64 | + (lift x : FreeM P α) ≠ PFunctor.FreeM.pure y := by simp [lift] |
| 65 | + |
| 66 | +@[simp] lemma pure_ne_lift (x : P α) (y : α) : |
| 67 | + PFunctor.FreeM.pure y ≠ (lift x : FreeM P α) := by simp [lift] |
| 68 | + |
| 69 | +lemma monadLift_eq_lift (x : P.Obj α) : (x : FreeM P α) = FreeM.lift x := rfl |
| 70 | + |
| 71 | +/-- Bind operator on `FreeM P` used in the monad definition. -/ |
| 72 | +@[always_inline, inline] |
| 73 | +protected def bind : FreeM P α → (α → FreeM P β) → FreeM P β |
| 74 | + | FreeM.pure a, f => f a |
| 75 | + | FreeM.roll a cont, f => FreeM.roll a (fun u ↦ FreeM.bind (cont u) f) |
| 76 | + |
| 77 | +@[simp] |
| 78 | +lemma bind_pure (a : α) (f : α → FreeM P β) : |
| 79 | + FreeM.bind (FreeM.pure a) f = f a := rfl |
| 80 | + |
| 81 | +@[simp] |
| 82 | +lemma bind_roll (a : P.A) (cont : P.B a → FreeM P β) (f : β → FreeM P γ) : |
| 83 | + FreeM.bind (FreeM.roll a cont) f = FreeM.roll a (fun u ↦ FreeM.bind (cont u) f) := rfl |
| 84 | + |
| 85 | +@[simp] |
| 86 | +lemma bind_lift (x : P.Obj α) (f : α → FreeM P β) : |
| 87 | + FreeM.bind (FreeM.lift x) f = FreeM.roll x.1 (fun a ↦ f (x.2 a)) := rfl |
| 88 | + |
| 89 | +@[simp] lemma bind_eq_pure_iff (x : FreeM P α) (f : α → FreeM P β) (b : β) : |
| 90 | + FreeM.bind x f = FreeM.pure b ↔ ∃ a, x = pure a ∧ f a = pure b := by |
| 91 | + cases x <;> simp |
| 92 | + |
| 93 | +@[simp] lemma pure_eq_bind_iff (x : FreeM P α) (f : α → FreeM P β) (b : β) : |
| 94 | + FreeM.pure b = FreeM.bind x f ↔ ∃ a, x = pure a ∧ pure b = f a := by |
| 95 | + cases x <;> simp |
| 96 | + |
| 97 | +instance : Monad (FreeM P) where |
| 98 | + pure := FreeM.pure |
| 99 | + bind := FreeM.bind |
| 100 | + |
| 101 | +lemma monad_pure_def (a : α) : (pure a : FreeM P α) = FreeM.pure a := rfl |
| 102 | + |
| 103 | +lemma monad_bind_def (x : FreeM P α) (f : α → FreeM P β) : |
| 104 | + x >>= f = FreeM.bind x f := rfl |
| 105 | + |
| 106 | +instance : LawfulMonad (FreeM P) := |
| 107 | + LawfulMonad.mk' (FreeM P) |
| 108 | + (fun x ↦ by |
| 109 | + induction x with |
| 110 | + | pure _ => rfl |
| 111 | + | roll a _ h => refine congr_arg (FreeM.roll a) (funext fun i ↦ h i)) |
| 112 | + (fun x f ↦ rfl) |
| 113 | + (fun x f g ↦ by |
| 114 | + induction x with |
| 115 | + | pure _ => rfl |
| 116 | + | roll a _ h => refine congr_arg (FreeM.roll a) (funext fun i ↦ h i)) |
| 117 | + |
| 118 | +lemma pure_inj (a b : α) : FreeM.pure (P := P) a = FreeM.pure b ↔ a = b := by simp |
| 119 | + |
| 120 | +@[simp] lemma roll_inj (a a' : P.A) |
| 121 | + (cont : P.B a → P.FreeM α) (cont' : P.B a' → P.FreeM α) : |
| 122 | + FreeM.roll a cont = FreeM.roll a' cont' ↔ ∃ h : a = a', h ▸ cont = cont' := by |
| 123 | + simp |
| 124 | + by_cases ha : a = a' |
| 125 | + · cases ha |
| 126 | + simp |
| 127 | + · simp [ha] |
| 128 | + |
| 129 | +/-- Proving a predicate `C` of `FreeM P α` requires two cases: |
| 130 | +* `pure a` for some `a : α` |
| 131 | +* `roll a cont h` for some `a : P.A`, `cont : P.B a → FreeM P α`, and `h : ∀ y, C (cont y)` -/ |
| 132 | +@[elab_as_elim] |
| 133 | +protected def inductionOn {C : FreeM P α → Prop} |
| 134 | + (pure : ∀ a, C (pure a)) |
| 135 | + (roll : (a : P.A) → (cont : P.B a → FreeM P α) → (∀ y, C (cont y)) → |
| 136 | + C (FreeM.roll a cont)) : |
| 137 | + (x : FreeM P α) → C x |
| 138 | + | FreeM.pure a => pure a |
| 139 | + | FreeM.roll a cont => roll a _ (fun u ↦ FreeM.inductionOn pure roll (cont u)) |
| 140 | + |
| 141 | +section construct |
| 142 | + |
| 143 | +/-- Dependent eliminator for `FreeM P`. -/ |
| 144 | +@[elab_as_elim] |
| 145 | +protected def construct {C : FreeM P α → Type*} |
| 146 | + (pure : (a : α) → C (pure a)) |
| 147 | + (roll : (a : P.A) → (cont : P.B a → FreeM P α) → ((y : P.B a) → C (cont y)) → |
| 148 | + C (FreeM.roll a cont)) : |
| 149 | + (x : FreeM P α) → C x |
| 150 | + | .pure a => pure a |
| 151 | + | .roll a cont => roll a _ (fun u ↦ FreeM.construct pure roll (cont u)) |
| 152 | + |
| 153 | +variable {C : FreeM P α → Type*} (h_pure : (a : α) → C (pure a)) |
| 154 | + (h_roll : (a : P.A) → (cont : P.B a → FreeM P α) → ((y : P.B a) → C (cont y)) → |
| 155 | + C (FreeM.roll a cont)) |
| 156 | + |
| 157 | +@[simp] |
| 158 | +lemma construct_pure (a : α) : FreeM.construct h_pure h_roll (pure a) = h_pure a := rfl |
| 159 | + |
| 160 | +@[simp] |
| 161 | +lemma construct_roll (a : P.A) (cont : P.B a → FreeM P α) : |
| 162 | + (FreeM.construct h_pure h_roll (FreeM.roll a cont) : C (FreeM.roll a cont)) = |
| 163 | + (h_roll a cont (fun u => FreeM.construct h_pure h_roll (cont u))) := rfl |
| 164 | + |
| 165 | +end construct |
| 166 | + |
| 167 | +section mapM |
| 168 | + |
| 169 | +variable {m : Type uB → Type v} {α : Type uB} |
| 170 | + |
| 171 | +/-- Canonical mapping of `FreeM P` into any other monad, given a map `s : (a : P.A) → m (P.B a)`. |
| 172 | +-/ |
| 173 | +protected def mapM [Pure m] [Bind m] (s : (a : P.A) → m (P.B a)) : FreeM P α → m α |
| 174 | + | .pure a => Pure.pure a |
| 175 | + | .roll a cont => (s a) >>= (fun u ↦ (cont u).mapM s) |
| 176 | + |
| 177 | +variable [Monad m] (s : (a : P.A) → m (P.B a)) |
| 178 | + |
| 179 | +@[simp] |
| 180 | +lemma mapM_pure' (a : α) : (FreeM.pure a : FreeM P α).mapM s = Pure.pure a := rfl |
| 181 | + |
| 182 | +@[simp] |
| 183 | +lemma mapM_roll (a : P.A) (cont : P.B a → FreeM P α) : |
| 184 | + (FreeM.roll a cont).mapM s = s a >>= fun u => (cont u).mapM s := rfl |
| 185 | + |
| 186 | +@[simp] lemma mapM_pure (a : α) : (Pure.pure a : FreeM P α).mapM s = Pure.pure a := rfl |
| 187 | + |
| 188 | +variable [LawfulMonad m] |
| 189 | + |
| 190 | +@[simp] |
| 191 | +lemma mapM_bind {α β : Type uB} (x : FreeM P α) (f : α → FreeM P β) : |
| 192 | + (FreeM.bind x f).mapM s = x.mapM s >>= fun u => (f u).mapM s := by |
| 193 | + induction x using FreeM.inductionOn with |
| 194 | + | pure _ => simp |
| 195 | + | roll a cont h => simp [h] |
| 196 | + |
| 197 | +@[simp] |
| 198 | +lemma mapM_bind' {α β : Type uB} (x : FreeM P α) (f : α → FreeM P β) : |
| 199 | + (x >>= f).mapM s = x.mapM s >>= fun u => (f u).mapM s := |
| 200 | + mapM_bind _ _ _ |
| 201 | + |
| 202 | +@[simp] |
| 203 | +lemma mapM_map {α β : Type uB} (x : FreeM P α) (f : α → β) : |
| 204 | + FreeM.mapM s (f <$> x) = f <$> FreeM.mapM s x := by |
| 205 | + simp [← bind_pure_comp] |
| 206 | + |
| 207 | +@[simp] |
| 208 | +lemma mapM_seq {α β : Type uB} |
| 209 | + (s : (a : P.A) → m (P.B a)) (x : FreeM P (α → β)) (y : FreeM P α) : |
| 210 | + FreeM.mapM s (x <*> y) = (FreeM.mapM s x) <*> (FreeM.mapM s y) := by |
| 211 | + simp [seq_eq_bind_map] |
| 212 | + |
| 213 | +@[simp] |
| 214 | +lemma mapM_lift (s : (a : P.A) → m (P.B a)) (x : P.Obj α) : |
| 215 | + FreeM.mapM s (FreeM.lift x) = s x.1 >>= (fun u ↦ (pure (x.2 u)).mapM s) := by |
| 216 | + simp [FreeM.mapM] |
| 217 | + |
| 218 | +@[simp] |
| 219 | +lemma mapM_liftA (s : (a : P.A) → m (P.B a)) (a : P.A) : |
| 220 | + FreeM.mapM s (FreeM.liftA a) = s a := by simp [liftA] |
| 221 | + |
| 222 | +end mapM |
| 223 | + |
| 224 | +end FreeM |
| 225 | + |
| 226 | +end PFunctor |
0 commit comments