feat(PFunctor): add resumptions and relate W-types, M-types and free monads

Cslib.lean

@@ -139,6 +139,8 @@ public import Cslib.Foundations.Data.PFunctor.Basic
 public import Cslib.Foundations.Data.PFunctor.Free
 public import Cslib.Foundations.Data.PFunctor.Free.Fold
 public import Cslib.Foundations.Data.PFunctor.Free.W
+public import Cslib.Foundations.Data.PFunctor.M
+public import Cslib.Foundations.Data.PFunctor.Resumption
 public import Cslib.Foundations.Data.Set.Saturation
 public import Cslib.Foundations.Data.StackTape
 public import Cslib.Foundations.Lint.Basic

Cslib/Foundations/Data/PFunctor/Basic.lean

@@ -21,12 +21,14 @@ Special cases `C`, `linear`, `selfMonomial`, `purePower`, the indeterminate `y`,
 and canonical choices of `0` and `1` are defined as abbreviations or instances over `monomial`.
 The scoped notations `A y^ B` and `y^ B` denote `monomial A B` and `purePower B`, respectively.
 
-The child-map API includes `const`, `Unary`, and `DecidableEqChildren`.
+The extensions of `P + Q` and `C A` are described by `addObjEquiv` and `constObjEquiv`.
+The child-map API includes `const`, `Unary`, and `DecidableEqChildren`. `W.induction` is an
+induction principle for `P.W` through `W.mk`.
 -/
 
 @[expose] public section
 
-universe uA uB uA₁ uA₂ uB₁ uB₂
+universe uA uB uA₁ uA₂ uB₁ uB₂ v w
 
 namespace PFunctor
 
@@ -155,6 +157,75 @@ defined as the product of the head types and the sum of the child types. -/
 
 end prod
 
+section Obj
+
+variable {X : Type v} {Y : Type w}
+
+@[simp]
+theorem map_id' (P : PFunctor.{uA, uB}) : P.map (id : X → X) = id :=
+  funext P.id_map
+
+@[simp]
+theorem map_comp_map (P : PFunctor.{uA, uB}) {Z : Type*} (f : X → Y) (g : Y → Z) :
+    P.map g ∘ P.map f = P.map (g ∘ f) :=
+  funext (P.map_map f g)
+
+/-- The extension of a sum of polynomial functors is the sum of their extensions. -/
+def addObjEquiv (P : PFunctor.{uA₁, uB}) (Q : PFunctor.{uA₂, uB}) (X : Type v) :
+    (P + Q).Obj X ≃ P.Obj X ⊕ Q.Obj X where
+  toFun
+    | .mk (.inl a) f => .inl (.mk a f)
+    | .mk (.inr a) f => .inr (.mk a f)
+  invFun
+    | .inl (.mk a f) => .mk (.inl a) f
+    | .inr (.mk a f) => .mk (.inr a) f
+  left_inv x := by cases x with | mk s f => cases s <;> rfl
+  right_inv x := by rcases x with (x | x) <;> cases x <;> rfl
+
+@[simp]
+theorem addObjEquiv_mk_inl (P : PFunctor.{uA₁, uB}) (Q : PFunctor.{uA₂, uB}) (a : P.A)
+    (f : P.B a → X) : addObjEquiv P Q X (.mk (.inl a) f) = .inl (.mk a f) := rfl
+
+@[simp]
+theorem addObjEquiv_mk_inr (P : PFunctor.{uA₁, uB}) (Q : PFunctor.{uA₂, uB}) (a : Q.A)
+    (f : Q.B a → X) : addObjEquiv P Q X (.mk (.inr a) f) = .inr (.mk a f) := rfl
+
+theorem addObjEquiv_map (P : PFunctor.{uA₁, uB}) (Q : PFunctor.{uA₂, uB}) (f : X → Y)
+    (x : (P + Q).Obj X) :
+    addObjEquiv P Q Y ((P + Q).map f x) = Sum.map (P.map f) (Q.map f) (addObjEquiv P Q X x) := by
+  cases x with | mk s g => cases s <;> rfl
+
+/-- The extension of a constant polynomial functor is the constant. -/
+def constObjEquiv (A : Type uA) (X : Type v) : (C.{uA, uB} A).Obj X ≃ A where
+  toFun x := x.fst
+  invFun a := .mk a PEmpty.elim
+  left_inv x := by cases x with | mk a f => exact congrArg (Obj.mk a) (funext (·.elim))
+  right_inv _ := rfl
+
+@[simp]
+theorem constObjEquiv_mk (A : Type uA) (a : A) (f : PEmpty.{uB + 1} → X) :
+    constObjEquiv A X (.mk a f) = a := rfl
+
+@[simp]
+theorem constObjEquiv_map (A : Type uA) (f : X → Y) (x : (C.{uA, uB} A).Obj X) :
+    constObjEquiv A Y ((C A).map f x) = constObjEquiv A X x := rfl
+
+end Obj
+
+section W
+
+variable {P : PFunctor.{uA, uB}}
+
+/-- Induction on `P.W` through `W.mk`, keeping subtrees typed as `P.W` rather than `WType P.B`. -/
+@[elab_as_elim, induction_eliminator]
+protected theorem W.induction {motive : P.W → Prop}
+    (mk : ∀ (a : P.A) (f : P.B a → P.W), (∀ i, motive (f i)) → motive (W.mk (.mk a f)))
+    (w : P.W) : motive w := by
+  induction w using WType.rec with
+  | mk a f ih => exact mk a f ih
+
+end W
+
 section Unary
 
 /-- A polynomial functor is unary if all child types have exactly one element. -/

Cslib/Foundations/Data/PFunctor/Free.lean

@@ -16,7 +16,9 @@ public import Mathlib.Data.PFunctor.Univariate.Basic
 We define the free monad on a **polynomial functor** (`PFunctor`), and prove some basic properties.
 
 The free monad `PFunctor.FreeM P` extends the W-type construction with an extra `pure`
-constructor, yielding a monad that is free over the polynomial functor `P`.
+constructor, yielding a monad that is free over the polynomial functor `P`. `FreeM.equivW`
+identifies `P.FreeM α` with the W-type of `C α + P`, and the M-type of the same polynomial gives
+the coinductive counterpart `PFunctor.Resumption`, whose programs may run forever.
 
 ## Comparison with `Cslib.FreeM`
 

Cslib/Foundations/Data/PFunctor/Free/W.lean

@@ -6,13 +6,18 @@ Authors: Quang Dao, Devon Tuma
 
 module
 
+public import Cslib.Foundations.Data.PFunctor.Basic
 public import Cslib.Foundations.Data.PFunctor.Free
 
 /-!
-# Polynomial free monads with an empty result type
+# Polynomial free monads as W-types
 
-When `α` is empty, a tree in `P.FreeM α` has only operation nodes. The equivalence
-`PFunctor.FreeM.equivWOfIsEmpty` identifies these trees with the W-type `P.W`.
+A tree in `P.FreeM α` is a W-tree whose nodes either return a value of `α` or perform an operation
+of `P`: `PFunctor.FreeM.equivW` identifies `P.FreeM α` with the W-type of `C α + P`, the
+polynomial of the functor `X ↦ α ⊕ P X` whose initial algebra is the free monad.
+
+When `α` is empty, a tree has only operation nodes, and `PFunctor.FreeM.equivWOfIsEmpty` identifies
+these trees with the W-type `P.W` itself.
 -/
 
 @[expose] public section
@@ -69,4 +74,50 @@ def FreeM.equivWOfIsEmpty [IsEmpty α] : P.FreeM α ≃ P.W where
   left_inv := W.toFreeM_toWOfIsEmpty
   right_inv := toWOfIsEmpty_toFreeM
 
+/-- Regard a free program as a W-tree of `C α + P`, whose leaves carry the returned values. -/
+def FreeM.toW : P.FreeM α → (C.{u, uB} α + P).W
+  | .pure a => W.mk (.mk (.inl a) PEmpty.elim)
+  | .liftBind a cont => W.mk (.mk (.inr a) fun b => FreeM.toW (cont b))
+
+/-- Read a W-tree of `C α + P` as a free program. -/
+def FreeM.ofW : (C.{u, uB} α + P).W → P.FreeM α
+  | WType.mk (.inl a) _ => .pure a
+  | WType.mk (.inr a) cont => .liftBind a fun b => FreeM.ofW (cont b)
+
+@[simp]
+theorem FreeM.toW_pure (a : α) :
+    toW (pure a : P.FreeM α) = (W.mk (.mk (.inl a) PEmpty.elim) : (C.{u, uB} α + P).W) := rfl
+
+@[simp]
+theorem FreeM.toW_lift_bind (a : P.A) (cont : P.B a → P.FreeM α) :
+    toW ((lift a).bind (α := no_index (P.B a)) cont) =
+      (W.mk (.mk (.inr a) fun b => toW (cont b)) : (C.{u, uB} α + P).W) := rfl
+
+@[simp]
+theorem FreeM.toW_lift_bind' {α : Type uB} (a : P.A) (cont : P.B a → P.FreeM α) :
+    toW (Bind.bind (α := no_index (P.B a)) (lift a) cont) =
+      (W.mk (.mk (.inr a) fun b => toW (cont b)) : (C.{uB, uB} α + P).W) := rfl
+
+@[simp]
+theorem FreeM.ofW_toW (x : P.FreeM α) : ofW (toW x) = x := by
+  induction x with
+  | pure a => rfl
+  | lift_bind a cont ih => exact congrArg (liftBind a) (funext ih)
+
+@[simp]
+theorem FreeM.toW_ofW (w : (C.{u, uB} α + P).W) : toW (ofW w) = w := by
+  induction w with
+  | mk a f ih =>
+    cases a with
+    | inl a => exact congrArg (fun f => (W.mk (.mk (.inl a) f) : (C α + P).W)) (funext (·.elim))
+    | inr a => exact congrArg (fun f => (W.mk (.mk (.inr a) f) : (C α + P).W)) (funext ih)
+
+/-- Free programs are the W-trees of `C α + P`. -/
+@[simps]
+def FreeM.equivW : P.FreeM α ≃ (C.{u, uB} α + P).W where
+  toFun := toW
+  invFun := ofW
+  left_inv := ofW_toW
+  right_inv := toW_ofW
+
 end PFunctor

Cslib/Foundations/Data/PFunctor/M.lean

@@ -0,0 +1,120 @@
+/-
+Copyright (c) 2026 PolyFun Contributors. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Devon Tuma, Quang Dao
+-/
+
+module
+
+public import Cslib.Foundations.Data.PFunctor.Basic
+public import Mathlib.Data.PFunctor.Univariate.M
+
+/-!
+# W-types as well-founded M-types
+
+For a polynomial functor `P`, the W-type `P.W` is its initial algebra (well-founded trees) and the
+M-type `P.M` is its final coalgebra (possibly infinite trees). The canonical map `W.toM` regards a
+well-founded tree as a possibly infinite one, and `W.equivM` identifies `P.W` with the M-trees that
+are accessible for the immediate-subtree relation, i.e. have no infinite descending path.
+
+Lambek's lemma `M.destEquiv` packages the destructor of `P.M` as an equivalence.
+-/
+
+@[expose] public section
+
+universe uA uB
+
+namespace PFunctor
+
+variable {P : PFunctor.{uA, uB}}
+
+/-- Lambek's lemma: the destructor of the final coalgebra is an equivalence. -/
+@[simps]
+def M.destEquiv : P.M ≃ P P.M where
+  toFun := M.dest
+  invFun := M.mk
+  left_inv := M.mk_dest
+  right_inv := M.dest_mk
+
+theorem M.dest_injective : Function.Injective (M.dest (F := P)) := M.destEquiv.injective
+
+@[simp]
+theorem M.dest_inj {x y : P.M} : M.dest x = M.dest y ↔ x = y := M.dest_injective.eq_iff
+
+@[simp]
+theorem M.corec_dest : M.corec (M.dest (F := P)) = id :=
+  (M.corec_unique _ _ fun _ => (P.id_map _).symm).symm
+
+/-- Corecursion is natural in maps of states that commute with the step functions. -/
+theorem M.corec_comp {X Y : Type*} (g : X → P X) (h : Y → P Y) (f : X → Y)
+    (hf : ∀ x, h (f x) = P.map f (g x)) : M.corec h ∘ f = M.corec g :=
+  M.corec_unique g _ fun x => by simp [M.dest_corec, hf]
+
+/-- The canonical map from the initial algebra into the final coalgebra, regarding a well-founded
+tree as a possibly infinite one. -/
+def W.toM : P.W → P.M :=
+  M.corec W.dest
+
+@[simp]
+theorem W.toM_mk (a : P.A) (f : P.B a → P.W) :
+    (W.mk (.mk a f)).toM = M.mk (.mk a fun i => (f i).toM) :=
+  M.dest_injective (M.dest_corec _ _)
+
+namespace M
+
+/-- `c` is an immediate subtree of `t`. -/
+def IsChild (c t : P.M) : Prop :=
+  ∃ i, (M.dest t).snd i = c
+
+/-- An M-tree is well-founded when it is accessible for the immediate-subtree relation, i.e. it has
+no infinite descending path. Its branches need not have a common depth bound. -/
+abbrev IsWellFounded (t : P.M) : Prop :=
+  Acc IsChild t
+
+theorem isWellFounded_mk {a : P.A} {f : P.B a → P.M} :
+    (M.mk (.mk a f)).IsWellFounded ↔ ∀ i, (f i).IsWellFounded :=
+  ⟨fun h i => h.inv ⟨i, rfl⟩, fun h => ⟨_, fun _ ⟨i, hi⟩ => hi ▸ h i⟩⟩
+
+/-- The W-tree represented by a well-founded M-tree. -/
+def toW (t : P.M) (h : t.IsWellFounded) : P.W :=
+  Acc.rec (motive := fun _ _ => P.W)
+    (fun t _ ih => W.mk (.mk (M.dest t).fst fun i => ih _ ⟨i, rfl⟩)) h
+
+@[simp]
+theorem toW_mk (a : P.A) (f : P.B a → P.M) (h : (M.mk (.mk a f)).IsWellFounded) :
+    (M.mk (.mk a f)).toW h = W.mk (.mk a fun i => (f i).toW (isWellFounded_mk.1 h i)) := by
+  cases h; rfl
+
+end M
+
+theorem W.isWellFounded_toM (w : P.W) : w.toM.IsWellFounded := by
+  induction w with
+  | mk a f ih => simpa [M.isWellFounded_mk] using ih
+
+@[simp]
+theorem M.toW_toM (w : P.W) (h : w.toM.IsWellFounded) : w.toM.toW h = w := by
+  induction w with
+  | mk a f ih => simp [ih]
+
+theorem W.toM_toW (t : P.M) (h : t.IsWellFounded) : (t.toW h).toM = t := by
+  induction h with
+  | intro t _ ih =>
+    induction t using M.cases with
+    | f x =>
+      cases x with
+      | mk a f =>
+        rw [M.toW_mk, W.toM_mk]
+        exact congrArg (M.mk <| .mk a ·) (funext fun i => ih _ ⟨i, rfl⟩)
+
+/-- W-trees are exactly the well-founded M-trees. -/
+@[simps]
+def W.equivM : P.W ≃ {t : P.M // t.IsWellFounded} where
+  toFun w := ⟨w.toM, w.isWellFounded_toM⟩
+  invFun t := t.1.toW t.2
+  left_inv w := M.toW_toM w _
+  right_inv t := Subtype.ext (W.toM_toW t.1 t.2)
+
+theorem W.toM_injective : Function.Injective (W.toM (P := P)) :=
+  fun _ _ h => W.equivM.injective (Subtype.ext h)
+
+end PFunctor