[theories/modules] law of total probability

Add module-based law of total probability for distributions
with finite support.

Author: alleystoughton

theories/modules/TotalProb.ec

@@ -0,0 +1,281 @@
+(*^ 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.
+
+(*& theory parameter: lossess distribution on `t` &*)
+
+op [lossless] dt : t distr.
+
+(*& 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 an input `i`,
+    `Rand(M).f(i)` samples a value `x` from the distribution `dt`,
+    passes it to `M.main` along with `i`, and returns the
+    boolean `M.main` returns. &*)
+
+module Rand(M : T) = {
+  proc f(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(i : input) : bool = {
+    var b : bool;
+    x <$ dt;
+    b <@ M.main(i, x);
+    return b;
+  }
+}.
+
+local clone EventPartitioning as EP with
+  type input <- 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 (_ : input) (g : glob RandAux) (_ : bool) : t  = x_of_glob_RA g.
+local op E (_ : input) (_ : glob RandAux) (o : bool) : bool = o.
+
+local lemma RandAux_partition_eq (i' : input) (x' : t) &m :
+  Pr[RandAux.f(i') @ &m : res /\ x_of_glob_RA (glob RandAux) = x'] =
+  mu1 dt x' * Pr[M.main(i', x') @ &m : res].
+proof.
+byphoare
+  (_ :
+   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' (supp : t list) (i' : input) &m :
+  uniq supp => support_is dt supp =>
+  Pr[Rand(M).f(i') @ &m : res] =
+  big predT
+  (fun x' => mu1 dt x' * Pr[M.main(i', x') @ &m : res])
+  supp.
+proof.
+move => uniq_supp support_is_dt_supp.
+have -> :
+  Pr[Rand(M).f(i') @ &m : res] = Pr[RandAux.f(i') @ &m : res].
+  byequiv (_ : ={i, glob M} ==> ={res}) => //; proc; sim.
+rewrite (EPL.list_partitioning RandAux i' E phi supp &m) //.
+have -> /= :
+  Pr[RandAux.f(i') @ &m: res /\ ! (x_of_glob_RA (glob RandAux) \in supp)] =
+  0%r.
+  byphoare => //; 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) (supp : t list) (i : input) &m :
+  uniq supp => support_is dt supp =>
+  Pr[Rand(M).f(i) @ &m : res] =
+  big predT
+  (fun (x : t) => mu1 dt x * Pr[M.main(i, x) @ &m : res])
+  supp.
+proof.
+move => 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,
+  op dt  <- DBool.dbool
+proof *.
+realize dt_ll. apply dbool_ll. qed.
+
+(*& total probability lemma for `DBool.dbool` &*)
+
+lemma total_prob_bool (M <: T) (i : input) &m :
+  Pr[Rand(M).f(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 [true; false]) //.
+smt(supp_dbool).
+by rewrite 2!big_cons big_nil /= /predT /predF /= 2!dbool1E.
+qed.
+
+end TotalBool.
+
+(*& Specialization to `drange` &*)
+
+abstract theory TotalRange.
+
+(*& theory parameter: lower end of range of integers (inclusive) &*)
+
+op m : int.
+
+(*& theory parameter: upper end of range of integers (exclusive) &*)
+
+op n : int.
+
+(*& theory axiom: range is nonempty &*)
+
+axiom lt_m_n : m < n.
+
+clone include TotalGeneral with
+  type t <- int,
+  op dt  <- drange m n
+proof *.
+realize dt_ll. rewrite drange_ll lt_m_n. qed.
+
+lemma big_weight_simp (M <: T) (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 m n` &*)
+
+lemma total_prob_drange (M <: T) (i : input) &m :
+  Pr[Rand(M).f(i) @ &m : res] =
+  bigi predT
+  (fun (j : int) => Pr[M.main(i, j) @ &m : res] / (n - m)%r)
+  m n.
+proof.
+rewrite (total_prob M (range m n)) 1: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.
+
+(*& theory parameter: list of elements from `t` &*)
+
+op xs : t list.
+
+(*& theory axiom: list is nonempty &*)
+
+axiom xs_ne_nil : xs <> [].
+
+(*& theory axiom: list has no duplicates &*)
+
+axiom uniq_xs : uniq xs.
+
+clone include TotalGeneral with
+  type t <- t,
+  op dt  <- duniform xs
+proof *.
+realize dt_ll. by rewrite duniform_ll xs_ne_nil. qed.
+
+lemma big_weight_simp (M <: T) (i : input, ys : t list) &m :
+  (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.
+elim ys => [// | y ys IH mem_impl].
+smt(big_cons uniq_xs duniform1E_uniq).
+qed.
+
+(*& total probability lemma for `duniform xs` &*)
+
+lemma total_prob_uniform (M <: T) (i : input) &m :
+  Pr[Rand(M).f(i) @ &m : res] =
+  big predT
+  (fun (x : t) => Pr[M.main(i, x) @ &m : res] / (size xs)%r)
+  xs.
+proof.
+rewrite (total_prob M xs) 1:uniq_xs //.
+rewrite /support_is => y; smt(supp_duniform).
+by rewrite (big_weight_simp M).
+qed.
+
+end TotalUniform.