refactor(Probability): define statistical distance on arbitrary types

Cslib/Crypto/Protocols/Commitment/Basic.lean

@@ -44,14 +44,12 @@ open scoped NNReal
 variable {Message Commitment Opening : Type*}
 
 /-- Perfect hiding is exactly statistical hiding with zero error. -/
-theorem perfectlyHiding_iff_statisticallyHiding_zero
-    [Fintype Commitment] (scheme : Scheme Message Commitment Opening) :
+theorem perfectlyHiding_iff_statisticallyHiding_zero (scheme : Scheme Message Commitment Opening) :
     scheme.PerfectlyHiding ↔ scheme.StatisticallyHiding 0 := by
   simp [PerfectlyHiding, StatisticallyHiding]
 
 /-- Enlarging the permitted error preserves statistical hiding. -/
-theorem StatisticallyHiding.mono [Fintype Commitment]
-    {scheme : Scheme Message Commitment Opening} :
+theorem StatisticallyHiding.mono {scheme : Scheme Message Commitment Opening} :
     Monotone scheme.StatisticallyHiding :=
   fun _ _ hεδ h message₀ message₁ =>
     StatisticallyClose.mono hεδ (h message₀ message₁)
@@ -74,7 +72,7 @@ theorem PerfectlyBinding.disjoint_support_commitmentDist
 messages are at the maximum statistical distance: an unbounded observer can
 read the message off the commitment. -/
 theorem PerfectlyBinding.dist_commitmentDist_eq_one
-    [Fintype Commitment] {scheme : Scheme Message Commitment Opening}
+    {scheme : Scheme Message Commitment Opening}
     (hbind : scheme.PerfectlyBinding) {message₀ message₁ : Message}
     (hne : message₀ ≠ message₁) :
     dist (scheme.commitmentDist message₀)
@@ -86,7 +84,7 @@ hiding with error below one and perfectly binding unless any two messages are
 equal. The error bound is sharp: statistical hiding with error one holds
 vacuously for every scheme. -/
 theorem subsingleton_of_statisticallyHiding_of_perfectlyBinding
-    [Fintype Commitment] (scheme : Scheme Message Commitment Opening) {ε : ℝ≥0}
+    (scheme : Scheme Message Commitment Opening) {ε : ℝ≥0}
     (hε : ε < 1) (hhide : scheme.StatisticallyHiding ε)
     (hbind : scheme.PerfectlyBinding) : Subsingleton Message := by
   refine ⟨fun message₀ message₁ => ?_⟩
@@ -97,21 +95,12 @@ theorem subsingleton_of_statisticallyHiding_of_perfectlyBinding
   exact absurd hle (by exact_mod_cast hε.not_ge)
 
 /-- A scheme cannot be both perfectly hiding and perfectly binding unless any
-two messages are equal. Unlike the statistical version, this needs no
-finiteness assumption on the commitment type. -/
+two messages are equal. -/
 theorem subsingleton_of_perfectlyHiding_of_perfectlyBinding
     (scheme : Scheme Message Commitment Opening)
     (hhide : scheme.PerfectlyHiding) (hbind : scheme.PerfectlyBinding) :
-    Subsingleton Message := by
-  refine ⟨fun message₀ message₁ => ?_⟩
-  obtain ⟨commitment, hcommitment⟩ :=
-    (scheme.commitmentDist message₀).support_nonempty
-  obtain ⟨opening₀, hpair₀⟩ :=
-    scheme.mem_support_commitmentDist_iff.mp hcommitment
-  rw [hhide message₀ message₁] at hcommitment
-  obtain ⟨opening₁, hpair₁⟩ :=
-    scheme.mem_support_commitmentDist_iff.mp hcommitment
-  exact hbind commitment message₀ opening₀ message₁ opening₁
-    (scheme.accepts_of_mem_support hpair₀) (scheme.accepts_of_mem_support hpair₁)
+    Subsingleton Message :=
+  scheme.subsingleton_of_statisticallyHiding_of_perfectlyBinding zero_lt_one
+    (scheme.perfectlyHiding_iff_statisticallyHiding_zero.mp hhide) hbind
 
 end Cslib.Crypto.Protocols.Commitment.Scheme

Cslib/Crypto/Protocols/Commitment/Defs.lean

@@ -58,8 +58,7 @@ def PerfectlyHiding (scheme : Scheme Message Commitment Opening) : Prop :=
 /-- A scheme is statistically hiding with error `ε` when the commitment
 distributions of any two messages are within statistical distance `ε`
 ([BonehShoup2023], Definition 3.5 and Section 8.12). -/
-def StatisticallyHiding [Fintype Commitment]
-    (scheme : Scheme Message Commitment Opening) (ε : ℝ≥0) : Prop :=
+def StatisticallyHiding (scheme : Scheme Message Commitment Opening) (ε : ℝ≥0) : Prop :=
   ∀ message₀ message₁ : Message,
     StatisticallyClose (scheme.commitmentDist message₀)
       (scheme.commitmentDist message₁) ε

Cslib/Probability/StatisticalDistance.lean

@@ -6,19 +6,22 @@ Authors: Samuel Schlesinger
 
 module
 
-public import Cslib.Init
+public import Cslib.Probability.PMF
+public import Mathlib.Analysis.Normed.Group.InfiniteSum
 public import Mathlib.Probability.ProbabilityMassFunction.Constructions
+public import Mathlib.Topology.Algebra.InfiniteSum.Real
 public import Mathlib.Topology.MetricSpace.Defs
 
 /-!
-# Statistical Distance of Finite Probability Mass Functions
+# Statistical Distance of Probability Mass Functions
 
-For PMFs `p` and `q` on a finite type, their statistical distance is
+For PMFs `p` and `q`, their statistical distance is
 
-`(1 / 2) * ∑ a, |p a - q a|`.
+`(1 / 2) * ∑' a, |p a - q a|`.
 
-This is [BonehShoup2023], Definition 3.5. The probabilities are converted from
-`ℝ≥0∞`, Mathlib's codomain for a `PMF`, to `ℝ` before taking the finite sum.
+On finite types this is [BonehShoup2023], Definition 3.5. The probabilities are converted from
+`ℝ≥0∞`, Mathlib's codomain for a `PMF`, to `ℝ` before summing. Absolute summability follows
+from the unit mass of each distribution, so the same API applies to infinite types such as `ℕ`.
 
 Statistical distance is packaged as a scoped `MetricSpace` instance on
 `PMF α`, so it is spelled `dist p q` and the general metric API applies:
@@ -63,56 +66,50 @@ universe u v
 
 variable {α : Type u} {β : Type v}
 
-private theorem sum_toReal [Fintype α] (p : PMF α) :
-    ∑ a, (p a).toReal = 1 := by
-  rw [← ENNReal.toReal_one, ← p.tsum_coe, tsum_fintype,
-    ENNReal.toReal_sum fun a _ => p.apply_ne_top a]
-
-private theorem bind_apply_toReal [Fintype α] (p : PMF α)
-    (kernel : α → PMF β) (b : β) :
-    (p.bind kernel b).toReal =
-      ∑ a, (p a).toReal * (kernel a b).toReal := by
-  rw [PMF.bind_apply, tsum_fintype,
-    ENNReal.toReal_sum fun a _ =>
-      ENNReal.mul_ne_top (p.apply_ne_top a) ((kernel a).apply_ne_top b)]
-  simp
-
-/-- Statistical distance makes the PMFs on a finite type a metric space
-([BonehShoup2023], Definition 3.5). -/
-noncomputable scoped instance instMetricSpace [Fintype α] :
+private theorem summable_abs_sub (p q : PMF α) :
+    Summable (fun a => |(p a).toReal - (q a).toReal|) :=
+  ((summable_toReal p).sub (summable_toReal q)).abs
+
+/-- Statistical distance makes discrete probability distributions a metric space. -/
+noncomputable scoped instance instMetricSpace :
     MetricSpace (PMF α) where
-  dist p q := (∑ a, |(p a).toReal - (q a).toReal|) / 2
+  dist p q := (∑' a, |(p a).toReal - (q a).toReal|) / 2
   dist_self p := by simp
   dist_comm p q := by simp [abs_sub_comm]
   dist_triangle p q r := by
-    simp only [← add_div, ← Finset.sum_add_distrib]
-    gcongr with a
-    exact abs_sub_le (p a).toReal (q a).toReal (r a).toReal
+    rw [← add_div, ← (summable_abs_sub p q).tsum_add (summable_abs_sub q r)]
+    gcongr ?_ / _
+    exact (summable_abs_sub p r).tsum_le_tsum (fun a => abs_sub_le _ _ _)
+      ((summable_abs_sub p q).add (summable_abs_sub q r))
   eq_of_dist_eq_zero {p q} h := by
-    have hsum : ∑ a, |(p a).toReal - (q a).toReal| = 0 := by
+    have hsum : ∑' a, |(p a).toReal - (q a).toReal| = 0 := by
       simpa [div_eq_zero_iff] using h
     ext a
     apply (ENNReal.toReal_eq_toReal_iff' (p.apply_ne_top a) (q.apply_ne_top a)).mp
-    simpa [sub_eq_zero] using congr_fun
-      ((Fintype.sum_eq_zero_iff_of_nonneg fun _ => abs_nonneg _).mp hsum) a
+    have ha := (summable_abs_sub p q).le_tsum a (fun _ _ => abs_nonneg _)
+    simpa [hsum, sub_eq_zero] using ha
+
+/-- Statistical distance is half the sum of the absolute differences of the masses. -/
+theorem dist_eq_tsum (p q : PMF α) :
+    dist p q = (∑' a, |(p a).toReal - (q a).toReal|) / 2 := rfl
 
 /-- The distance between two PMFs on a finite type is their statistical
 distance ([BonehShoup2023], Definition 3.5). -/
 theorem dist_eq [Fintype α] (p q : PMF α) :
-    dist p q = (∑ a, |(p a).toReal - (q a).toReal|) / 2 :=
-  rfl
+    dist p q = (∑ a, |(p a).toReal - (q a).toReal|) / 2 := by
+  simp [dist_eq_tsum]
 
 /-- Statistical distance is at most one. -/
-theorem dist_le_one [Fintype α] (p q : PMF α) : dist p q ≤ 1 := by
-  rw [dist_eq]
-  have h := Finset.sum_le_sum fun a (_ : a ∈ Finset.univ) =>
-    abs_sub_le (p a).toReal 0 (q a).toReal
-  simp only [sub_zero, zero_sub, abs_neg, abs_of_nonneg ENNReal.toReal_nonneg,
-    Finset.sum_add_distrib, sum_toReal] at h
+theorem dist_le_one (p q : PMF α) : dist p q ≤ 1 := by
+  rw [dist_eq_tsum]
+  have h := Summable.tsum_le_tsum
+    (fun a => by simpa using abs_sub_le (p a).toReal 0 (q a).toReal)
+    (summable_abs_sub p q) ((summable_toReal p).add (summable_toReal q))
+  rw [(summable_toReal p).tsum_add (summable_toReal q), tsum_toReal, tsum_toReal] at h
   linarith
 
 /-- PMFs with disjoint supports are at the maximum statistical distance. -/
-theorem dist_eq_one_of_disjoint_support [Fintype α] {p q : PMF α}
+theorem dist_eq_one_of_disjoint_support {p q : PMF α}
     (h : Disjoint p.support q.support) : dist p q = 1 := by
   have key : ∀ a, |(p a).toReal - (q a).toReal| = (p a).toReal + (q a).toReal := by
     intro a
@@ -123,73 +120,91 @@ theorem dist_eq_one_of_disjoint_support [Fintype α] {p q : PMF α}
         exact Set.disjoint_left.mp h ((p.mem_support_iff a).mpr hp)
           ((q.mem_support_iff a).mpr hq)
       simp [hq]
-  simp [dist_eq, key, Finset.sum_add_distrib, sum_toReal]
+  rw [dist_eq_tsum, tsum_congr key, (summable_toReal p).tsum_add (summable_toReal q),
+    tsum_toReal, tsum_toReal]
+  norm_num
+
+private theorem summable_weighted_kernel {weight : α → ℝ} (hweight : Summable weight)
+    (hnonneg : ∀ a, 0 ≤ weight a) (kernel : α → PMF β) :
+    Summable (fun pair : α × β => weight pair.1 * (kernel pair.1 pair.2).toReal) := by
+  rw [summable_prod_of_nonneg (fun pair => mul_nonneg (hnonneg pair.1) ENNReal.toReal_nonneg)]
+  exact ⟨fun a => (summable_toReal (kernel a)).mul_left (weight a),
+    by simpa [tsum_mul_left] using hweight⟩
+
+/-- Weighting by the masses of a kernel at a fixed outcome preserves summability. -/
+private theorem summable_mul_kernel {weight : α → ℝ} (hweight : Summable weight)
+    (kernel : α → PMF β) (b : β) : Summable (fun a => weight a * (kernel a b).toReal) := by
+  refine Summable.of_norm_bounded (g := fun a => |weight a|) hweight.abs fun a => ?_
+  rw [Real.norm_eq_abs, abs_mul, abs_of_nonneg ENNReal.toReal_nonneg]
+  exact mul_le_of_le_one_right (abs_nonneg _)
+    (ENNReal.toReal_le_of_le_ofReal zero_le_one (by simpa using (kernel a).coe_le_one b))
 
 /-- Applying the same randomized kernel to two PMFs cannot increase their
 statistical distance. -/
-theorem dist_bind_le [Fintype α] [Fintype β]
-    (p q : PMF α) (kernel : α → PMF β) :
+theorem dist_bind_le (p q : PMF α) (kernel : α → PMF β) :
     dist (p.bind kernel) (q.bind kernel) ≤ dist p q := by
-  simp only [dist_eq, bind_apply_toReal]
-  gcongr
+  have hd := summable_weighted_kernel (summable_abs_sub p q) (fun _ => abs_nonneg _) kernel
+  have hbound (b : β) : |((p.bind kernel) b).toReal - ((q.bind kernel) b).toReal| ≤
+      ∑' a, |(p a).toReal - (q a).toReal| * (kernel a b).toReal := by
+    rw [bind_apply_toReal_tsum, bind_apply_toReal_tsum, ← (summable_mul_kernel
+      (summable_toReal p) kernel b).tsum_sub (summable_mul_kernel (summable_toReal q) kernel b)]
+    simp_rw [← sub_mul]
+    have hs := summable_mul_kernel ((summable_toReal p).sub (summable_toReal q)) kernel b
+    simpa using
+      norm_tsum_le_tsum_norm (f := fun a => ((p a).toReal - (q a).toReal) * (kernel a b).toReal)
+        (by simpa using hs.abs)
+  rw [dist_eq_tsum, dist_eq_tsum]
+  gcongr ?_ / _
   calc
-    (∑ b, |(∑ a, (p a).toReal * (kernel a b).toReal) -
-        ∑ a, (q a).toReal * (kernel a b).toReal|)
-        ≤ ∑ b, ∑ a, |(p a).toReal * (kernel a b).toReal -
-          (q a).toReal * (kernel a b).toReal| :=
-      Finset.sum_le_sum fun _ _ => by
-        rw [← Finset.sum_sub_distrib]
-        exact Finset.abs_sum_le_sum_abs _ _
-    _ = ∑ b, ∑ a, (kernel a b).toReal *
-        |(p a).toReal - (q a).toReal| := by
-      simp_rw [← sub_mul, abs_mul, abs_of_nonneg ENNReal.toReal_nonneg, mul_comm]
-    _ = ∑ a, |(p a).toReal - (q a).toReal| := by
-      rw [Finset.sum_comm]
-      simp [← Finset.sum_mul, sum_toReal]
+    _ ≤ ∑' b, ∑' a, |(p a).toReal - (q a).toReal| * (kernel a b).toReal :=
+      Summable.tsum_le_tsum hbound (summable_abs_sub _ _) hd.prod_symm.prod
+    _ = _ := by
+      rw [Summable.tsum_comm (f := fun a b =>
+        |(p a).toReal - (q a).toReal| * (kernel a b).toReal) hd]
+      simp [tsum_mul_left]
 
 /-- Deterministic postprocessing cannot increase statistical distance
 ([BonehShoup2023], Theorem 3.13). -/
-theorem dist_map_le [Fintype α] [Fintype β]
-    (p q : PMF α) (f : α → β) :
+theorem dist_map_le (p q : PMF α) (f : α → β) :
     dist (p.map f) (q.map f) ≤ dist p q := by
   simpa [PMF.bind_pure_comp] using dist_bind_le p q (PMF.pure ∘ f)
 
 /-- Two PMFs are `ε`-statistically close when their statistical distance is at
 most `ε`. The `ℝ≥0` parameter rules out meaningless negative bounds. -/
-def StatisticallyClose [Fintype α] (p q : PMF α) (ε : ℝ≥0) : Prop :=
+def StatisticallyClose (p q : PMF α) (ε : ℝ≥0) : Prop :=
   dist p q ≤ (ε : ℝ)
 
 namespace StatisticallyClose
 
 /-- Every PMF is statistically close to itself with zero error. -/
-theorem refl [Fintype α] (p : PMF α) : StatisticallyClose p p 0 := by
+theorem refl (p : PMF α) : StatisticallyClose p p 0 := by
   simp [StatisticallyClose]
 
 /-- Statistical closeness is symmetric. -/
-theorem symm [Fintype α] {p q : PMF α} {ε : ℝ≥0}
+theorem symm {p q : PMF α} {ε : ℝ≥0}
     (h : StatisticallyClose p q ε) : StatisticallyClose q p ε := by
   simpa [StatisticallyClose, dist_comm] using h
 
 /-- A statistical-closeness bound remains valid when its error is enlarged. -/
-theorem mono [Fintype α] {p q : PMF α} : Monotone (StatisticallyClose p q) :=
+theorem mono {p q : PMF α} : Monotone (StatisticallyClose p q) :=
   fun _ _ hεδ h => le_trans h (by exact_mod_cast hεδ)
 
 /-- Closeness bounds chain through an intermediate distribution, adding the
 errors. -/
-theorem trans [Fintype α] {p q r : PMF α} {ε δ : ℝ≥0}
+theorem trans {p q r : PMF α} {ε δ : ℝ≥0}
     (hpq : StatisticallyClose p q ε) (hqr : StatisticallyClose q r δ) :
     StatisticallyClose p r (ε + δ) :=
   (dist_triangle p q r).trans (by simpa using add_le_add hpq hqr)
 
 /-- A shared randomized postprocessing kernel preserves statistical
 closeness. -/
-theorem bind [Fintype α] [Fintype β] {p q : PMF α} {ε : ℝ≥0}
+theorem bind {p q : PMF α} {ε : ℝ≥0}
     (h : StatisticallyClose p q ε) (kernel : α → PMF β) :
     StatisticallyClose (p.bind kernel) (q.bind kernel) ε :=
   (dist_bind_le p q kernel).trans h
 
 /-- Deterministic postprocessing preserves statistical closeness. -/
-theorem map [Fintype α] [Fintype β] {p q : PMF α} {ε : ℝ≥0}
+theorem map {p q : PMF α} {ε : ℝ≥0}
     (h : StatisticallyClose p q ε) (f : α → β) :
     StatisticallyClose (p.map f) (q.map f) ε :=
   (dist_map_le p q f).trans h
@@ -198,7 +213,7 @@ end StatisticallyClose
 
 /-- Statistical closeness with zero error is equality. -/
 @[simp]
-theorem statisticallyClose_zero_iff [Fintype α] (p q : PMF α) :
+theorem statisticallyClose_zero_iff (p q : PMF α) :
     StatisticallyClose p q 0 ↔ p = q := by
   simp [StatisticallyClose]
 

CslibTests/Commitment.lean

@@ -12,7 +12,7 @@ open Cslib.Crypto.Protocols.Commitment
 open Cslib.Probability.PMF
 open scoped NNReal
 
-example {α β : Type*} [Fintype α] [Fintype β] {p q : PMF α} {ε : ℝ≥0}
+example {α β : Type*} {p q : PMF α} {ε : ℝ≥0}
     (h : StatisticallyClose p q ε) (kernel : α → PMF β) :
     StatisticallyClose (p.bind kernel) (q.bind kernel) ε :=
   h.bind kernel
@@ -41,7 +41,7 @@ theorem revealingScheme_perfectlyBinding (Message : Type*) [DecidableEq Message]
 
 /-- The hiding–binding trade-off: no scheme with two distinct messages is both
 statistically hiding with error below one and perfectly binding. -/
-example {ε : ℝ≥0} (hε : ε < 1) (scheme : Scheme Bool Bool Unit)
+example {ε : ℝ≥0} (hε : ε < 1) (scheme : Scheme Bool ℕ Unit)
     (hhide : scheme.StatisticallyHiding ε) (hbind : scheme.PerfectlyBinding) :
     False := by
   have h := scheme.subsingleton_of_statisticallyHiding_of_perfectlyBinding hε hhide hbind