From 494275e9e9515a9b3bd1af0b431c2aebb9eab23e Mon Sep 17 00:00:00 2001
From: Marc Ilunga <marc.ilunga@trailofbits.com>
Date: Tue, 6 Oct 2026 12:55:32 +0200
Subject: [PATCH 1/4] Probability: kernels, pushforward and sequential
 composition of laws
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`Distribution.bindK μ f`, the mixture of the kernel `f` over the law `μ`
(the composition of conditional distributions, CR18 §3.6.1), with the
lemmas the probabilistic automata of the next layer use: linearity,
composition of mixtures, the point kernel, pushforwards as point kernels
and through mixtures, restriction through mixtures, weight and
nonnegativity, and independent products as mixtures (`prod_eq_bindK`,
`pi_eq_bindK_update`). Probability laws get `ProbDist.map` with
`map_map`, and `ProbDist.bind`, defined from `bindK` (the one bind).

Part of #14.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 README.md                   |   2 +-
 THEORY.md                   |  14 +-
 src/Probability.lean        |   3 +-
 src/Probability/Kernel.lean | 369 ++++++++++++++++++++++++++++++++++++
 4 files changed, 385 insertions(+), 3 deletions(-)
 create mode 100644 src/Probability/Kernel.lean

diff --git a/README.md b/README.md
index ac648f8..28ee142 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@ summarizes the theory layer by layer, with links to the Lean declarations.
 
 | Library | Content |
 |---|---|
-| `Probability` | finitely supported distributions, statistical distance, couplings, expectation, counting |
+| `Probability` | finitely supported distributions, kernels, statistical distance, couplings, expectation, counting |
 | `RandomSystems` | systems and automata; deterministic converters and programs; distributions over deterministic systems; random systems given by their cumulative probabilities; probabilistic converters, attachment and parallel composition; transcript distance and distinguishers; games |
 | `ConstructiveCryptography.Specification`, `.Construction` | specifications, constructions and relaxations over any carrier |
 | `ConstructiveCryptography.CryptographicAlgebra` | the classes of cryptographic algebras: resource theories, parallel composition, compatible pseudo-metrics and distinguisher classes, relaxations |
diff --git a/THEORY.md b/THEORY.md
index b717af8..44ed8e9 100644
--- a/THEORY.md
+++ b/THEORY.md
@@ -842,7 +842,13 @@ R \;\simeq_{\mathrm{ptxt}}\; P R \;\simeq_{\mathrm{cca}}\; P C R \;\simeq_{\math
 A [**distribution**][Distribution] is a finitely supported function to $`\mathbb R`$, and a
 [**probability distribution**][ProbDist] one that is [nonnegative with weight one][isProbDist].
 Distributions have [masses][mass] of events, [pushforwards][fTransform], [products][prod] of two
-laws and of [families][pi], and [uniform][uniform] laws. The [**statistical distance**][statDist]
+laws and of [families][pi], and [uniform][uniform] laws. A [**kernel**][bindK] $`f`$ mixed over a
+law $`\mu`$ gives $`\sum_a \mu(a)\, f(a)`$, the composition of conditional distributions
+[3, §3.6.1]: [mixtures compose][bindK_bindK], a pushforward
+[is the point kernel of its function][bindK_single_comp], and an independent product
+[disintegrates][pi_eq_bindK_update] at a coordinate. Probability laws have
+[pushforwards][ProbDist.map] and [sequential compositions][ProbDist.bind]. The
+[**statistical distance**][statDist]
 satisfies the [triangle inequality][statDist_triangle] and [data processing][statDist_fTransform_le],
 $`\delta(fX, fY) \le \delta(X, Y)`$; with [positive parts][statDist_eq_weight_posPart] it is
 $`\delta(X, Y) = \lvert (X - Y)^{+} \rvert`$, and with [overlaps][statDist_eq_weight_sub_weight_inf]
@@ -1290,6 +1296,12 @@ agreeing with a map on a finite set $`S`$: there are $`\lvert Y \rvert^{\lvert X
 [fTransform]: src/Probability/Distribution.lean#L587
 [prod]: src/Probability/Distribution.lean#L942
 [pi]: src/Probability/Distribution.lean#L1583
+[bindK]: src/Probability/Kernel.lean#L110
+[bindK_bindK]: src/Probability/Kernel.lean#L156
+[bindK_single_comp]: src/Probability/Kernel.lean#L173
+[pi_eq_bindK_update]: src/Probability/Kernel.lean#L309
+[ProbDist.map]: src/Probability/Kernel.lean#L347
+[ProbDist.bind]: src/Probability/Kernel.lean#L359
 [uniform]: src/Probability/Distribution.lean#L521
 [statDist]: src/Probability/StatisticalDistance.lean#L88
 [statDist_triangle]: src/Probability/StatisticalDistance.lean#L243
diff --git a/src/Probability.lean b/src/Probability.lean
index 9e99c2c..1e583b3 100644
--- a/src/Probability.lean
+++ b/src/Probability.lean
@@ -2,8 +2,9 @@ import Probability.Counting
 import Probability.Coupling
 import Probability.Distribution
 import Probability.Expectation
+import Probability.Kernel
 import Probability.Lift
 import Probability.MultiCoupling
 import Probability.StatisticalDistance
 
-/-! Finitely supported distributions, statistical distance, couplings, expectation and counting. -/
+/-! Finitely supported distributions, kernels, statistical distance, couplings, expectation and counting. -/
diff --git a/src/Probability/Kernel.lean b/src/Probability/Kernel.lean
new file mode 100644
index 0000000..f929f56
--- /dev/null
+++ b/src/Probability/Kernel.lean
@@ -0,0 +1,369 @@
+import Probability.Distribution
+import Mathlib.LinearAlgebra.Finsupp.LinearCombination
+
+/-!
+# Kernels
+
+A kernel `f : A → Distribution B` gives a law of a second outcome for every value of a first one.
+Its mixture over a law `μ` of the first outcome is the law of the second,
+
+```
+  (bindK μ f)(b) = ∑ₐ μ(a) · f(a)(b),
+```
+
+the composition of conditional probability distributions [1, §3.6.1, p. 65] (`bindK`). A
+pushforward is the point kernel of its function, an independent product samples its factors one
+after the other, and kernels of probability laws mix probability laws into probability laws.
+
+## Main definitions
+
+* `Distribution.bindK μ f`: the mixture of the kernel `f` over `μ`
+* `Distribution.ProbDist.map f P`: the law of `f a` for `a ← P`
+* `Distribution.ProbDist.bind P f`: the law of `b ← f a` for `a ← P`
+
+## Main results
+
+* `Distribution.bindK_bindK`, `Distribution.bindK_single_one`: mixtures compose, and the point
+  kernel leaves a law unchanged
+* `Distribution.bindK_single_comp`, `Distribution.fTransform_bindK`,
+  `Distribution.bindK_fTransform`: a pushforward is the point kernel of its function, and passes
+  through mixtures
+* `Distribution.weight_bindK_of_weight`, `Distribution.bindK_nonNeg`: kernels of laws preserve the
+  weight and nonnegativity
+* `Distribution.prod_eq_bindK`, `Distribution.pi_eq_bindK_update`: independent products as
+  mixtures
+* `Distribution.ProbDist.map_map`: pushing forward twice is pushing forward along the composite
+
+## References
+
+1. U. Maurer. *Cryptography Foundations*. Lecture notes, ETH Zurich, Spring 2018.
+-/
+
+namespace Probability.Distribution
+
+open Classical
+
+variable {A B C : Type*}
+
+/-! ### Restriction and pushforward -/
+
+/-- Restricting keeps only mass that was there. -/
+theorem support_restrict_subset (μ : Distribution A) (P : A → Prop) :
+    (μ.restrict P).support ⊆ μ.support := fun a ha => by
+  -- The restriction at `a` is `μ a` or zero; it is nonzero, so it is `μ a`.
+  rw [Finsupp.mem_support_iff, restrict_apply] at ha
+  rw [Finsupp.mem_support_iff]
+  split_ifs at ha
+  · exact ha
+  · exact absurd rfl ha
+
+/-- Restricting twice is restricting to both events. -/
+theorem restrict_restrict (μ : Distribution A) (P Q : A → Prop) :
+    (μ.restrict P).restrict Q = μ.restrict fun a => P a ∧ Q a := by
+  -- Pointwise: `μ a` exactly when both events hold.
+  ext a
+  simp only [restrict_apply]
+  by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
+
+/-- Restriction depends only on the event up to equivalence. -/
+theorem restrict_congr (μ : Distribution A) {P Q : A → Prop} (h : ∀ a, P a ↔ Q a) :
+    μ.restrict P = μ.restrict Q := by
+  -- Pointwise, the two events are the same condition.
+  ext a
+  simp only [restrict_apply, h]
+
+/-- Restricting to the impossible event leaves nothing. -/
+theorem restrict_false (μ : Distribution A) : μ.restrict (fun _ => False) = 0 := by
+  -- Pointwise, the condition never holds.
+  ext a
+  simp
+
+/-- A point mass restricted to an event: itself or nothing. -/
+theorem single_restrict (a : A) (w : ℝ) (P : A → Prop) :
+    Distribution.restrict (Finsupp.single a w : Distribution A) P =
+      if P a then Finsupp.single a w else 0 := by
+  -- The only point with mass is `a`: kept if `P a`, dropped otherwise.
+  by_cases h : P a
+  · simp [restrict, Finsupp.filter_single_of_pos _ h, h]
+  · simp [restrict, Finsupp.filter_single_of_neg _ h, h]
+
+/-- The pushforward of a point mass. -/
+theorem fTransform_single (f : A → B) (a : A) (w : ℝ) :
+    fTransform f (Finsupp.single a w) = Finsupp.single (f a) w :=
+  Finsupp.mapDomain_single
+
+/-- Restricting a pushforward is pushing forward the restriction to the preimage. -/
+theorem restrict_fTransform (f : A → B) (μ : Distribution A) (P : B → Prop) :
+    (fTransform f μ).restrict P = fTransform f (μ.restrict fun a => P (f a)) := by
+  -- At `b`, both sides are masses of fibres of `f`.
+  ext b
+  rw [restrict_apply, fTransform_apply_eq_mass, fTransform_apply_eq_mass, mass_restrict]
+  split_ifs with h
+  · -- `P b`: on the fibre of `b`, the condition `P (f a)` holds
+    exact mass_congr _ fun a => ⟨fun ha => ⟨ha, ha ▸ h⟩, fun ha => ha.1⟩
+  · -- not `P b`: the fibre of `b` meets no point with `P (f a)`
+    exact (mass_eq_zero_of_forall_not _ fun a (ha : f a = b ∧ P (f a)) => h (ha.1 ▸ ha.2)).symm
+
+/-! ### Kernels -/
+
+/-- **The mixture of a kernel** over a distribution: `∑ a, μ a • f a`. -/
+noncomputable def bindK (μ : Distribution A) (f : A → Distribution B) : Distribution B :=
+  Finsupp.linearCombination ℝ f μ
+
+/-- Nothing mixes to nothing. -/
+@[simp] theorem bindK_zero (f : A → Distribution B) : bindK 0 f = 0 :=
+  map_zero _
+
+/-- A point mass of weight `w` mixes to `w` times its kernel law. -/
+@[simp] theorem bindK_single (a : A) (w : ℝ) (f : A → Distribution B) :
+    bindK (Finsupp.single a w) f = w • f a :=
+  Finsupp.linearCombination_single _ _ _
+
+/-- The mixture is additive in the first law. -/
+theorem bindK_add (μ ν : Distribution A) (f : A → Distribution B) :
+    bindK (μ + ν) f = bindK μ f + bindK ν f :=
+  map_add _ _ _
+
+/-- The mixture at `b`: `∑ₐ μ(a) · f(a)(b)`. -/
+theorem bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) :
+    bindK μ f b = μ.sum fun a w => w * f a b := by
+  -- A linear combination evaluated at `b` is the combination of the values at `b`.
+  unfold bindK
+  rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply]
+  rfl
+
+/-- A kernel that vanishes mixes to nothing. -/
+@[simp] theorem bindK_zero_right (μ : Distribution A) :
+    bindK μ (fun _ => (0 : Distribution B)) = 0 := by
+  -- Induction on `μ` as a sum of point masses.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp
+  | add μ ν hμ hν => rw [bindK_add, hμ, hν, add_zero]
+  | single a w => simp
+
+/-- A mixture depends only on the kernel on the support. -/
+theorem bindK_congr {μ : Distribution A} {f g : A → Distribution B}
+    (h : ∀ a ∈ μ.support, f a = g a) : bindK μ f = bindK μ g := by
+  -- Both are sums over the support of `μ`, term by term equal.
+  unfold bindK
+  rw [Finsupp.linearCombination_apply, Finsupp.linearCombination_apply]
+  exact Finset.sum_congr rfl fun a ha => by
+    show μ a • f a = μ a • g a
+    rw [h a ha]
+
+/-- **Mixtures compose**: mixing `g` over the mixture of `f` is mixing, over `μ`, the mixtures of
+`g` over each `f a`. -/
+theorem bindK_bindK (μ : Distribution A) (f : A → Distribution B) (g : B → Distribution C) :
+    bindK (bindK μ f) g = bindK μ fun a => bindK (f a) g := by
+  -- Induction on `μ`: both sides are additive in `μ`, and agree on point masses.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp
+  | add μ ν hμ hν => rw [bindK_add, bindK_add, bindK_add, hμ, hν]
+  | single a w => simp [bindK]
+
+/-- **The point kernel** leaves a distribution unchanged. -/
+theorem bindK_single_one (μ : Distribution A) : bindK μ (fun a => Finsupp.single a 1) = μ := by
+  -- Induction on `μ`: a point mass of weight `w` mixes to `w` times the point mass.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp
+  | add μ ν hμ hν => rw [bindK_add, hμ, hν]
+  | single a w => simp
+
+/-- **A point kernel along `φ` is the pushforward** along `φ`. -/
+theorem bindK_single_comp (μ : Distribution A) (φ : A → B) :
+    bindK μ (fun a => Finsupp.single (φ a) 1) = fTransform φ μ := by
+  -- Induction on `μ`: both sides are additive, and move a point mass from `a` to `φ a`.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp [fTransform]
+  | add μ ν hμ hν =>
+    rw [bindK_add, hμ, hν]
+    exact (Finsupp.mapDomain_add).symm
+  | single a w => simp [fTransform]
+
+/-- **Pushing a mixture forward** pushes each law of the kernel forward. -/
+theorem fTransform_bindK (φ : B → C) (μ : Distribution A) (f : A → Distribution B) :
+    fTransform φ (bindK μ f) = bindK μ fun a => fTransform φ (f a) := by
+  -- Induction on `μ`: the pushforward is linear.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp [fTransform]
+  | add μ ν hμ hν =>
+    rw [bindK_add, bindK_add, ← hμ, ← hν]
+    exact Finsupp.mapDomain_add
+  | single a w =>
+    rw [bindK_single, bindK_single]
+    exact Finsupp.mapDomain_smul _ _
+
+/-- **Mixing over a pushforward** is mixing the composite kernel. -/
+theorem bindK_fTransform (φ : A → B) (μ : Distribution A) (f : B → Distribution C) :
+    bindK (fTransform φ μ) f = bindK μ (f ∘ φ) := by
+  -- Induction on `μ`: a point mass at `a` is moved to `φ a`, where the kernel is `f (φ a)`.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp [fTransform]
+  | add μ ν hμ hν =>
+    rw [show fTransform φ (μ + ν) = fTransform φ μ + fTransform φ ν from Finsupp.mapDomain_add,
+      bindK_add, bindK_add, hμ, hν]
+  | single a w => simp [fTransform]
+
+/-- A pushforward depends only on the function on the support. -/
+theorem fTransform_congr_support {f g : A → B} {μ : Distribution A}
+    (h : ∀ a ∈ μ.support, f a = g a) : fTransform f μ = fTransform g μ := by
+  -- Both are point kernels, which agree on the support.
+  rw [← bindK_single_comp, ← bindK_single_comp]
+  exact bindK_congr fun a ha => by rw [h a ha]
+
+/-- **Restricting a mixture** restricts each law of the kernel. -/
+theorem restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribution B) :
+    (bindK μ f).restrict P = bindK μ fun a => (f a).restrict P := by
+  -- Induction on `μ`: restriction is linear.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp [restrict, Finsupp.filter_zero]
+  | add μ ν hμ hν => rw [bindK_add, bindK_add, restrict, Finsupp.filter_add, ← restrict,
+      ← restrict, hμ, hν]
+  | single a w => simp [restrict, Finsupp.filter_smul]
+
+/-- A kernel that vanishes off an event is the kernel on the restriction to it. -/
+theorem bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution B) :
+    bindK μ (fun a => if P a then f a else 0) = bindK (μ.restrict P) f := by
+  -- Induction on `μ`: a point mass at `a` contributes `f a` exactly when `P a`.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp [restrict, Finsupp.filter_zero]
+  | add μ ν hμ hν => rw [bindK_add, hμ, hν, ← bindK_add, restrict, restrict, restrict,
+      Finsupp.filter_add]
+  | single a w =>
+    by_cases h : P a
+    · simp [h, restrict, Finsupp.filter_single_of_pos _ h]
+    · simp [h, restrict, Finsupp.filter_single_of_neg _ h]
+
+/-- **The weight of a mixture**: each weight of the kernel, weighted by `μ`. -/
+theorem weight_bindK (μ : Distribution A) (f : A → Distribution B) :
+    (bindK μ f).weight = μ.sum fun a w => w * (f a).weight := by
+  -- Induction on `μ`: the weight is additive, and scales with a point mass.
+  induction μ using Finsupp.induction_linear with
+  | zero => simp [weight]
+  | add μ ν hμ hν =>
+    rw [bindK_add, weight_add, hμ, hν,
+      Finsupp.sum_add_index' (by simp) (fun _ _ _ => add_mul _ _ _)]
+  | single a w => simp [weight_smul]
+
+/-- **A kernel of weight one preserves the weight.** -/
+theorem weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B}
+    (hf : ∀ a ∈ μ.support, (f a).weight = 1) : (bindK μ f).weight = μ.weight := by
+  -- Each `μ a` is weighted by the weight of `f a`, which is one.
+  rw [weight_bindK, weight]
+  exact Finset.sum_congr rfl fun a ha => by
+    show μ a * (f a).weight = μ a
+    rw [hf a ha, mul_one]
+
+/-- **A nonnegative kernel mixed over a nonnegative distribution** is nonnegative. -/
+theorem bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.NonNeg)
+    (hf : ∀ a, (f a).NonNeg) : (bindK μ f).NonNeg := by
+  -- At each `b`, a sum of products of nonnegative numbers.
+  intro b
+  unfold bindK
+  rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply]
+  exact Finset.sum_nonneg fun a _ => by
+    simpa [Finsupp.smul_apply] using mul_nonneg (hμ a) (hf a b)
+
+/-- A possible outcome of a mixture is possible under the kernel at a possible first outcome. -/
+theorem mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b : B}
+    (hb : b ∈ (bindK μ f).support) : ∃ a ∈ μ.support, b ∈ (f a).support := by
+  -- Otherwise every term `μ a · f a b` of the mixture at `b` vanishes.
+  by_contra hn
+  push Not at hn
+  apply Finsupp.mem_support_iff.mp hb
+  unfold bindK
+  rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply]
+  exact Finset.sum_eq_zero fun a ha => by
+    simp [Finsupp.notMem_support_iff.mp (hn a ha)]
+
+/-! ### Independent products -/
+
+/-- **An independent product as a mixture**: sample the first factor, then the second beside it. -/
+theorem prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
+    prod μ ν = bindK μ fun a => fTransform (fun b => (a, b)) ν := by
+  -- Compare at a pair `(a, b)`: `μ a · ν b` on the left.
+  ext ⟨a, b⟩
+  rw [prod_apply, bindK_apply, Finsupp.sum]
+  -- The law `ν` placed beside `a'` puts `ν b` at `(a, b)` if `a' = a`, and nothing otherwise.
+  have hpt : ∀ a', fTransform (fun b' => (a', b')) ν (a, b) = if a' = a then ν b else 0 := by
+    intro a'
+    rw [fTransform_apply_eq_mass]
+    split_ifs with h
+    · subst h
+      rw [mass, Finsupp.sum]
+      simp only [Prod.mk.injEq, true_and]
+      rw [Finset.sum_ite_eq']
+      split_ifs with hb
+      · rfl
+      · exact (Finsupp.notMem_support_iff.mp hb).symm
+    · exact mass_eq_zero_of_forall_not _ fun b' h' => h (Prod.mk.inj h').1
+  -- So the mixture at `(a, b)` keeps only the term `a' = a`: `μ a · ν b`.
+  simp only [hpt, mul_ite, mul_zero]
+  rw [Finset.sum_ite_eq']
+  split_ifs with ha
+  · rfl
+  · rw [Finsupp.notMem_support_iff.mp ha, zero_mul]
+
+/-- **Disintegration of an independent product** at one coordinate: sample that coordinate first,
+then the product with that coordinate fixed. -/
+theorem pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι → Distribution A)
+    (n : ι) : pi μ = bindK (μ n) fun z => pi (Function.update μ n (Finsupp.single z 1)) := by
+  -- Compare at a sample `t`: `∏ᵢ μᵢ(tᵢ)` on the left.
+  ext t
+  -- With coordinate `n` fixed to `z`, the product is `[t n = z] · ∏_{i ≠ n} μᵢ(tᵢ)`.
+  have hrest : ∀ z : A, (∏ i, Function.update μ n (Finsupp.single z 1) i (t i)) =
+      (Finsupp.single z (1 : ℝ)) (t n) * ∏ i ∈ Finset.univ.erase n, μ i (t i) := by
+    intro z
+    rw [← Finset.mul_prod_erase Finset.univ _ (Finset.mem_univ n), Function.update_self]
+    congr 1
+    exact Finset.prod_congr rfl fun i hi => by
+      rw [Function.update_of_ne (Finset.ne_of_mem_erase hi)]
+  -- Mixing over `z ← μ n` keeps the term `z = t n`: `μ n (t n) · ∏_{i ≠ n} μᵢ(tᵢ)`.
+  rw [pi_apply, bindK_apply]
+  simp only [pi_apply, hrest, Finsupp.single_apply]
+  rw [← Finset.mul_prod_erase Finset.univ _ (Finset.mem_univ n), Finsupp.sum]
+  simp only [← mul_assoc, mul_ite, mul_one, mul_zero, ite_mul, zero_mul]
+  rw [Finset.sum_ite_eq']
+  split_ifs with h
+  · rfl
+  · rw [Finsupp.notMem_support_iff.mp h, zero_mul]
+
+/-- With one coordinate fixed to `z`, every possible sample has `z` there. -/
+theorem eq_of_mem_support_pi_update {ι : Type*} [Fintype ι] [DecidableEq ι]
+    (μ : ι → Distribution A) (n : ι) (z : A) {t : ι → A}
+    (ht : t ∈ (pi (Function.update μ n (Finsupp.single z 1))).support) : t n = z := by
+  -- Otherwise the factor at `n`, the point mass at `z`, vanishes at `t n`.
+  rw [Finsupp.mem_support_iff, pi_apply] at ht
+  by_contra hne
+  apply ht
+  refine Finset.prod_eq_zero (Finset.mem_univ n) ?_
+  rw [Function.update_self, Finsupp.single_apply, if_neg (Ne.symm hne)]
+
+/-! ### Probability laws -/
+
+namespace ProbDist
+
+/-- **The pushforward** of a law: the law of `f a` for `a ← P`. -/
+noncomputable def map (f : A → B) (P : ProbDist A) : ProbDist B :=
+  ⟨fTransform f P.1, fTransform_isProbDist f P.2⟩
+
+@[simp] theorem map_val (f : A → B) (P : ProbDist A) : (map f P).1 = fTransform f P.1 := rfl
+
+/-- **Pushing forward twice** is pushing forward along the composite. -/
+theorem map_map (f : B → C) (g : A → B) (P : ProbDist A) :
+    map f (map g P) = map (f ∘ g) P :=
+  -- Both sides have the same underlying distribution: `fTransform f ∘ fTransform g`.
+  Subtype.ext (fTransform_fTransform _ _ _)
+
+/-- **Sequential composition**: the law of `b ← f a` for `a ← P`, the mixture of `f` over `P`. -/
+noncomputable def bind (P : ProbDist A) (f : A → ProbDist B) : ProbDist B :=
+  ⟨bindK P.1 fun a => (f a).1,
+    bindK_nonNeg P.2.1 fun a => (f a).2.1,
+    (weight_bindK_of_weight fun a _ => (f a).2.2).trans P.2.2⟩
+
+@[simp] theorem bind_val (P : ProbDist A) (f : A → ProbDist B) :
+    (bind P f).1 = bindK P.1 fun a => (f a).1 := rfl
+
+end ProbDist
+
+end Probability.Distribution

From 8846092403cccff3bc78a03d5baca9fb666ea1ae Mon Sep 17 00:00:00 2001
From: Marc Ilunga <marc.ilunga@trailofbits.com>
Date: Tue, 6 Oct 2026 13:03:58 +0200
Subject: [PATCH 2/4] =?UTF-8?q?Kernels:=20cite=20CR18=20=C2=A73.6.1=20for?=
 =?UTF-8?q?=20conditional=20distributions=20only;=20lemma=20for=20steps?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CR18 §3.6.1 defines a conditional distribution p_{Y|X}; it does not define its
mixture over a law, so the module doc and THEORY.md now say only that. The root
doc of Probability is wrapped to 100 characters, and the steps of Kernel.lean
are lemmas (#18).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 THEORY.md                   |  6 ++--
 src/Probability.lean        |  3 +-
 src/Probability/Kernel.lean | 72 ++++++++++++++++++-------------------
 3 files changed, 41 insertions(+), 40 deletions(-)

diff --git a/THEORY.md b/THEORY.md
index 44ed8e9..50b8537 100644
--- a/THEORY.md
+++ b/THEORY.md
@@ -842,9 +842,9 @@ R \;\simeq_{\mathrm{ptxt}}\; P R \;\simeq_{\mathrm{cca}}\; P C R \;\simeq_{\math
 A [**distribution**][Distribution] is a finitely supported function to $`\mathbb R`$, and a
 [**probability distribution**][ProbDist] one that is [nonnegative with weight one][isProbDist].
 Distributions have [masses][mass] of events, [pushforwards][fTransform], [products][prod] of two
-laws and of [families][pi], and [uniform][uniform] laws. A [**kernel**][bindK] $`f`$ mixed over a
-law $`\mu`$ gives $`\sum_a \mu(a)\, f(a)`$, the composition of conditional distributions
-[3, §3.6.1]: [mixtures compose][bindK_bindK], a pushforward
+laws and of [families][pi], and [uniform][uniform] laws. A [**kernel**][bindK] $`f`$ is a conditional
+distribution [3, §3.6.1]; mixed over a law $`\mu`$ it gives $`\sum_a \mu(a)\, f(a)`$:
+[mixtures compose][bindK_bindK], a pushforward
 [is the point kernel of its function][bindK_single_comp], and an independent product
 [disintegrates][pi_eq_bindK_update] at a coordinate. Probability laws have
 [pushforwards][ProbDist.map] and [sequential compositions][ProbDist.bind]. The
diff --git a/src/Probability.lean b/src/Probability.lean
index 1e583b3..9447640 100644
--- a/src/Probability.lean
+++ b/src/Probability.lean
@@ -7,4 +7,5 @@ import Probability.Lift
 import Probability.MultiCoupling
 import Probability.StatisticalDistance
 
-/-! Finitely supported distributions, kernels, statistical distance, couplings, expectation and counting. -/
+/-! Finitely supported distributions, kernels, statistical distance, couplings, expectation and
+counting. -/
diff --git a/src/Probability/Kernel.lean b/src/Probability/Kernel.lean
index f929f56..b04e093 100644
--- a/src/Probability/Kernel.lean
+++ b/src/Probability/Kernel.lean
@@ -4,15 +4,15 @@ import Mathlib.LinearAlgebra.Finsupp.LinearCombination
 /-!
 # Kernels
 
-A kernel `f : A → Distribution B` gives a law of a second outcome for every value of a first one.
-Its mixture over a law `μ` of the first outcome is the law of the second,
+A kernel `f : A → Distribution B` gives a law of a second outcome for every value of a first one:
+a conditional probability distribution `p_{Y|X}` [1, §3.6.1, p. 65]. Its mixture over a law `μ`
+of the first outcome is the law of the second (`bindK`),
 
 ```
-  (bindK μ f)(b) = ∑ₐ μ(a) · f(a)(b),
+  (bindK μ f)(b) = ∑ₐ μ(a) · f(a)(b).
 ```
 
-the composition of conditional probability distributions [1, §3.6.1, p. 65] (`bindK`). A
-pushforward is the point kernel of its function, an independent product samples its factors one
+A pushforward is the point kernel of its function, an independent product samples its factors one
 after the other, and kernels of probability laws mix probability laws into probability laws.
 
 ## Main definitions
@@ -48,7 +48,7 @@ variable {A B C : Type*}
 /-! ### Restriction and pushforward -/
 
 /-- Restricting keeps only mass that was there. -/
-theorem support_restrict_subset (μ : Distribution A) (P : A → Prop) :
+lemma support_restrict_subset (μ : Distribution A) (P : A → Prop) :
     (μ.restrict P).support ⊆ μ.support := fun a ha => by
   -- The restriction at `a` is `μ a` or zero; it is nonzero, so it is `μ a`.
   rw [Finsupp.mem_support_iff, restrict_apply] at ha
@@ -58,7 +58,7 @@ theorem support_restrict_subset (μ : Distribution A) (P : A → Prop) :
   · exact absurd rfl ha
 
 /-- Restricting twice is restricting to both events. -/
-theorem restrict_restrict (μ : Distribution A) (P Q : A → Prop) :
+lemma restrict_restrict (μ : Distribution A) (P Q : A → Prop) :
     (μ.restrict P).restrict Q = μ.restrict fun a => P a ∧ Q a := by
   -- Pointwise: `μ a` exactly when both events hold.
   ext a
@@ -66,20 +66,20 @@ theorem restrict_restrict (μ : Distribution A) (P Q : A → Prop) :
   by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
 
 /-- Restriction depends only on the event up to equivalence. -/
-theorem restrict_congr (μ : Distribution A) {P Q : A → Prop} (h : ∀ a, P a ↔ Q a) :
+lemma restrict_congr (μ : Distribution A) {P Q : A → Prop} (h : ∀ a, P a ↔ Q a) :
     μ.restrict P = μ.restrict Q := by
   -- Pointwise, the two events are the same condition.
   ext a
   simp only [restrict_apply, h]
 
 /-- Restricting to the impossible event leaves nothing. -/
-theorem restrict_false (μ : Distribution A) : μ.restrict (fun _ => False) = 0 := by
+lemma restrict_false (μ : Distribution A) : μ.restrict (fun _ => False) = 0 := by
   -- Pointwise, the condition never holds.
   ext a
   simp
 
 /-- A point mass restricted to an event: itself or nothing. -/
-theorem single_restrict (a : A) (w : ℝ) (P : A → Prop) :
+lemma single_restrict (a : A) (w : ℝ) (P : A → Prop) :
     Distribution.restrict (Finsupp.single a w : Distribution A) P =
       if P a then Finsupp.single a w else 0 := by
   -- The only point with mass is `a`: kept if `P a`, dropped otherwise.
@@ -88,12 +88,12 @@ theorem single_restrict (a : A) (w : ℝ) (P : A → Prop) :
   · simp [restrict, Finsupp.filter_single_of_neg _ h, h]
 
 /-- The pushforward of a point mass. -/
-theorem fTransform_single (f : A → B) (a : A) (w : ℝ) :
+lemma fTransform_single (f : A → B) (a : A) (w : ℝ) :
     fTransform f (Finsupp.single a w) = Finsupp.single (f a) w :=
   Finsupp.mapDomain_single
 
 /-- Restricting a pushforward is pushing forward the restriction to the preimage. -/
-theorem restrict_fTransform (f : A → B) (μ : Distribution A) (P : B → Prop) :
+lemma restrict_fTransform (f : A → B) (μ : Distribution A) (P : B → Prop) :
     (fTransform f μ).restrict P = fTransform f (μ.restrict fun a => P (f a)) := by
   -- At `b`, both sides are masses of fibres of `f`.
   ext b
@@ -111,21 +111,21 @@ noncomputable def bindK (μ : Distribution A) (f : A → Distribution B) : Distr
   Finsupp.linearCombination ℝ f μ
 
 /-- Nothing mixes to nothing. -/
-@[simp] theorem bindK_zero (f : A → Distribution B) : bindK 0 f = 0 :=
+@[simp] lemma bindK_zero (f : A → Distribution B) : bindK 0 f = 0 :=
   map_zero _
 
 /-- A point mass of weight `w` mixes to `w` times its kernel law. -/
-@[simp] theorem bindK_single (a : A) (w : ℝ) (f : A → Distribution B) :
+@[simp] lemma bindK_single (a : A) (w : ℝ) (f : A → Distribution B) :
     bindK (Finsupp.single a w) f = w • f a :=
   Finsupp.linearCombination_single _ _ _
 
 /-- The mixture is additive in the first law. -/
-theorem bindK_add (μ ν : Distribution A) (f : A → Distribution B) :
+lemma bindK_add (μ ν : Distribution A) (f : A → Distribution B) :
     bindK (μ + ν) f = bindK μ f + bindK ν f :=
   map_add _ _ _
 
 /-- The mixture at `b`: `∑ₐ μ(a) · f(a)(b)`. -/
-theorem bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) :
+lemma bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) :
     bindK μ f b = μ.sum fun a w => w * f a b := by
   -- A linear combination evaluated at `b` is the combination of the values at `b`.
   unfold bindK
@@ -133,7 +133,7 @@ theorem bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) :
   rfl
 
 /-- A kernel that vanishes mixes to nothing. -/
-@[simp] theorem bindK_zero_right (μ : Distribution A) :
+@[simp] lemma bindK_zero_right (μ : Distribution A) :
     bindK μ (fun _ => (0 : Distribution B)) = 0 := by
   -- Induction on `μ` as a sum of point masses.
   induction μ using Finsupp.induction_linear with
@@ -142,7 +142,7 @@ theorem bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) :
   | single a w => simp
 
 /-- A mixture depends only on the kernel on the support. -/
-theorem bindK_congr {μ : Distribution A} {f g : A → Distribution B}
+lemma bindK_congr {μ : Distribution A} {f g : A → Distribution B}
     (h : ∀ a ∈ μ.support, f a = g a) : bindK μ f = bindK μ g := by
   -- Both are sums over the support of `μ`, term by term equal.
   unfold bindK
@@ -153,7 +153,7 @@ theorem bindK_congr {μ : Distribution A} {f g : A → Distribution B}
 
 /-- **Mixtures compose**: mixing `g` over the mixture of `f` is mixing, over `μ`, the mixtures of
 `g` over each `f a`. -/
-theorem bindK_bindK (μ : Distribution A) (f : A → Distribution B) (g : B → Distribution C) :
+lemma bindK_bindK (μ : Distribution A) (f : A → Distribution B) (g : B → Distribution C) :
     bindK (bindK μ f) g = bindK μ fun a => bindK (f a) g := by
   -- Induction on `μ`: both sides are additive in `μ`, and agree on point masses.
   induction μ using Finsupp.induction_linear with
@@ -162,7 +162,7 @@ theorem bindK_bindK (μ : Distribution A) (f : A → Distribution B) (g : B →
   | single a w => simp [bindK]
 
 /-- **The point kernel** leaves a distribution unchanged. -/
-theorem bindK_single_one (μ : Distribution A) : bindK μ (fun a => Finsupp.single a 1) = μ := by
+lemma bindK_single_one (μ : Distribution A) : bindK μ (fun a => Finsupp.single a 1) = μ := by
   -- Induction on `μ`: a point mass of weight `w` mixes to `w` times the point mass.
   induction μ using Finsupp.induction_linear with
   | zero => simp
@@ -170,7 +170,7 @@ theorem bindK_single_one (μ : Distribution A) : bindK μ (fun a => Finsupp.sing
   | single a w => simp
 
 /-- **A point kernel along `φ` is the pushforward** along `φ`. -/
-theorem bindK_single_comp (μ : Distribution A) (φ : A → B) :
+lemma bindK_single_comp (μ : Distribution A) (φ : A → B) :
     bindK μ (fun a => Finsupp.single (φ a) 1) = fTransform φ μ := by
   -- Induction on `μ`: both sides are additive, and move a point mass from `a` to `φ a`.
   induction μ using Finsupp.induction_linear with
@@ -181,7 +181,7 @@ theorem bindK_single_comp (μ : Distribution A) (φ : A → B) :
   | single a w => simp [fTransform]
 
 /-- **Pushing a mixture forward** pushes each law of the kernel forward. -/
-theorem fTransform_bindK (φ : B → C) (μ : Distribution A) (f : A → Distribution B) :
+lemma fTransform_bindK (φ : B → C) (μ : Distribution A) (f : A → Distribution B) :
     fTransform φ (bindK μ f) = bindK μ fun a => fTransform φ (f a) := by
   -- Induction on `μ`: the pushforward is linear.
   induction μ using Finsupp.induction_linear with
@@ -194,7 +194,7 @@ theorem fTransform_bindK (φ : B → C) (μ : Distribution A) (f : A → Distrib
     exact Finsupp.mapDomain_smul _ _
 
 /-- **Mixing over a pushforward** is mixing the composite kernel. -/
-theorem bindK_fTransform (φ : A → B) (μ : Distribution A) (f : B → Distribution C) :
+lemma bindK_fTransform (φ : A → B) (μ : Distribution A) (f : B → Distribution C) :
     bindK (fTransform φ μ) f = bindK μ (f ∘ φ) := by
   -- Induction on `μ`: a point mass at `a` is moved to `φ a`, where the kernel is `f (φ a)`.
   induction μ using Finsupp.induction_linear with
@@ -205,14 +205,14 @@ theorem bindK_fTransform (φ : A → B) (μ : Distribution A) (f : B → Distrib
   | single a w => simp [fTransform]
 
 /-- A pushforward depends only on the function on the support. -/
-theorem fTransform_congr_support {f g : A → B} {μ : Distribution A}
+lemma fTransform_congr_support {f g : A → B} {μ : Distribution A}
     (h : ∀ a ∈ μ.support, f a = g a) : fTransform f μ = fTransform g μ := by
   -- Both are point kernels, which agree on the support.
   rw [← bindK_single_comp, ← bindK_single_comp]
   exact bindK_congr fun a ha => by rw [h a ha]
 
 /-- **Restricting a mixture** restricts each law of the kernel. -/
-theorem restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribution B) :
+lemma restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribution B) :
     (bindK μ f).restrict P = bindK μ fun a => (f a).restrict P := by
   -- Induction on `μ`: restriction is linear.
   induction μ using Finsupp.induction_linear with
@@ -222,7 +222,7 @@ theorem restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distrib
   | single a w => simp [restrict, Finsupp.filter_smul]
 
 /-- A kernel that vanishes off an event is the kernel on the restriction to it. -/
-theorem bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution B) :
+lemma bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution B) :
     bindK μ (fun a => if P a then f a else 0) = bindK (μ.restrict P) f := by
   -- Induction on `μ`: a point mass at `a` contributes `f a` exactly when `P a`.
   induction μ using Finsupp.induction_linear with
@@ -235,7 +235,7 @@ theorem bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution
     · simp [h, restrict, Finsupp.filter_single_of_neg _ h]
 
 /-- **The weight of a mixture**: each weight of the kernel, weighted by `μ`. -/
-theorem weight_bindK (μ : Distribution A) (f : A → Distribution B) :
+lemma weight_bindK (μ : Distribution A) (f : A → Distribution B) :
     (bindK μ f).weight = μ.sum fun a w => w * (f a).weight := by
   -- Induction on `μ`: the weight is additive, and scales with a point mass.
   induction μ using Finsupp.induction_linear with
@@ -246,7 +246,7 @@ theorem weight_bindK (μ : Distribution A) (f : A → Distribution B) :
   | single a w => simp [weight_smul]
 
 /-- **A kernel of weight one preserves the weight.** -/
-theorem weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B}
+lemma weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B}
     (hf : ∀ a ∈ μ.support, (f a).weight = 1) : (bindK μ f).weight = μ.weight := by
   -- Each `μ a` is weighted by the weight of `f a`, which is one.
   rw [weight_bindK, weight]
@@ -255,7 +255,7 @@ theorem weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B}
     rw [hf a ha, mul_one]
 
 /-- **A nonnegative kernel mixed over a nonnegative distribution** is nonnegative. -/
-theorem bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.NonNeg)
+lemma bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.NonNeg)
     (hf : ∀ a, (f a).NonNeg) : (bindK μ f).NonNeg := by
   -- At each `b`, a sum of products of nonnegative numbers.
   intro b
@@ -265,7 +265,7 @@ theorem bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.
     simpa [Finsupp.smul_apply] using mul_nonneg (hμ a) (hf a b)
 
 /-- A possible outcome of a mixture is possible under the kernel at a possible first outcome. -/
-theorem mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b : B}
+lemma mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b : B}
     (hb : b ∈ (bindK μ f).support) : ∃ a ∈ μ.support, b ∈ (f a).support := by
   -- Otherwise every term `μ a · f a b` of the mixture at `b` vanishes.
   by_contra hn
@@ -279,7 +279,7 @@ theorem mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b :
 /-! ### Independent products -/
 
 /-- **An independent product as a mixture**: sample the first factor, then the second beside it. -/
-theorem prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
+lemma prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
     prod μ ν = bindK μ fun a => fTransform (fun b => (a, b)) ν := by
   -- Compare at a pair `(a, b)`: `μ a · ν b` on the left.
   ext ⟨a, b⟩
@@ -306,7 +306,7 @@ theorem prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
 
 /-- **Disintegration of an independent product** at one coordinate: sample that coordinate first,
 then the product with that coordinate fixed. -/
-theorem pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι → Distribution A)
+lemma pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι → Distribution A)
     (n : ι) : pi μ = bindK (μ n) fun z => pi (Function.update μ n (Finsupp.single z 1)) := by
   -- Compare at a sample `t`: `∏ᵢ μᵢ(tᵢ)` on the left.
   ext t
@@ -329,7 +329,7 @@ theorem pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι 
   · rw [Finsupp.notMem_support_iff.mp h, zero_mul]
 
 /-- With one coordinate fixed to `z`, every possible sample has `z` there. -/
-theorem eq_of_mem_support_pi_update {ι : Type*} [Fintype ι] [DecidableEq ι]
+lemma eq_of_mem_support_pi_update {ι : Type*} [Fintype ι] [DecidableEq ι]
     (μ : ι → Distribution A) (n : ι) (z : A) {t : ι → A}
     (ht : t ∈ (pi (Function.update μ n (Finsupp.single z 1))).support) : t n = z := by
   -- Otherwise the factor at `n`, the point mass at `z`, vanishes at `t n`.
@@ -347,10 +347,10 @@ namespace ProbDist
 noncomputable def map (f : A → B) (P : ProbDist A) : ProbDist B :=
   ⟨fTransform f P.1, fTransform_isProbDist f P.2⟩
 
-@[simp] theorem map_val (f : A → B) (P : ProbDist A) : (map f P).1 = fTransform f P.1 := rfl
+@[simp] lemma map_val (f : A → B) (P : ProbDist A) : (map f P).1 = fTransform f P.1 := rfl
 
 /-- **Pushing forward twice** is pushing forward along the composite. -/
-theorem map_map (f : B → C) (g : A → B) (P : ProbDist A) :
+lemma map_map (f : B → C) (g : A → B) (P : ProbDist A) :
     map f (map g P) = map (f ∘ g) P :=
   -- Both sides have the same underlying distribution: `fTransform f ∘ fTransform g`.
   Subtype.ext (fTransform_fTransform _ _ _)
@@ -361,7 +361,7 @@ noncomputable def bind (P : ProbDist A) (f : A → ProbDist B) : ProbDist B :=
     bindK_nonNeg P.2.1 fun a => (f a).2.1,
     (weight_bindK_of_weight fun a _ => (f a).2.2).trans P.2.2⟩
 
-@[simp] theorem bind_val (P : ProbDist A) (f : A → ProbDist B) :
+@[simp] lemma bind_val (P : ProbDist A) (f : A → ProbDist B) :
     (bind P f).1 = bindK P.1 fun a => (f a).1 := rfl
 
 end ProbDist

From f272bb102e18575ba790bec04f1cd2246aced7af Mon Sep 17 00:00:00 2001
From: Marc Ilunga <marc.ilunga@trailofbits.com>
Date: Tue, 13 Oct 2026 11:39:49 +0100
Subject: [PATCH 3/4] Kernels: review fixes

The restriction and pushforward facts that do not involve kernels
(`support_restrict_subset`, `restrict_restrict`, `restrict_congr`,
`restrict_false`, `single_restrict`, `fTransform_single`,
`restrict_fTransform`) move to Distribution.lean, next to `restrict` and
`fTransform`. Kernel.lean no longer opens `Classical`: `bindK_ite` states its
condition under `open scoped Classical in`, and the four proofs that decide
equalities or events start with `classical`. Two `show` steps become direct
terms.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01E7GPy1oUx164NSZsxJT3N9
---
 THEORY.md                         | 20 ++++-----
 src/Probability/Distribution.lean | 63 ++++++++++++++++++++++++++
 src/Probability/Kernel.lean       | 74 +++----------------------------
 3 files changed, 80 insertions(+), 77 deletions(-)

diff --git a/THEORY.md b/THEORY.md
index 50b8537..8bf93e3 100644
--- a/THEORY.md
+++ b/THEORY.md
@@ -1293,16 +1293,16 @@ agreeing with a map on a finite set $`S`$: there are $`\lvert Y \rvert^{\lvert X
 [ProbDist]: src/Probability/Distribution.lean#L149
 [isProbDist]: src/Probability/Distribution.lean#L128
 [mass]: src/Probability/Distribution.lean#L217
-[fTransform]: src/Probability/Distribution.lean#L587
-[prod]: src/Probability/Distribution.lean#L942
-[pi]: src/Probability/Distribution.lean#L1583
-[bindK]: src/Probability/Kernel.lean#L110
-[bindK_bindK]: src/Probability/Kernel.lean#L156
-[bindK_single_comp]: src/Probability/Kernel.lean#L173
-[pi_eq_bindK_update]: src/Probability/Kernel.lean#L309
-[ProbDist.map]: src/Probability/Kernel.lean#L347
-[ProbDist.bind]: src/Probability/Kernel.lean#L359
-[uniform]: src/Probability/Distribution.lean#L521
+[fTransform]: src/Probability/Distribution.lean#L632
+[prod]: src/Probability/Distribution.lean#L1005
+[pi]: src/Probability/Distribution.lean#L1646
+[bindK]: src/Probability/Kernel.lean#L49
+[bindK_bindK]: src/Probability/Kernel.lean#L93
+[bindK_single_comp]: src/Probability/Kernel.lean#L110
+[pi_eq_bindK_update]: src/Probability/Kernel.lean#L247
+[ProbDist.map]: src/Probability/Kernel.lean#L287
+[ProbDist.bind]: src/Probability/Kernel.lean#L299
+[uniform]: src/Probability/Distribution.lean#L566
 [statDist]: src/Probability/StatisticalDistance.lean#L88
 [statDist_triangle]: src/Probability/StatisticalDistance.lean#L243
 [statDist_fTransform_le]: src/Probability/StatisticalDistance.lean#L615
diff --git a/src/Probability/Distribution.lean b/src/Probability/Distribution.lean
index 1c74683..ede7487 100644
--- a/src/Probability/Distribution.lean
+++ b/src/Probability/Distribution.lean
@@ -280,6 +280,51 @@ theorem mass_restrict {A : Type*} (X : Distribution A) (P Q : A → Prop) :
   intro a _ha
   by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
 
+/-- Restricting keeps only mass that was there. -/
+lemma support_restrict_subset {A : Type*} (μ : Distribution A) (P : A → Prop) :
+    (μ.restrict P).support ⊆ μ.support := fun a ha => by
+  -- The restriction at `a` is `μ a` or zero; it is nonzero, so it is `μ a`.
+  rw [Finsupp.mem_support_iff, restrict_apply] at ha
+  rw [Finsupp.mem_support_iff]
+  split_ifs at ha
+  · exact ha
+  · exact absurd rfl ha
+
+
+/-- Restricting twice is restricting to both events. -/
+lemma restrict_restrict {A : Type*} (μ : Distribution A) (P Q : A → Prop) :
+    (μ.restrict P).restrict Q = μ.restrict fun a => P a ∧ Q a := by
+  -- Pointwise: `μ a` exactly when both events hold.
+  ext a
+  simp only [restrict_apply]
+  by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
+
+
+/-- Restriction depends only on the event up to equivalence. -/
+lemma restrict_congr {A : Type*} (μ : Distribution A) {P Q : A → Prop} (h : ∀ a, P a ↔ Q a) :
+    μ.restrict P = μ.restrict Q := by
+  -- Pointwise, the two events are the same condition.
+  ext a
+  simp only [restrict_apply, h]
+
+
+/-- Restricting to the impossible event leaves nothing. -/
+lemma restrict_false {A : Type*} (μ : Distribution A) : μ.restrict (fun _ => False) = 0 := by
+  -- Pointwise, the condition never holds.
+  ext a
+  simp
+
+
+/-- A point mass restricted to an event: itself or nothing. -/
+lemma single_restrict {A : Type*} (a : A) (w : ℝ) (P : A → Prop) :
+    Distribution.restrict (Finsupp.single a w : Distribution A) P =
+      if P a then Finsupp.single a w else 0 := by
+  -- The only point with mass is `a`: kept if `P a`, dropped otherwise.
+  by_cases h : P a
+  · simp [restrict, Finsupp.filter_single_of_pos _ h, h]
+  · simp [restrict, Finsupp.filter_single_of_neg _ h, h]
+
+
 theorem mass_true {A : Type*} (X : Distribution A) :
     X.mass (fun _ => True) = X.weight := by
   unfold mass weight
@@ -651,6 +696,24 @@ theorem mass_fTransform {A B : Type*}
     (fun b => by by_cases hb : P b <;> simp [hb])
     (fun b m₁ m₂ => by by_cases hb : P b <;> simp [hb])]
 
+/-- The pushforward of a point mass. -/
+lemma fTransform_single {A B : Type*} (f : A → B) (a : A) (w : ℝ) :
+    fTransform f (Finsupp.single a w) = Finsupp.single (f a) w :=
+  Finsupp.mapDomain_single
+
+
+/-- Restricting a pushforward is pushing forward the restriction to the preimage. -/
+lemma restrict_fTransform {A B : Type*} (f : A → B) (μ : Distribution A) (P : B → Prop) :
+    (fTransform f μ).restrict P = fTransform f (μ.restrict fun a => P (f a)) := by
+  -- At `b`, both sides are masses of fibres of `f`.
+  ext b
+  rw [restrict_apply, fTransform_apply_eq_mass, fTransform_apply_eq_mass, mass_restrict]
+  split_ifs with h
+  · -- `P b`: on the fibre of `b`, the condition `P (f a)` holds
+    exact mass_congr _ fun a => ⟨fun ha => ⟨ha, ha ▸ h⟩, fun ha => ha.1⟩
+  · -- not `P b`: the fibre of `b` meets no point with `P (f a)`
+    exact (mass_eq_zero_of_forall_not _ fun a (ha : f a = b ∧ P (f a)) => h (ha.1 ▸ ha.2)).symm
+
 /-- Pushforward preserves total mass. -/
 theorem weight_fTransform {A B : Type*}
     (f : A → B) (X : Distribution A) :
diff --git a/src/Probability/Kernel.lean b/src/Probability/Kernel.lean
index b04e093..f7ea347 100644
--- a/src/Probability/Kernel.lean
+++ b/src/Probability/Kernel.lean
@@ -41,69 +41,8 @@ after the other, and kernels of probability laws mix probability laws into proba
 
 namespace Probability.Distribution
 
-open Classical
-
 variable {A B C : Type*}
 
-/-! ### Restriction and pushforward -/
-
-/-- Restricting keeps only mass that was there. -/
-lemma support_restrict_subset (μ : Distribution A) (P : A → Prop) :
-    (μ.restrict P).support ⊆ μ.support := fun a ha => by
-  -- The restriction at `a` is `μ a` or zero; it is nonzero, so it is `μ a`.
-  rw [Finsupp.mem_support_iff, restrict_apply] at ha
-  rw [Finsupp.mem_support_iff]
-  split_ifs at ha
-  · exact ha
-  · exact absurd rfl ha
-
-/-- Restricting twice is restricting to both events. -/
-lemma restrict_restrict (μ : Distribution A) (P Q : A → Prop) :
-    (μ.restrict P).restrict Q = μ.restrict fun a => P a ∧ Q a := by
-  -- Pointwise: `μ a` exactly when both events hold.
-  ext a
-  simp only [restrict_apply]
-  by_cases hP : P a <;> by_cases hQ : Q a <;> simp [hP, hQ]
-
-/-- Restriction depends only on the event up to equivalence. -/
-lemma restrict_congr (μ : Distribution A) {P Q : A → Prop} (h : ∀ a, P a ↔ Q a) :
-    μ.restrict P = μ.restrict Q := by
-  -- Pointwise, the two events are the same condition.
-  ext a
-  simp only [restrict_apply, h]
-
-/-- Restricting to the impossible event leaves nothing. -/
-lemma restrict_false (μ : Distribution A) : μ.restrict (fun _ => False) = 0 := by
-  -- Pointwise, the condition never holds.
-  ext a
-  simp
-
-/-- A point mass restricted to an event: itself or nothing. -/
-lemma single_restrict (a : A) (w : ℝ) (P : A → Prop) :
-    Distribution.restrict (Finsupp.single a w : Distribution A) P =
-      if P a then Finsupp.single a w else 0 := by
-  -- The only point with mass is `a`: kept if `P a`, dropped otherwise.
-  by_cases h : P a
-  · simp [restrict, Finsupp.filter_single_of_pos _ h, h]
-  · simp [restrict, Finsupp.filter_single_of_neg _ h, h]
-
-/-- The pushforward of a point mass. -/
-lemma fTransform_single (f : A → B) (a : A) (w : ℝ) :
-    fTransform f (Finsupp.single a w) = Finsupp.single (f a) w :=
-  Finsupp.mapDomain_single
-
-/-- Restricting a pushforward is pushing forward the restriction to the preimage. -/
-lemma restrict_fTransform (f : A → B) (μ : Distribution A) (P : B → Prop) :
-    (fTransform f μ).restrict P = fTransform f (μ.restrict fun a => P (f a)) := by
-  -- At `b`, both sides are masses of fibres of `f`.
-  ext b
-  rw [restrict_apply, fTransform_apply_eq_mass, fTransform_apply_eq_mass, mass_restrict]
-  split_ifs with h
-  · -- `P b`: on the fibre of `b`, the condition `P (f a)` holds
-    exact mass_congr _ fun a => ⟨fun ha => ⟨ha, ha ▸ h⟩, fun ha => ha.1⟩
-  · -- not `P b`: the fibre of `b` meets no point with `P (f a)`
-    exact (mass_eq_zero_of_forall_not _ fun a (ha : f a = b ∧ P (f a)) => h (ha.1 ▸ ha.2)).symm
-
 /-! ### Kernels -/
 
 /-- **The mixture of a kernel** over a distribution: `∑ a, μ a • f a`. -/
@@ -147,9 +86,7 @@ lemma bindK_congr {μ : Distribution A} {f g : A → Distribution B}
   -- Both are sums over the support of `μ`, term by term equal.
   unfold bindK
   rw [Finsupp.linearCombination_apply, Finsupp.linearCombination_apply]
-  exact Finset.sum_congr rfl fun a ha => by
-    show μ a • f a = μ a • g a
-    rw [h a ha]
+  exact Finset.sum_congr rfl fun a ha => congrArg (μ a • ·) (h a ha)
 
 /-- **Mixtures compose**: mixing `g` over the mixture of `f` is mixing, over `μ`, the mixtures of
 `g` over each `f a`. -/
@@ -214,6 +151,7 @@ lemma fTransform_congr_support {f g : A → B} {μ : Distribution A}
 /-- **Restricting a mixture** restricts each law of the kernel. -/
 lemma restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribution B) :
     (bindK μ f).restrict P = bindK μ fun a => (f a).restrict P := by
+  classical
   -- Induction on `μ`: restriction is linear.
   induction μ using Finsupp.induction_linear with
   | zero => simp [restrict, Finsupp.filter_zero]
@@ -221,6 +159,7 @@ lemma restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribut
       ← restrict, hμ, hν]
   | single a w => simp [restrict, Finsupp.filter_smul]
 
+open scoped Classical in
 /-- A kernel that vanishes off an event is the kernel on the restriction to it. -/
 lemma bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution B) :
     bindK μ (fun a => if P a then f a else 0) = bindK (μ.restrict P) f := by
@@ -250,9 +189,7 @@ lemma weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B}
     (hf : ∀ a ∈ μ.support, (f a).weight = 1) : (bindK μ f).weight = μ.weight := by
   -- Each `μ a` is weighted by the weight of `f a`, which is one.
   rw [weight_bindK, weight]
-  exact Finset.sum_congr rfl fun a ha => by
-    show μ a * (f a).weight = μ a
-    rw [hf a ha, mul_one]
+  exact Finset.sum_congr rfl fun a ha => by simp only [hf a ha, mul_one]
 
 /-- **A nonnegative kernel mixed over a nonnegative distribution** is nonnegative. -/
 lemma bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.NonNeg)
@@ -281,6 +218,7 @@ lemma mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b : B}
 /-- **An independent product as a mixture**: sample the first factor, then the second beside it. -/
 lemma prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
     prod μ ν = bindK μ fun a => fTransform (fun b => (a, b)) ν := by
+  classical
   -- Compare at a pair `(a, b)`: `μ a · ν b` on the left.
   ext ⟨a, b⟩
   rw [prod_apply, bindK_apply, Finsupp.sum]
@@ -308,6 +246,7 @@ lemma prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
 then the product with that coordinate fixed. -/
 lemma pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι → Distribution A)
     (n : ι) : pi μ = bindK (μ n) fun z => pi (Function.update μ n (Finsupp.single z 1)) := by
+  classical
   -- Compare at a sample `t`: `∏ᵢ μᵢ(tᵢ)` on the left.
   ext t
   -- With coordinate `n` fixed to `z`, the product is `[t n = z] · ∏_{i ≠ n} μᵢ(tᵢ)`.
@@ -332,6 +271,7 @@ lemma pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι →
 lemma eq_of_mem_support_pi_update {ι : Type*} [Fintype ι] [DecidableEq ι]
     (μ : ι → Distribution A) (n : ι) (z : A) {t : ι → A}
     (ht : t ∈ (pi (Function.update μ n (Finsupp.single z 1))).support) : t n = z := by
+  classical
   -- Otherwise the factor at `n`, the point mass at `z`, vanishes at `t n`.
   rw [Finsupp.mem_support_iff, pi_apply] at ht
   by_contra hne

From f64df8ec7de0b41e1f71bbfb8a4507e7ff4e324a Mon Sep 17 00:00:00 2001
From: Marc Ilunga <marc.ilunga@trailofbits.com>
Date: Tue, 13 Oct 2026 12:00:33 +0100
Subject: [PATCH 4/4] Kernels: proof comments

Each induction on a law says what both sides are at a sum of two laws and
at a point mass; the product lemmas comment each step (the factor split off,
the terms that survive the sum, the support cases).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01E7GPy1oUx164NSZsxJT3N9
---
 THEORY.md                   | 10 ++++-----
 src/Probability/Kernel.lean | 44 +++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/THEORY.md b/THEORY.md
index 8bf93e3..448c9b1 100644
--- a/THEORY.md
+++ b/THEORY.md
@@ -1297,11 +1297,11 @@ agreeing with a map on a finite set $`S`$: there are $`\lvert Y \rvert^{\lvert X
 [prod]: src/Probability/Distribution.lean#L1005
 [pi]: src/Probability/Distribution.lean#L1646
 [bindK]: src/Probability/Kernel.lean#L49
-[bindK_bindK]: src/Probability/Kernel.lean#L93
-[bindK_single_comp]: src/Probability/Kernel.lean#L110
-[pi_eq_bindK_update]: src/Probability/Kernel.lean#L247
-[ProbDist.map]: src/Probability/Kernel.lean#L287
-[ProbDist.bind]: src/Probability/Kernel.lean#L299
+[bindK_bindK]: src/Probability/Kernel.lean#L95
+[bindK_single_comp]: src/Probability/Kernel.lean#L117
+[pi_eq_bindK_update]: src/Probability/Kernel.lean#L283
+[ProbDist.map]: src/Probability/Kernel.lean#L329
+[ProbDist.bind]: src/Probability/Kernel.lean#L341
 [uniform]: src/Probability/Distribution.lean#L566
 [statDist]: src/Probability/StatisticalDistance.lean#L88
 [statDist_triangle]: src/Probability/StatisticalDistance.lean#L243
diff --git a/src/Probability/Kernel.lean b/src/Probability/Kernel.lean
index f7ea347..3177344 100644
--- a/src/Probability/Kernel.lean
+++ b/src/Probability/Kernel.lean
@@ -77,7 +77,9 @@ lemma bindK_apply (μ : Distribution A) (f : A → Distribution B) (b : B) :
   -- Induction on `μ` as a sum of point masses.
   induction μ using Finsupp.induction_linear with
   | zero => simp
+  -- A sum of two laws: the mixture is additive, and `0 + 0 = 0`.
   | add μ ν hμ hν => rw [bindK_add, hμ, hν, add_zero]
+  -- A point mass of weight `w`: it mixes to `w • 0 = 0`.
   | single a w => simp
 
 /-- A mixture depends only on the kernel on the support. -/
@@ -95,7 +97,10 @@ lemma bindK_bindK (μ : Distribution A) (f : A → Distribution B) (g : B → Di
   -- Induction on `μ`: both sides are additive in `μ`, and agree on point masses.
   induction μ using Finsupp.induction_linear with
   | zero => simp
+  -- A sum of two laws: every mixture involved is additive in its first law.
   | add μ ν hμ hν => rw [bindK_add, bindK_add, bindK_add, hμ, hν]
+  -- A point mass of weight `w` at `a`: both sides are `w • bindK (f a) g`, as the mixture is
+  -- linear in its first law.
   | single a w => simp [bindK]
 
 /-- **The point kernel** leaves a distribution unchanged. -/
@@ -103,7 +108,9 @@ lemma bindK_single_one (μ : Distribution A) : bindK μ (fun a => Finsupp.single
   -- Induction on `μ`: a point mass of weight `w` mixes to `w` times the point mass.
   induction μ using Finsupp.induction_linear with
   | zero => simp
+  -- A sum of two laws: the mixture is additive.
   | add μ ν hμ hν => rw [bindK_add, hμ, hν]
+  -- A point mass of weight `w` at `a`: `w • single a 1 = single a w`.
   | single a w => simp
 
 /-- **A point kernel along `φ` is the pushforward** along `φ`. -/
@@ -113,8 +120,11 @@ lemma bindK_single_comp (μ : Distribution A) (φ : A → B) :
   induction μ using Finsupp.induction_linear with
   | zero => simp [fTransform]
   | add μ ν hμ hν =>
+    -- the mixture is additive; then the induction hypotheses ...
     rw [bindK_add, hμ, hν]
+    -- ... and the pushforward of a sum is the sum of the pushforwards
     exact (Finsupp.mapDomain_add).symm
+  -- A point mass of weight `w` at `a`: both sides are the point mass of weight `w` at `φ a`.
   | single a w => simp [fTransform]
 
 /-- **Pushing a mixture forward** pushes each law of the kernel forward. -/
@@ -124,10 +134,14 @@ lemma fTransform_bindK (φ : B → C) (μ : Distribution A) (f : A → Distribut
   induction μ using Finsupp.induction_linear with
   | zero => simp [fTransform]
   | add μ ν hμ hν =>
+    -- both mixtures split over `μ + ν`; the right one is folded back by the hypotheses ...
     rw [bindK_add, bindK_add, ← hμ, ← hν]
+    -- ... leaving the additivity of the pushforward
     exact Finsupp.mapDomain_add
   | single a w =>
+    -- a point mass of weight `w` at `a` mixes to `w • f a` and `w • fTransform φ (f a)` ...
     rw [bindK_single, bindK_single]
+    -- ... and the pushforward commutes with scaling by `w`
     exact Finsupp.mapDomain_smul _ _
 
 /-- **Mixing over a pushforward** is mixing the composite kernel. -/
@@ -137,8 +151,10 @@ lemma bindK_fTransform (φ : A → B) (μ : Distribution A) (f : B → Distribut
   induction μ using Finsupp.induction_linear with
   | zero => simp [fTransform]
   | add μ ν hμ hν =>
+    -- the pushforward and both mixtures are additive; then the induction hypotheses
     rw [show fTransform φ (μ + ν) = fTransform φ μ + fTransform φ ν from Finsupp.mapDomain_add,
       bindK_add, bindK_add, hμ, hν]
+  -- A point mass of weight `w` at `a`: both sides are `w • f (φ a)`.
   | single a w => simp [fTransform]
 
 /-- A pushforward depends only on the function on the support. -/
@@ -155,8 +171,10 @@ lemma restrict_bindK (P : B → Prop) (μ : Distribution A) (f : A → Distribut
   -- Induction on `μ`: restriction is linear.
   induction μ using Finsupp.induction_linear with
   | zero => simp [restrict, Finsupp.filter_zero]
+  -- A sum of two laws: the mixtures and the restriction (a filter) are all additive.
   | add μ ν hμ hν => rw [bindK_add, bindK_add, restrict, Finsupp.filter_add, ← restrict,
       ← restrict, hμ, hν]
+  -- A point mass of weight `w` at `a`: restricting `w • f a` is `w •` the restriction of `f a`.
   | single a w => simp [restrict, Finsupp.filter_smul]
 
 open scoped Classical in
@@ -166,11 +184,16 @@ lemma bindK_ite (μ : Distribution A) (P : A → Prop) (f : A → Distribution B
   -- Induction on `μ`: a point mass at `a` contributes `f a` exactly when `P a`.
   induction μ using Finsupp.induction_linear with
   | zero => simp [restrict, Finsupp.filter_zero]
+  -- A sum of two laws: the mixtures are additive, and the restriction of a sum is the sum of
+  -- the restrictions.
   | add μ ν hμ hν => rw [bindK_add, hμ, hν, ← bindK_add, restrict, restrict, restrict,
       Finsupp.filter_add]
   | single a w =>
+    -- a point mass of weight `w` at `a`: case on the event at `a`
     by_cases h : P a
+    -- `P a`: the restriction keeps the point mass, and both sides are `w • f a`
     · simp [h, restrict, Finsupp.filter_single_of_pos _ h]
+    -- not `P a`: the restriction drops the point mass, and both sides are zero
     · simp [h, restrict, Finsupp.filter_single_of_neg _ h]
 
 /-- **The weight of a mixture**: each weight of the kernel, weighted by `μ`. -/
@@ -180,8 +203,11 @@ lemma weight_bindK (μ : Distribution A) (f : A → Distribution B) :
   induction μ using Finsupp.induction_linear with
   | zero => simp [weight]
   | add μ ν hμ hν =>
+    -- the mixture and its weight are additive; the sum over `μ + ν` splits, as its summand
+    -- `w * (f a).weight` is additive in `w`
     rw [bindK_add, weight_add, hμ, hν,
       Finsupp.sum_add_index' (by simp) (fun _ _ _ => add_mul _ _ _)]
+  -- A point mass of weight `w` at `a`: the weight of `w • f a` is `w` times that of `f a`.
   | single a w => simp [weight_smul]
 
 /-- **A kernel of weight one preserves the weight.** -/
@@ -189,6 +215,7 @@ lemma weight_bindK_of_weight {μ : Distribution A} {f : A → Distribution B}
     (hf : ∀ a ∈ μ.support, (f a).weight = 1) : (bindK μ f).weight = μ.weight := by
   -- Each `μ a` is weighted by the weight of `f a`, which is one.
   rw [weight_bindK, weight]
+  -- Term by term over the support of `μ`: `μ(a) · 1 = μ(a)`.
   exact Finset.sum_congr rfl fun a ha => by simp only [hf a ha, mul_one]
 
 /-- **A nonnegative kernel mixed over a nonnegative distribution** is nonnegative. -/
@@ -196,8 +223,10 @@ lemma bindK_nonNeg {μ : Distribution A} {f : A → Distribution B} (hμ : μ.No
     (hf : ∀ a, (f a).NonNeg) : (bindK μ f).NonNeg := by
   -- At each `b`, a sum of products of nonnegative numbers.
   intro b
+  -- The mixture at `b` is `∑ₐ μ(a) · f(a)(b)`.
   unfold bindK
   rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply]
+  -- Each term is `μ(a) · f(a)(b)`, with `μ(a) ≥ 0` and `f(a)(b) ≥ 0`.
   exact Finset.sum_nonneg fun a _ => by
     simpa [Finsupp.smul_apply] using mul_nonneg (hμ a) (hf a b)
 
@@ -207,9 +236,12 @@ lemma mem_support_bindK {μ : Distribution A} {f : A → Distribution B} {b : B}
   -- Otherwise every term `μ a · f a b` of the mixture at `b` vanishes.
   by_contra hn
   push Not at hn
+  -- So `f(a)(b) = 0` at every possible `a`; we show the mixture at `b` is zero, against `hb`.
   apply Finsupp.mem_support_iff.mp hb
+  -- The mixture at `b` is `∑ₐ μ(a) · f(a)(b)` over the support of `μ` ...
   unfold bindK
   rw [Finsupp.linearCombination_apply, Finsupp.sum, Finsupp.finsetSum_apply]
+  -- ... and every term vanishes, since `f(a)(b) = 0`.
   exact Finset.sum_eq_zero fun a ha => by
     simp [Finsupp.notMem_support_iff.mp (hn a ha)]
 
@@ -225,8 +257,10 @@ lemma prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
   -- The law `ν` placed beside `a'` puts `ν b` at `(a, b)` if `a' = a`, and nothing otherwise.
   have hpt : ∀ a', fTransform (fun b' => (a', b')) ν (a, b) = if a' = a then ν b else 0 := by
     intro a'
+    -- The pushforward at `(a, b)` is the `ν`-mass of the event `(a', b') = (a, b)`.
     rw [fTransform_apply_eq_mass]
     split_ifs with h
+    -- `a' = a`: the event is `b' = b`, of mass `ν b` (zero when `b` is outside the support).
     · subst h
       rw [mass, Finsupp.sum]
       simp only [Prod.mk.injEq, true_and]
@@ -234,10 +268,12 @@ lemma prod_eq_bindK (μ : Distribution A) (ν : Distribution B) :
       split_ifs with hb
       · rfl
       · exact (Finsupp.notMem_support_iff.mp hb).symm
+    -- `a' ≠ a`: no pair `(a', b')` is `(a, b)`, so the event is empty.
     · exact mass_eq_zero_of_forall_not _ fun b' h' => h (Prod.mk.inj h').1
   -- So the mixture at `(a, b)` keeps only the term `a' = a`: `μ a · ν b`.
   simp only [hpt, mul_ite, mul_zero]
   rw [Finset.sum_ite_eq']
+  -- The sum keeps that term if `a` is in the support of `μ`; otherwise `μ a = 0`.
   split_ifs with ha
   · rfl
   · rw [Finsupp.notMem_support_iff.mp ha, zero_mul]
@@ -253,16 +289,21 @@ lemma pi_eq_bindK_update {ι : Type*} [Fintype ι] [DecidableEq ι] (μ : ι →
   have hrest : ∀ z : A, (∏ i, Function.update μ n (Finsupp.single z 1) i (t i)) =
       (Finsupp.single z (1 : ℝ)) (t n) * ∏ i ∈ Finset.univ.erase n, μ i (t i) := by
     intro z
+    -- Split off the factor at `n`, which the update sets to the point mass at `z` ...
     rw [← Finset.mul_prod_erase Finset.univ _ (Finset.mem_univ n), Function.update_self]
     congr 1
+    -- ... and off `n` the update leaves the factors `μ i` unchanged.
     exact Finset.prod_congr rfl fun i hi => by
       rw [Function.update_of_ne (Finset.ne_of_mem_erase hi)]
   -- Mixing over `z ← μ n` keeps the term `z = t n`: `μ n (t n) · ∏_{i ≠ n} μᵢ(tᵢ)`.
   rw [pi_apply, bindK_apply]
   simp only [pi_apply, hrest, Finsupp.single_apply]
+  -- Split the factor at `n` off the left side too: `μ n (t n) · ∏_{i ≠ n} μᵢ(tᵢ)`.
   rw [← Finset.mul_prod_erase Finset.univ _ (Finset.mem_univ n), Finsupp.sum]
+  -- On the right, the term at `z` is `μ n z · ∏_{i ≠ n} μᵢ(tᵢ)` if `z = t n`, and zero otherwise.
   simp only [← mul_assoc, mul_ite, mul_one, mul_zero, ite_mul, zero_mul]
   rw [Finset.sum_ite_eq']
+  -- The sum keeps that term if `t n` is in the support of `μ n`; otherwise `μ n (t n) = 0`.
   split_ifs with h
   · rfl
   · rw [Finsupp.notMem_support_iff.mp h, zero_mul]
@@ -276,6 +317,7 @@ lemma eq_of_mem_support_pi_update {ι : Type*} [Fintype ι] [DecidableEq ι]
   rw [Finsupp.mem_support_iff, pi_apply] at ht
   by_contra hne
   apply ht
+  -- The product `∏ᵢ` vanishes: its factor at `n`, `single z 1 (t n)`, is zero as `t n ≠ z`.
   refine Finset.prod_eq_zero (Finset.mem_univ n) ?_
   rw [Function.update_self, Finsupp.single_apply, if_neg (Ne.symm hne)]
 
@@ -298,7 +340,9 @@ lemma map_map (f : B → C) (g : A → B) (P : ProbDist A) :
 /-- **Sequential composition**: the law of `b ← f a` for `a ← P`, the mixture of `f` over `P`. -/
 noncomputable def bind (P : ProbDist A) (f : A → ProbDist B) : ProbDist B :=
   ⟨bindK P.1 fun a => (f a).1,
+    -- nonnegative: a mixture of nonnegative laws over the nonnegative law `P`
     bindK_nonNeg P.2.1 fun a => (f a).2.1,
+    -- weight one: each `f a` has weight one, so the mixture keeps the weight of `P`, which is one
     (weight_bindK_of_weight fun a _ => (f a).2.2).trans P.2.2⟩
 
 @[simp] lemma bind_val (P : ProbDist A) (f : A → ProbDist B) :