feat(PFunctor): measure semantics for free monads and resumptions

Cslib.lean

@@ -138,14 +138,19 @@ public import Cslib.Foundations.Data.OmegaSequence.Topology
 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.Measure
 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.PFunctor.Resumption.Measure
 public import Cslib.Foundations.Data.Set.Saturation
 public import Cslib.Foundations.Data.StackTape
 public import Cslib.Foundations.Lint.Basic
 public import Cslib.Foundations.Logic.InferenceSystem
 public import Cslib.Foundations.Logic.LogicalEquivalence
 public import Cslib.Foundations.Logic.Operators
 public import Cslib.Foundations.MeasureTheory.FiniteSupport
+public import Cslib.Foundations.MeasureTheory.Monotone
 public import Cslib.Foundations.Relation.Attr
 public import Cslib.Foundations.Relation.Basic
 public import Cslib.Foundations.Relation.Confluence

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/Measure.lean

@@ -0,0 +1,108 @@
+/-
+Copyright (c) 2026 Devon Tuma. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Devon Tuma
+-/
+
+module
+
+public import Cslib.Foundations.Data.PFunctor.Free.Fold
+public import Mathlib.MeasureTheory.Measure.ProbabilityMeasure
+
+/-!
+# Output measures of polynomial programs
+
+Given a measure `μ a` on the responses of each operation `a`, a program `x : P.FreeM α` has the
+output measure `x.toMeasure μ`: each operation is answered according to `μ`, and the returned value
+is observed. It is the fold of `x` into Mathlib's Giry bind, so it follows `Measure.bind`'s
+convention for continuations that are not measurable. When the response spaces are discrete, every
+continuation is measurable, `toMeasure` sends `bind` to the Giry bind (`toMeasure_bind`), and it
+sends programs to probability measures whenever each `μ a` is one.
+-/
+
+@[expose] public section
+
+open MeasureTheory
+
+universe uA uB u v
+
+namespace PFunctor.FreeM
+
+variable {P : PFunctor.{uA, uB}} [∀ a, MeasurableSpace (P.B a)] {α : Type u} {β : Type v}
+  [MeasurableSpace α] [MeasurableSpace β]
+
+/-- The output measure of `x` when each operation `a` is answered according to `μ a`. -/
+noncomputable def toMeasure (x : P.FreeM α) (μ : (a : P.A) → Measure (P.B a)) : Measure α :=
+  x.foldFreeM .dirac fun a k => (μ a).bind k
+
+variable (μ : (a : P.A) → Measure (P.B a))
+
+@[simp]
+theorem toMeasure_pure (a : α) : (pure a : P.FreeM α).toMeasure μ = .dirac a := rfl
+
+@[simp]
+theorem toMeasure_lift_bind (a : P.A) (k : P.B a → P.FreeM α) :
+    ((lift a).bind (α := no_index (P.B a)) k).toMeasure μ =
+      (μ a).bind fun b => (k b).toMeasure μ := rfl
+
+@[simp]
+theorem toMeasure_lift_bind' {α : Type uB} [MeasurableSpace α] (a : P.A) (k : P.B a → P.FreeM α) :
+    (Bind.bind (α := no_index (P.B a)) (lift a) k).toMeasure μ =
+      (μ a).bind fun b => (k b).toMeasure μ := rfl
+
+@[simp]
+theorem toMeasure_lift (a : P.A) : toMeasure (α := no_index (P.B a)) (lift a) μ = μ a :=
+  Measure.bind_dirac
+
+section Discrete
+
+variable [∀ a, DiscreteMeasurableSpace (P.B a)]
+
+theorem toMeasure_bind (x : P.FreeM α) {f : α → P.FreeM β}
+    (hf : Measurable fun a => (f a).toMeasure μ) :
+    (x.bind f).toMeasure μ = (x.toMeasure μ).bind fun a => (f a).toMeasure μ := by
+  induction x with
+  | pure a => exact (Measure.dirac_bind hf a).symm
+  | lift_bind a k ih =>
+    rw [liftBind_bind, toMeasure_lift_bind, toMeasure_lift_bind,
+      Measure.bind_bind Measurable.of_discrete.aemeasurable hf.aemeasurable]
+    exact congrArg _ (funext ih)
+
+theorem toMeasure_map (x : P.FreeM α) {f : α → β} (hf : Measurable f) :
+    (x.map f).toMeasure μ = (x.toMeasure μ).map f := by
+  rw [← bind_pure_comp, toMeasure_bind μ x (f := pure ∘ f) (Measure.measurable_dirac.comp hf)]
+  exact Measure.bind_dirac_eq_map _ hf
+
+@[simp]
+theorem toMeasure_bind_of_discrete [DiscreteMeasurableSpace α] (x : P.FreeM α)
+    (f : α → P.FreeM β) :
+    (x.bind f).toMeasure μ = (x.toMeasure μ).bind fun a => (f a).toMeasure μ :=
+  toMeasure_bind μ x .of_discrete
+
+@[simp]
+theorem toMeasure_bind_of_discrete' {α β : Type u} [MeasurableSpace α] [MeasurableSpace β]
+    [DiscreteMeasurableSpace α] (x : P.FreeM α) (f : α → P.FreeM β) :
+    (x >>= f).toMeasure μ = (x.toMeasure μ).bind fun a => (f a).toMeasure μ :=
+  toMeasure_bind_of_discrete μ x f
+
+@[simp]
+theorem toMeasure_map_of_discrete [DiscreteMeasurableSpace α] (x : P.FreeM α) (f : α → β) :
+    (x.map f).toMeasure μ = (x.toMeasure μ).map f :=
+  toMeasure_map μ x .of_discrete
+
+@[simp]
+theorem toMeasure_map_of_discrete' {α β : Type u} [MeasurableSpace α] [MeasurableSpace β]
+    [DiscreteMeasurableSpace α] (x : P.FreeM α) (f : α → β) :
+    (f <$> x).toMeasure μ = (x.toMeasure μ).map f :=
+  toMeasure_map_of_discrete μ x f
+
+instance [∀ a, IsProbabilityMeasure (μ a)] (x : P.FreeM α) :
+    IsProbabilityMeasure (x.toMeasure μ) := by
+  induction x with
+  | pure a => exact Measure.dirac.isProbabilityMeasure
+  | lift_bind a k ih =>
+    exact isProbabilityMeasure_bind Measurable.of_discrete.aemeasurable (.of_forall ih)
+
+end Discrete
+
+end PFunctor.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