From ad5e2b63332630e4aaacc3dcb2b18ca5df267120 Mon Sep 17 00:00:00 2001
From: Alley Stoughton <alley.stoughton@icloud.com>
Date: Mon, 6 Apr 2026 13:11:07 -0400
Subject: [PATCH 1/6] [theories/modules] law of total probability

Add module-based law of total probability for lossless distributions
with finite support.
---
 theories/modules/TotalProb.ec | 259 ++++++++++++++++++++++++++++++++++
 1 file changed, 259 insertions(+)
 create mode 100644 theories/modules/TotalProb.ec

diff --git a/theories/modules/TotalProb.ec b/theories/modules/TotalProb.ec
new file mode 100644
index 000000000..226765647
--- /dev/null
+++ b/theories/modules/TotalProb.ec
@@ -0,0 +1,259 @@
+(*^ This theory proves a module-based law of total probability for
+    distributions with finite support, and then specializes it to
+    `DBool.dbool`, `drange` and `duniform`. ^*)
+
+require import AllCore List StdOrder StdBigop Distr DBool.
+require EventPartitioning.
+import RealOrder Bigreal BRA.
+
+(*& Are the elements in the support of a distribution the
+    same as the elements of a list? &*)
+
+op support_is (d : 'a distr) (xs : 'a list) : bool =
+  forall (y : 'a), y \in d <=> y \in xs.
+
+(*& General abstract theory &*)
+
+abstract theory TotalGeneral.
+
+(*& theory parameter: additional input &*)
+
+type input.
+
+(*& theory parameter: type of distribution &*)
+
+type t.
+
+(*& A module with a `main` procedure taking in an input plus
+    a value of type `t`, and returning a boolean result. &*)
+
+module type T = {
+  proc main(i : input, x : t) : bool
+}.
+
+(*& For a module `M` with module type `T` and a distribution `dt` for
+    `t` plus an input `i`, `Rand(M).f(dt, i)` samples a value `x` from
+    `dt`, passes it to `M.main` along with `i`, and returns the
+    boolean `M.main` returns. &*)
+
+module Rand(M : T) = {
+  proc f(dt : t distr, i : input) : bool = {
+    var b : bool; var x : t;
+    x <$ dt;
+    b <@ M.main(i, x);
+    return b;
+  }
+}.
+
+(* skip to end of section for main lemma *)
+
+section.
+
+declare module M <: T.
+
+local module RandAux = {
+  var x : t  (* a global variable *)
+
+  proc f(dt : t distr, i : input) : bool = {
+    var b : bool;
+    x <$ dt;
+    b <@ M.main(i, x);
+    return b;
+  }
+}.
+
+local clone EventPartitioning as EP with
+  type input <- t distr * input,
+  type output <- bool
+proof *.
+
+local clone EP.ListPartitioning as EPL with
+  type partition <- t
+proof *.
+
+local op x_of_glob_RA (g : (glob RandAux)) : t = g.`2.
+local op phi (_ : t distr * input) (g : glob RandAux) (_ : bool) : t =
+  x_of_glob_RA g.
+local op E (_ : t distr * input) (_ : glob RandAux) (o : bool) : bool = o.
+
+local lemma RandAux_partition_eq (dt' : t distr) (i' : input) (x' : t) &m :
+  Pr[RandAux.f(dt', i') @ &m : res /\ x_of_glob_RA (glob RandAux) = x'] =
+  mu1 dt' x' * Pr[M.main(i', x') @ &m : res].
+proof.
+byphoare
+  (_ :
+   dt = dt' /\ i = i' /\ (glob M) = (glob M){m} ==>
+   res /\ x_of_glob_RA (glob RandAux) = x') => //.
+proc; rewrite /x_of_glob_RA /=.
+seq 1 :
+  (RandAux.x = x')
+  (mu1 dt' x')
+  Pr[M.main(i', x') @ &m : res]
+  (1%r - mu1 dt x')
+  0%r
+  (i = i' /\ (glob M) = (glob M){m}) => //.
+auto.
+rnd (pred1 x'); auto.
+call (_ : i = i' /\ x = x' /\ glob M = (glob M){m} ==> res).
+bypr => &hr [#] -> -> glob_eq.
+byequiv (_ : ={i, x, glob M} ==> ={res}) => //; sim.
+auto.
+hoare; call (_ : true); auto; smt().
+qed.
+
+lemma total_prob' (dt' : t distr) (supp : t list) (i' : input) &m :
+  is_lossless dt' => uniq supp => support_is dt' supp =>
+  Pr[Rand(M).f(dt', i') @ &m : res] =
+  big predT
+  (fun x' => mu1 dt' x' * Pr[M.main(i', x') @ &m : res])
+  supp.
+proof.
+move => dt'_ll uniq_supp support_is_dt'_supp.
+have -> :
+  Pr[Rand(M).f(dt', i') @ &m : res] = Pr[RandAux.f(dt', i') @ &m : res].
+  byequiv (_ : ={dt, i, glob M} ==> ={res}) => //; proc; sim.
+rewrite (EPL.list_partitioning RandAux (dt', i') E phi supp &m) //.
+have -> /= :
+  Pr[RandAux.f(dt', i') @ &m: res /\ ! (x_of_glob_RA (glob RandAux) \in supp)] =
+  0%r.
+  byphoare (_ : dt = dt' /\ i = i' ==> _) => //; proc.
+  seq 1 :
+    (RandAux.x \in supp)
+    1%r
+    0%r
+    0%r
+    1%r => //.
+  auto.
+  hoare; call (_ : true); auto; smt().
+  rnd pred0; auto; smt(mu0).
+congr; apply fun_ext => x'; by rewrite RandAux_partition_eq.
+qed.
+
+end section.
+
+(*& total probability lemma for distributions with finite support &*)
+
+lemma total_prob (M <: T) (dt : t distr) (supp : t list) (i : input) &m :
+  is_lossless dt => uniq supp => support_is dt supp =>
+  Pr[Rand(M).f(dt, i) @ &m : res] =
+  big predT
+  (fun (x : t) => mu1 dt x * Pr[M.main(i, x) @ &m : res])
+  supp.
+proof.
+move => dt_ll uniq_supp support_is_dt_supp.
+by apply (total_prob' M).
+qed.
+
+end TotalGeneral.
+
+(*& Specialization to `DBool.dbool` &*)
+
+abstract theory TotalBool.
+
+clone include TotalGeneral with
+  type t <- bool
+proof *.
+
+(*& total probability lemma for `DBool.dbool` &*)
+
+lemma total_prob_bool (M <: T) (i : input) &m :
+  Pr[Rand(M).f(DBool.dbool, i) @ &m : res] =
+  Pr[M.main(i, true) @ &m : res]  / 2%r +
+  Pr[M.main(i, false) @ &m : res] / 2%r.
+proof.
+rewrite (total_prob M DBool.dbool [true; false]) //.
+rewrite dbool_ll.
+smt(supp_dbool).
+by rewrite 2!big_cons big_nil /= /predT /predF /= 2!dbool1E.
+qed.
+
+end TotalBool.
+
+(*& Specialization to `drange` &*)
+
+abstract theory TotalRange.
+
+clone include TotalGeneral with
+  type t <- int
+proof *.
+
+lemma big_weight_simp (M <: T) (m n : int, i : input, ys : int list) &m :
+  (forall y, y \in ys => m <= y < n) =>
+  big predT
+  (fun (x : int) => mu1 (drange m n) x * Pr[M.main(i, x) @ &m : res])
+  ys =
+  big predT
+  (fun (x : int) => Pr[M.main(i, x) @ &m : res] / (n - m)%r)
+  ys.
+proof.
+elim ys => [// | y ys IH mem_impl].
+rewrite 2!big_cons (_ : predT y) //=.
+rewrite drange1E (_ : m <= y < n) 1:mem_impl //=.
+rewrite IH => [z z_in_ys | //].
+by rewrite mem_impl /= z_in_ys.
+qed.
+
+(*& total probability lemma for `drange` &*)
+
+lemma total_prob_drange (M <: T) (m n : int, i : input) &m :
+  m < n =>
+  Pr[Rand(M).f(drange m n, i) @ &m : res] =
+  bigi predT
+  (fun (j : int) => Pr[M.main(i, j) @ &m : res] / (n - m)%r)
+  m n.
+proof.
+move => lt_m_n.
+rewrite (total_prob M (drange m n) (range m n)).
+by rewrite drange_ll.
+by rewrite range_uniq.
+rewrite /support_is => j; smt(supp_drange mem_range).
+rewrite (big_weight_simp M) //; by move => j /mem_range.
+qed.
+
+end TotalRange.
+
+(*& Specialization to `duniform` &*)
+
+abstract theory TotalUniform.
+
+(*& theory parameter: type &*)
+
+type t.
+
+clone include TotalGeneral with
+  type t <- t
+proof *.
+
+lemma big_weight_simp (M <: T) (xs : t list, i : input, ys : t list) &m :
+  uniq xs =>
+  (forall y, y \in ys => y \in xs) =>
+  big predT
+  (fun (x : t) => mu1 (duniform xs) x * Pr[M.main(i, x) @ &m : res])
+  ys =
+  big predT
+  (fun (x : t) => Pr[M.main(i, x) @ &m : res] / (size xs)%r)
+  ys.
+proof.
+move => uniq_xs.
+elim ys => [// | y ys IH mem_impl].
+smt(big_cons duniform1E_uniq).
+qed.
+
+(*& total probability lemma for `duniform` &*)
+
+lemma total_prob_uniform (M <: T) (xs : t list, i : input) &m :
+  uniq xs => xs <> [] =>
+  Pr[Rand(M).f(duniform xs, i) @ &m : res] =
+  big predT
+  (fun (x : t) => Pr[M.main(i, x) @ &m : res] / (size xs)%r)
+  xs.
+proof.
+move => uniq_xs xs_ne_nil.
+rewrite (total_prob M (duniform xs) xs).
+by rewrite duniform_ll.
+by rewrite uniq_xs.
+rewrite /support_is => y; smt(supp_duniform).
+by rewrite (big_weight_simp M).
+qed.
+
+end TotalUniform.

From 1206087c4bceda7f3bc1e9f09532a9fbb48710c2 Mon Sep 17 00:00:00 2001
From: Alley Stoughton <alley.stoughton@icloud.com>
Date: Mon, 6 Apr 2026 14:39:54 -0400
Subject: [PATCH 2/6] fixed to work with just Z3

---
 theories/modules/TotalProb.ec | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/theories/modules/TotalProb.ec b/theories/modules/TotalProb.ec
index 226765647..32df85d93 100644
--- a/theories/modules/TotalProb.ec
+++ b/theories/modules/TotalProb.ec
@@ -206,7 +206,9 @@ move => lt_m_n.
 rewrite (total_prob M (drange m n) (range m n)).
 by rewrite drange_ll.
 by rewrite range_uniq.
-rewrite /support_is => j; smt(supp_drange mem_range).
+rewrite /support_is => j; split => [j_in_drange | j_in_range].
+rewrite mem_range; by rewrite supp_drange in j_in_drange.
+rewrite supp_drange; by rewrite mem_range in j_in_range.
 rewrite (big_weight_simp M) //; by move => j /mem_range.
 qed.
 
@@ -236,7 +238,9 @@ lemma big_weight_simp (M <: T) (xs : t list, i : input, ys : t list) &m :
 proof.
 move => uniq_xs.
 elim ys => [// | y ys IH mem_impl].
-smt(big_cons duniform1E_uniq).
+rewrite 2!big_cons (_ : predT true) //=.
+rewrite duniform1E_uniq // (_ : y \in xs) 1:mem_impl //=.
+rewrite IH => [z z_in_ys | //]; by rewrite mem_impl /= z_in_ys.
 qed.
 
 (*& total probability lemma for `duniform` &*)

From a5267b07ac9b385103748224550b7127ace1efb2 Mon Sep 17 00:00:00 2001
From: Alley Stoughton <alley.stoughton@icloud.com>
Date: Tue, 7 Apr 2026 10:19:21 -0400
Subject: [PATCH 3/6] used is_finite_for instead of new operator

---
 theories/modules/TotalProb.ec | 25 ++++++++-----------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/theories/modules/TotalProb.ec b/theories/modules/TotalProb.ec
index 32df85d93..24b8e0a36 100644
--- a/theories/modules/TotalProb.ec
+++ b/theories/modules/TotalProb.ec
@@ -2,16 +2,10 @@
     distributions with finite support, and then specializes it to
     `DBool.dbool`, `drange` and `duniform`. ^*)
 
-require import AllCore List StdOrder StdBigop Distr DBool.
+require import AllCore List StdOrder StdBigop Distr DBool Finite.
 require EventPartitioning.
 import RealOrder Bigreal BRA.
 
-(*& Are the elements in the support of a distribution the
-    same as the elements of a list? &*)
-
-op support_is (d : 'a distr) (xs : 'a list) : bool =
-  forall (y : 'a), y \in d <=> y \in xs.
-
 (*& General abstract theory &*)
 
 abstract theory TotalGeneral.
@@ -102,13 +96,13 @@ hoare; call (_ : true); auto; smt().
 qed.
 
 lemma total_prob' (dt' : t distr) (supp : t list) (i' : input) &m :
-  is_lossless dt' => uniq supp => support_is dt' supp =>
+  is_lossless dt' => is_finite_for (support dt') supp =>
   Pr[Rand(M).f(dt', i') @ &m : res] =
   big predT
   (fun x' => mu1 dt' x' * Pr[M.main(i', x') @ &m : res])
   supp.
 proof.
-move => dt'_ll uniq_supp support_is_dt'_supp.
+move => dt'_ll [uniq_supp supp_iff]. 
 have -> :
   Pr[Rand(M).f(dt', i') @ &m : res] = Pr[RandAux.f(dt', i') @ &m : res].
   byequiv (_ : ={dt, i, glob M} ==> ={res}) => //; proc; sim.
@@ -134,13 +128,13 @@ end section.
 (*& total probability lemma for distributions with finite support &*)
 
 lemma total_prob (M <: T) (dt : t distr) (supp : t list) (i : input) &m :
-  is_lossless dt => uniq supp => support_is dt supp =>
+  is_lossless dt => is_finite_for (support dt) supp =>
   Pr[Rand(M).f(dt, i) @ &m : res] =
   big predT
   (fun (x : t) => mu1 dt x * Pr[M.main(i, x) @ &m : res])
   supp.
 proof.
-move => dt_ll uniq_supp support_is_dt_supp.
+move => dt_ll iff_supp_dt_supp.
 by apply (total_prob' M).
 qed.
 
@@ -205,10 +199,8 @@ proof.
 move => lt_m_n.
 rewrite (total_prob M (drange m n) (range m n)).
 by rewrite drange_ll.
-by rewrite range_uniq.
-rewrite /support_is => j; split => [j_in_drange | j_in_range].
-rewrite mem_range; by rewrite supp_drange in j_in_drange.
-rewrite supp_drange; by rewrite mem_range in j_in_range.
+rewrite /is_finite_for.
+smt(range_uniq mem_range supp_drange).
 rewrite (big_weight_simp M) //; by move => j /mem_range.
 qed.
 
@@ -255,8 +247,7 @@ proof.
 move => uniq_xs xs_ne_nil.
 rewrite (total_prob M (duniform xs) xs).
 by rewrite duniform_ll.
-by rewrite uniq_xs.
-rewrite /support_is => y; smt(supp_duniform).
+smt(supp_duniform).
 by rewrite (big_weight_simp M).
 qed.
 

From d2534bf4ab728329c74f1f88e404e7220cc707df Mon Sep 17 00:00:00 2001
From: Alley Stoughton <alley.stoughton@icloud.com>
Date: Mon, 13 Apr 2026 10:13:45 -0400
Subject: [PATCH 4/6] various improvements, per feedback

---
 theories/modules/TotalProb.ec | 93 ++++++++++-------------------------
 1 file changed, 27 insertions(+), 66 deletions(-)

diff --git a/theories/modules/TotalProb.ec b/theories/modules/TotalProb.ec
index 24b8e0a36..855c03bda 100644
--- a/theories/modules/TotalProb.ec
+++ b/theories/modules/TotalProb.ec
@@ -86,40 +86,39 @@ seq 1 :
   (1%r - mu1 dt x')
   0%r
   (i = i' /\ (glob M) = (glob M){m}) => //.
-auto.
-rnd (pred1 x'); auto.
-call (_ : i = i' /\ x = x' /\ glob M = (glob M){m} ==> res).
-bypr => &hr [#] -> -> glob_eq.
-byequiv (_ : ={i, x, glob M} ==> ={res}) => //; sim.
-auto.
++ auto.
++ rnd (pred1 x'); auto.
++ call (_ : i = i' /\ x = x' /\ glob M = (glob M){m} ==> res).
+  + bypr => &hr [#] -> -> glob_eq.
+    byequiv (_ : ={i, x, glob M} ==> ={res}) => //; sim.
++ auto.
 hoare; call (_ : true); auto; smt().
 qed.
 
 lemma total_prob' (dt' : t distr) (supp : t list) (i' : input) &m :
-  is_lossless dt' => is_finite_for (support dt') supp =>
+  is_finite_for (support dt') supp =>
   Pr[Rand(M).f(dt', i') @ &m : res] =
   big predT
   (fun x' => mu1 dt' x' * Pr[M.main(i', x') @ &m : res])
   supp.
 proof.
-move => dt'_ll [uniq_supp supp_iff]. 
+move => [uniq_supp supp_iff]. 
 have -> :
   Pr[Rand(M).f(dt', i') @ &m : res] = Pr[RandAux.f(dt', i') @ &m : res].
-  byequiv (_ : ={dt, i, glob M} ==> ={res}) => //; proc; sim.
+  + byequiv (_ : ={dt, i, glob M} ==> ={res}) => //; proc; sim.
 rewrite (EPL.list_partitioning RandAux (dt', i') E phi supp &m) //.
 have -> /= :
   Pr[RandAux.f(dt', i') @ &m: res /\ ! (x_of_glob_RA (glob RandAux) \in supp)] =
   0%r.
-  byphoare (_ : dt = dt' /\ i = i' ==> _) => //; proc.
-  seq 1 :
-    (RandAux.x \in supp)
-    1%r
-    0%r
-    0%r
-    1%r => //.
-  auto.
-  hoare; call (_ : true); auto; smt().
-  rnd pred0; auto; smt(mu0).
+  + byphoare (_ : dt = dt' /\ i = i' ==> _) => //; proc.
+    seq 1 :
+      (RandAux.x \in supp)
+      1%r
+      0%r
+      0%r
+      1%r => //.
+    + hoare; call (_ : true); auto; smt().
+    rnd pred0; auto; smt(mu0).
 congr; apply fun_ext => x'; by rewrite RandAux_partition_eq.
 qed.
 
@@ -128,15 +127,12 @@ end section.
 (*& total probability lemma for distributions with finite support &*)
 
 lemma total_prob (M <: T) (dt : t distr) (supp : t list) (i : input) &m :
-  is_lossless dt => is_finite_for (support dt) supp =>
+  is_finite_for (support dt) supp =>
   Pr[Rand(M).f(dt, i) @ &m : res] =
   big predT
   (fun (x : t) => mu1 dt x * Pr[M.main(i, x) @ &m : res])
   supp.
-proof.
-move => dt_ll iff_supp_dt_supp.
-by apply (total_prob' M).
-qed.
+proof. move => iff_supp_dt_supp; by apply (total_prob' M). qed.
 
 end TotalGeneral.
 
@@ -156,8 +152,7 @@ lemma total_prob_bool (M <: T) (i : input) &m :
   Pr[M.main(i, false) @ &m : res] / 2%r.
 proof.
 rewrite (total_prob M DBool.dbool [true; false]) //.
-rewrite dbool_ll.
-smt(supp_dbool).
++ smt(supp_dbool).
 by rewrite 2!big_cons big_nil /= /predT /predF /= 2!dbool1E.
 qed.
 
@@ -171,22 +166,6 @@ clone include TotalGeneral with
   type t <- int
 proof *.
 
-lemma big_weight_simp (M <: T) (m n : int, i : input, ys : int list) &m :
-  (forall y, y \in ys => m <= y < n) =>
-  big predT
-  (fun (x : int) => mu1 (drange m n) x * Pr[M.main(i, x) @ &m : res])
-  ys =
-  big predT
-  (fun (x : int) => Pr[M.main(i, x) @ &m : res] / (n - m)%r)
-  ys.
-proof.
-elim ys => [// | y ys IH mem_impl].
-rewrite 2!big_cons (_ : predT y) //=.
-rewrite drange1E (_ : m <= y < n) 1:mem_impl //=.
-rewrite IH => [z z_in_ys | //].
-by rewrite mem_impl /= z_in_ys.
-qed.
-
 (*& total probability lemma for `drange` &*)
 
 lemma total_prob_drange (M <: T) (m n : int, i : input) &m :
@@ -198,10 +177,9 @@ lemma total_prob_drange (M <: T) (m n : int, i : input) &m :
 proof.
 move => lt_m_n.
 rewrite (total_prob M (drange m n) (range m n)).
-by rewrite drange_ll.
-rewrite /is_finite_for.
-smt(range_uniq mem_range supp_drange).
-rewrite (big_weight_simp M) //; by move => j /mem_range.
++ rewrite /is_finite_for; smt(range_uniq mem_range supp_drange).
+apply eq_big_seq => x x_in_range_m_n /=.
+by rewrite drange1E (_ : m <= x < n) 1:-mem_range.
 qed.
 
 end TotalRange.
@@ -218,23 +196,6 @@ clone include TotalGeneral with
   type t <- t
 proof *.
 
-lemma big_weight_simp (M <: T) (xs : t list, i : input, ys : t list) &m :
-  uniq xs =>
-  (forall y, y \in ys => y \in xs) =>
-  big predT
-  (fun (x : t) => mu1 (duniform xs) x * Pr[M.main(i, x) @ &m : res])
-  ys =
-  big predT
-  (fun (x : t) => Pr[M.main(i, x) @ &m : res] / (size xs)%r)
-  ys.
-proof.
-move => uniq_xs.
-elim ys => [// | y ys IH mem_impl].
-rewrite 2!big_cons (_ : predT true) //=.
-rewrite duniform1E_uniq // (_ : y \in xs) 1:mem_impl //=.
-rewrite IH => [z z_in_ys | //]; by rewrite mem_impl /= z_in_ys.
-qed.
-
 (*& total probability lemma for `duniform` &*)
 
 lemma total_prob_uniform (M <: T) (xs : t list, i : input) &m :
@@ -246,9 +207,9 @@ lemma total_prob_uniform (M <: T) (xs : t list, i : input) &m :
 proof.
 move => uniq_xs xs_ne_nil.
 rewrite (total_prob M (duniform xs) xs).
-by rewrite duniform_ll.
-smt(supp_duniform).
-by rewrite (big_weight_simp M).
++ smt(supp_duniform).
+apply eq_big_seq => y y_in_xs /=.
+by rewrite duniform1E_uniq // (_ : y \in xs).
 qed.
 
 end TotalUniform.

From 3182db9a4bba8d5f660275276b798330047fe106 Mon Sep 17 00:00:00 2001
From: Pierre-Yves Strub <pierre-yves.strub@pqshield.com>
Date: Mon, 13 Apr 2026 16:35:04 +0200
Subject: [PATCH 5/6] Use `by` on closing tactics and bullet branches for
 maintainability

---
 theories/modules/TotalProb.ec | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/theories/modules/TotalProb.ec b/theories/modules/TotalProb.ec
index 855c03bda..da9b61bbf 100644
--- a/theories/modules/TotalProb.ec
+++ b/theories/modules/TotalProb.ec
@@ -117,9 +117,9 @@ have -> /= :
       0%r
       0%r
       1%r => //.
-    + hoare; call (_ : true); auto; smt().
-    rnd pred0; auto; smt(mu0).
-congr; apply fun_ext => x'; by rewrite RandAux_partition_eq.
+    + by hoare; call (_ : true); auto; smt().
+    by rnd pred0; auto; smt(mu0).
+by congr; apply fun_ext => x'; rewrite RandAux_partition_eq.
 qed.
 
 end section.
@@ -152,7 +152,7 @@ lemma total_prob_bool (M <: T) (i : input) &m :
   Pr[M.main(i, false) @ &m : res] / 2%r.
 proof.
 rewrite (total_prob M DBool.dbool [true; false]) //.
-+ smt(supp_dbool).
++ by smt(supp_dbool).
 by rewrite 2!big_cons big_nil /= /predT /predF /= 2!dbool1E.
 qed.
 
@@ -177,7 +177,7 @@ lemma total_prob_drange (M <: T) (m n : int, i : input) &m :
 proof.
 move => lt_m_n.
 rewrite (total_prob M (drange m n) (range m n)).
-+ rewrite /is_finite_for; smt(range_uniq mem_range supp_drange).
++ by rewrite /is_finite_for; smt(range_uniq mem_range supp_drange).
 apply eq_big_seq => x x_in_range_m_n /=.
 by rewrite drange1E (_ : m <= x < n) 1:-mem_range.
 qed.
@@ -207,7 +207,7 @@ lemma total_prob_uniform (M <: T) (xs : t list, i : input) &m :
 proof.
 move => uniq_xs xs_ne_nil.
 rewrite (total_prob M (duniform xs) xs).
-+ smt(supp_duniform).
++ by smt(supp_duniform).
 apply eq_big_seq => y y_in_xs /=.
 by rewrite duniform1E_uniq // (_ : y \in xs).
 qed.

From b24c56cd9471bbe15f57dce2781ce27f421cabbc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Dupressoir?= <fdupress@gmail.com>
Date: Mon, 13 Apr 2026 15:41:53 +0100
Subject: [PATCH 6/6] swap total_prob parameter, simplify proofs

---
 theories/modules/TotalProb.ec | 66 ++++++++++++++++-------------------
 1 file changed, 31 insertions(+), 35 deletions(-)

diff --git a/theories/modules/TotalProb.ec b/theories/modules/TotalProb.ec
index da9b61bbf..c513cc7d7 100644
--- a/theories/modules/TotalProb.ec
+++ b/theories/modules/TotalProb.ec
@@ -75,7 +75,7 @@ local lemma RandAux_partition_eq (dt' : t distr) (i' : input) (x' : t) &m :
   mu1 dt' x' * Pr[M.main(i', x') @ &m : res].
 proof.
 byphoare
-  (_ :
+  (:
    dt = dt' /\ i = i' /\ (glob M) = (glob M){m} ==>
    res /\ x_of_glob_RA (glob RandAux) = x') => //.
 proc; rewrite /x_of_glob_RA /=.
@@ -86,16 +86,16 @@ seq 1 :
   (1%r - mu1 dt x')
   0%r
   (i = i' /\ (glob M) = (glob M){m}) => //.
-+ auto.
-+ rnd (pred1 x'); auto.
-+ call (_ : i = i' /\ x = x' /\ glob M = (glob M){m} ==> res).
++ by auto.
++ by rnd (pred1 x'); auto.
++ call (: i = i' /\ x = x' /\ glob M = (glob M){m} ==> res).
   + bypr => &hr [#] -> -> glob_eq.
-    byequiv (_ : ={i, x, glob M} ==> ={res}) => //; sim.
-+ auto.
-hoare; call (_ : true); auto; smt().
+    by byequiv (: ={i, x, glob M} ==> ={res}) => //; sim.
++ by auto.
+by hoare; call (: true); auto; smt().
 qed.
 
-lemma total_prob' (dt' : t distr) (supp : t list) (i' : input) &m :
+lemma total_prob' (supp : t list) (dt' : t distr) (i' : input) &m :
   is_finite_for (support dt') supp =>
   Pr[Rand(M).f(dt', i') @ &m : res] =
   big predT
@@ -103,22 +103,22 @@ lemma total_prob' (dt' : t distr) (supp : t list) (i' : input) &m :
   supp.
 proof.
 move => [uniq_supp supp_iff]. 
-have -> :
+have ->:
   Pr[Rand(M).f(dt', i') @ &m : res] = Pr[RandAux.f(dt', i') @ &m : res].
-  + byequiv (_ : ={dt, i, glob M} ==> ={res}) => //; proc; sim.
+  + by byequiv (_ : ={dt, i, glob M} ==> ={res}) => //; proc; sim.
 rewrite (EPL.list_partitioning RandAux (dt', i') E phi supp &m) //.
 have -> /= :
   Pr[RandAux.f(dt', i') @ &m: res /\ ! (x_of_glob_RA (glob RandAux) \in supp)] =
   0%r.
-  + byphoare (_ : dt = dt' /\ i = i' ==> _) => //; proc.
-    seq 1 :
-      (RandAux.x \in supp)
-      1%r
-      0%r
-      0%r
-      1%r => //.
-    + by hoare; call (_ : true); auto; smt().
-    by rnd pred0; auto; smt(mu0).
++ byphoare (: dt = dt' /\ i = i' ==> _) => //; proc.
+  seq 1 :
+    (RandAux.x \in supp)
+    1%r
+    0%r
+    0%r
+    1%r => //.
+  + by hoare; call (: true); auto; smt().
+  by rnd pred0; auto; smt(mu0).
 by congr; apply fun_ext => x'; rewrite RandAux_partition_eq.
 qed.
 
@@ -126,13 +126,13 @@ end section.
 
 (*& total probability lemma for distributions with finite support &*)
 
-lemma total_prob (M <: T) (dt : t distr) (supp : t list) (i : input) &m :
+lemma total_prob (M <: T) (supp : t list) (dt : t distr) (i : input) &m :
   is_finite_for (support dt) supp =>
   Pr[Rand(M).f(dt, i) @ &m : res] =
   big predT
   (fun (x : t) => mu1 dt x * Pr[M.main(i, x) @ &m : res])
   supp.
-proof. move => iff_supp_dt_supp; by apply (total_prob' M). qed.
+proof. exact: (total_prob' M). qed.
 
 end TotalGeneral.
 
@@ -151,9 +151,9 @@ lemma total_prob_bool (M <: T) (i : input) &m :
   Pr[M.main(i, true) @ &m : res]  / 2%r +
   Pr[M.main(i, false) @ &m : res] / 2%r.
 proof.
-rewrite (total_prob M DBool.dbool [true; false]) //.
-+ by smt(supp_dbool).
-by rewrite 2!big_cons big_nil /= /predT /predF /= 2!dbool1E.
+rewrite (total_prob M [true; false]) //.
++ smt(supp_dbool).
+by rewrite 2!big_cons big_nil /= 2!dbool1E.
 qed.
 
 end TotalBool.
@@ -168,18 +168,16 @@ proof *.
 
 (*& total probability lemma for `drange` &*)
 
-lemma total_prob_drange (M <: T) (m n : int, i : input) &m :
+lemma total_prob_drange (M <: T) (m n : int) (i : input) &m :
   m < n =>
   Pr[Rand(M).f(drange m n, i) @ &m : res] =
   bigi predT
   (fun (j : int) => Pr[M.main(i, j) @ &m : res] / (n - m)%r)
   m n.
 proof.
-move => lt_m_n.
-rewrite (total_prob M (drange m n) (range m n)).
+move => lt_m_n; rewrite (total_prob M (range m n)).
 + by rewrite /is_finite_for; smt(range_uniq mem_range supp_drange).
-apply eq_big_seq => x x_in_range_m_n /=.
-by rewrite drange1E (_ : m <= x < n) 1:-mem_range.
+by apply: eq_big_seq=> |> y; rewrite drange1E mem_range=> ->.
 qed.
 
 end TotalRange.
@@ -198,18 +196,16 @@ proof *.
 
 (*& total probability lemma for `duniform` &*)
 
-lemma total_prob_uniform (M <: T) (xs : t list, i : input) &m :
+lemma total_prob_uniform (M <: T) (xs : t list) (i : input) &m :
   uniq xs => xs <> [] =>
   Pr[Rand(M).f(duniform xs, i) @ &m : res] =
   big predT
   (fun (x : t) => Pr[M.main(i, x) @ &m : res] / (size xs)%r)
   xs.
 proof.
-move => uniq_xs xs_ne_nil.
-rewrite (total_prob M (duniform xs) xs).
-+ by smt(supp_duniform).
-apply eq_big_seq => y y_in_xs /=.
-by rewrite duniform1E_uniq // (_ : y \in xs).
+move => uniq_xs xs_ne_nil; rewrite (total_prob M xs).
++ smt(supp_duniform).
+by apply: eq_big_seq=> |> x; rewrite duniform1E_uniq=> // ->.
 qed.
 
 end TotalUniform.