Sub-distributions

This file defines sub-distributions as probability mass functions over Option α. A sub-distribution can have total mass less than 1, representing computations that may fail (return none).

Main definitions

• SDistr α - Sub-distribution over type α, defined as PMF (Option α)
• SDistr.pure - Point mass distribution
• SDistr.fail - Failed computation (mass 0 on all some values)
• SDistr.bind - Monadic bind for sub-distributions
• SDistr.uniform - Uniform distribution over a finite type

References

• SSProve: https://github.com/SSProve/ssprove
• EasyCrypt sub-distributions

@[reducible, inline]

abbrev CatCrypt.Prob.SDistr (α : Type u_1) : Type u_1

Sub-distribution: a distribution that may have total mass < 1. Represented as a PMF over Option α, where none represents failure.

Equations

• CatCrypt.Prob.SDistr α = PMF (Option α)

Basic constructors

noncomputable def CatCrypt.Prob.SDistr.pure {α : Type u_1} (a : α) : SDistr α

Point mass distribution at a

Equations

• CatCrypt.Prob.SDistr.pure a = PMF.pure (some a)

noncomputable def CatCrypt.Prob.SDistr.fail {α : Type u_1} : SDistr α

Failed computation: all mass on none

Equations

• CatCrypt.Prob.SDistr.fail = PMF.pure none

noncomputable def CatCrypt.Prob.SDistr.bind {α : Type u_1} {β : Type u_2} (d : SDistr α) (f : α → SDistr β) : SDistr β

Monadic bind for sub-distributions. If the first computation fails, the result fails. Otherwise, apply f to the result.

Equations

• d.bind f = PMF.bind d fun (oa : Option α) => match oa with | some a => f a | none => CatCrypt.Prob.SDistr.fail

@[implicit_reducible]

noncomputable instance CatCrypt.Prob.SDistr.instMonad : Monad SDistr

Equations

• One or more equations did not get rendered due to their size.

Uniform distribution

noncomputable def CatCrypt.Prob.SDistr.uniform (α : Type u_4) [Fintype α] [Nonempty α] : SDistr α

Uniform distribution over a finite nonempty type

Equations

• CatCrypt.Prob.SDistr.uniform α = PMF.map some (PMF.uniformOfFintype α)

Measures and mass

noncomputable def CatCrypt.Prob.SDistr.mass {α : Type u_1} (d : SDistr α) : ENNReal

Total mass: probability of not failing

Equations

• d.mass = 1 - d none

def CatCrypt.Prob.SDistr.support {α : Type u_1} (d : SDistr α) : Set α

Support of the distribution: values with nonzero probability

Equations

• d.support = {a : α | d (some a) ≠ 0}

Basic lemmas

@[simp]

theorem CatCrypt.Prob.SDistr.pure_apply_some {α : Type u_1} (a b : α) [DecidableEq α] : (pure a) (some b) = if a = b then 1 else 0

@[simp]

theorem CatCrypt.Prob.SDistr.pure_apply_none {α : Type u_1} (a : α) : (pure a) none = 0

@[simp]

theorem CatCrypt.Prob.SDistr.fail_apply_some {α : Type u_1} (a : α) : fail (some a) = 0

@[simp]

theorem CatCrypt.Prob.SDistr.fail_apply_none {α : Type u_1} : fail none = 1

@[simp]

theorem CatCrypt.Prob.SDistr.mass_pure {α : Type u_1} (a : α) : (pure a).mass = 1

@[simp]

theorem CatCrypt.Prob.SDistr.mass_fail {α : Type u_1} : fail.mass = 0

Support lemmas

theorem CatCrypt.Prob.SDistr.mem_support_iff {α : Type u_1} (d : SDistr α) (a : α) : a ∈ d.support ↔ d (some a) ≠ 0

theorem CatCrypt.Prob.SDistr.mem_support_pure_iff {α : Type u_1} (a b : α) : b ∈ (pure a).support ↔ a = b

Membership in support of pure implies equality (without DecidableEq)

theorem CatCrypt.Prob.SDistr.support_pure {α : Type u_1} [DecidableEq α] (a : α) : (pure a).support = {a}

theorem CatCrypt.Prob.SDistr.support_fail {α : Type u_1} : fail.support = ∅

theorem CatCrypt.Prob.SDistr.bind_fail {α : Type u_1} {β : Type u_2} (f : α → SDistr β) : fail.bind f = fail

Binding fail with any function gives fail

theorem CatCrypt.Prob.SDistr.pure_bind {α : Type u_1} {β : Type u_2} (a : α) (f : α → SDistr β) : (pure a).bind f = f a

Binding pure with a function gives the function applied to the value

theorem CatCrypt.Prob.SDistr.bind_pure {α : Type u_1} (d : SDistr α) : d.bind pure = d

Binding with pure is identity

theorem CatCrypt.Prob.SDistr.bind_assoc {α : Type u_1} {β : Type u_2} {γ : Type u_3} (d : SDistr α) (f : α → SDistr β) (g : β → SDistr γ) : (d.bind f).bind g = d.bind fun (a : α) => (f a).bind g

Bind is associative

Lawful Monad Instance

instance CatCrypt.Prob.SDistr.instLawfulMonad : LawfulMonad SDistr

SDistr is a lawful monad. The laws follow from pure_bind, bind_pure, bind_assoc.

Support-based congruence (eq_in_dlet in Rocq)

theorem CatCrypt.Prob.SDistr.bind_congr_support {α : Type u_1} {β : Type u_2} {d : SDistr α} {f g : α → SDistr β} (h : ∀ (a : α), d (some a) ≠ 0 → f a = g a) : d.bind f = d.bind g

If functions agree on the support of a distribution, their binds are equal. This is the Lean equivalent of eq_in_dlet from mathcomp-experimental-reals: If f and g agree on {a | d (some a) ≠ 0}, then d.bind f = d.bind g.

This is a key lemma for coupling composition proofs.

theorem CatCrypt.Prob.SDistr.bind_congr_support' {α : Type u_1} {β : Type u_2} {d₁ d₂ : SDistr α} {f g : α → SDistr β} (hd : d₁ = d₂) (h : ∀ (a : α), d₁ (some a) ≠ 0 → f a = g a) : d₁.bind f = d₂.bind g

Variant: congruence with explicit equality of distributions

Uniform distribution lemmas

theorem CatCrypt.Prob.SDistr.uniform_apply_some {α : Type u_1} [Fintype α] [Nonempty α] (a : α) : (uniform α) (some a) = (↑(Fintype.card α))⁻¹

theorem CatCrypt.Prob.SDistr.uniform_apply_none {α : Type u_1} [Fintype α] [Nonempty α] : (uniform α) none = 0

@[simp]

theorem CatCrypt.Prob.SDistr.mass_uniform {α : Type u_1} [Fintype α] [Nonempty α] : (uniform α).mass = 1

theorem CatCrypt.Prob.SDistr.uniform_bind_bij {α : Type u_1} {β : Type u_2} [Fintype α] [Fintype β] [Nonempty α] [Nonempty β] (f : α ≃ β) : ((uniform α).bind fun (a : α) => pure (f a)) = uniform β

Applying a bijection to a uniform distribution gives a uniform distribution

theorem CatCrypt.Prob.SDistr.uniform_bind_equiv_comp {α : Type u_1} {β : Type u_2} {γ : Type u_3} [Fintype α] [Fintype β] [Nonempty α] [Nonempty β] (σ : α ≃ β) (g : β → SDistr γ) : ((uniform α).bind fun (a : α) => g (σ a)) = (uniform β).bind g

Compositional version of uniform_bind_bij: applying a bijection inside a uniform bind doesn't change the distribution. (uniform α).bind (fun a => g (σ a)) = (uniform β).bind g for any bijection σ : α ≃ β.

theorem CatCrypt.Prob.SDistr.uniform_bind_const {α : Type u_1} {β : Type u_2} [Fintype α] [Nonempty α] (d : SDistr β) : ((uniform α).bind fun (x : α) => d) = d

Binding uniform distribution with a constant function gives the constant. This is because averaging a constant gives the constant.

Mathematical Proof:

• LHS = ∑_a (1/|α|) * d = |α| * (1/|α|) * d = d
• The uniform distribution has |α| elements each with probability 1/|α|
• Summing |α| copies of d weighted by 1/|α| gives d

This is a fundamental property needed for the averaging coupling argument in pRHL sampling rules.

theorem CatCrypt.Prob.SDistr.bind_comm {α : Type u_1} {β : Type u_2} {γ : Type u_3} (d₁ : SDistr α) (d₂ : SDistr β) (f : α → β → SDistr γ) : (d₁.bind fun (a : α) => d₂.bind (f a)) = d₂.bind fun (b : β) => d₁.bind fun (a : α) => f a b

Commutativity of bind for sub-distributions. Two independent bind operations can be swapped.

Nested uniform sampling equals product uniform sampling

theorem CatCrypt.Prob.SDistr.uniform_nested_eq_prod (A B : Type) [Fintype A] [Fintype B] [Nonempty A] [Nonempty B] : ((uniform A).bind fun (a : A) => (uniform B).bind fun (b : B) => pure (a, b)) = uniform (A × B)

Nested uniform sampling equals product uniform sampling.

For finite types A and B, sampling a : A and b : B independently and uniformly, then pairing them, gives the uniform distribution over A × B.

Mathematical Proof:

• LHS: Pr[(a₀, b₀)] = Pr[a = a₀] × Pr[b = b₀] = (1/|A|) × (1/|B|) = 1/(|A| × |B|)
• RHS: Pr[(a₀, b₀)] = 1/|A × B| = 1/(|A| × |B|)

These are equal, so the distributions are equal.

This is a key lemma for cryptographic proofs involving uniform sampling over product types.

Option tsum helpers

theorem CatCrypt.Prob.SDistr.tsum_option_none_zero {γ : Type u_4} (f : Option γ → ENNReal) (hf : f none = 0) : ∑' (x : Option γ), f x = ∑' (a : γ), f (some a)

Helper lemma: For a function on Option where the none value is 0, the sum over Option equals the sum over some.

theorem CatCrypt.Prob.SDistr.tsum_option_eq_add {γ : Type u_4} (f : Option γ → ENNReal) : ∑' (x : Option γ), f x = f none + ∑' (a : γ), f (some a)

Helper: sum over Option α can be split into the none term plus sum over some terms

theorem CatCrypt.Prob.SDistr.eq_of_some_eq {α : Type u_1} {d₁ d₂ : SDistr α} (h : ∀ (a : α), d₁ (some a) = d₂ (some a)) : d₁ = d₂

If two SDistrs agree on all some values, they are equal. This follows from PMF having total mass 1.