require import AllCore Ring.
require import Distr List Aprhl StdRing StdOrder StdBigop.
  (*---*) import IntID IntOrder RField RealOrder.

  (* we define two positive constants eps, del  that we will 
     use to reason about our privacy budget*)

op eps: { real | 0%r <= eps } as ge0_eps.
op del: { real | 0%r <= del } as ge0_delta.

op gs: { int | 0 <= gs } as ge0_gs.


hint exact : ge0_eps.
hint exact : ge0_delta.
hint exact : ge0_gs.

(* used by trivial only if exact *)

(*
   EX 1   
*)

module M1 = {
   proc f (x:int): int = {
     var r;
     x<- x * 2;
     r <$ lap eps x;
    return r;
  }
}.



lemma ex1 : aequiv [ [(4%r*eps) & 0%r]
  M1.f ~ M1.f
 :  (`|x{1} - x{2}|<= 2)  
==> res{2} = res{1} ].
proof.
  proc.
  seq 1 1 :(`|x{1} - x{2}|<= 4 ). 
  wp; toequiv; auto; smt().
  lap 0 4.  
qed.



(*
   EX 2   
*)

module M2 = {
   proc f (x:int): int = {
     var r;
     x<- x ^ 2;
     r <$ lap eps x;
    return r;
  }
}.



lemma ex2 : aequiv [ [0%r & 0%r]
  M2.f ~ M2.f
 :  x{1} = x{2}  
==> res{2} = res{1} ].
proof.
  proc.
  seq 1 1 :(x{1} = x{2}). 
  wp; toequiv; auto; smt().
  lap 0 0.  
qed.



(*
   EX 3   
*)

module M3 = {
  proc f (d:int list , q1 q2 :int list -> int): int = {
      var r;
      var s;
      r <$ lap (eps) (q1 d);
      s <$ lap (eps) (q2 d);
  return r;
  }
}.

pred adj (ms ns : int list) =
  size ms = size ns /\
  (exists (i : int), forall (j : int),
  !(i=j)=> nth 0 ms j = nth 0 ns j).



lemma ex3 : aequiv [ [(3%r*eps) & 0%r]
  M3.f ~ M3.f
    :  (adj d{1} d{2} /\ ={q1} /\  ={q2} /\
        forall d d', (adj d d') => (`|(q1{1} d)-(q1{2} d')| <= 3) /\
        forall d d', (adj d d') => (`|(q2{1} d)-(q2{2} d')| <= 3)
       )  
==> res{2} = res{1} ].
proof.
  proc. simplify.
  seq 1 1 :(adj d{1} d{2} /\ ={r}  /\
           forall d d', (adj d d') => (`|(q2{1} d)-(q2{2} d')| <= 3))
            <[ (3%r* eps) & 0%r ]>. 
   lap 0 3=> //; smt.
   exists* (q2 d){1}; elim* => e1.
   exists* (q2 d){2}; elim* => e2.
  lap (e2-e1) 0. smt.
qed.


lemma ex3' : aequiv [ [(3%r*eps) & 0%r]
  M3.f ~ M3.f
    :  (adj d{1} d{2} /\ ={q1} /\  ={q2} /\
        forall d d', (adj d d') => (`|(q1{1} d)-(q1{2} d')| <= 3) /\
        forall d d', (adj d d') => (`|(q2{1} d)-(q2{2} d')| <= 3)
       )  
==> res{2} = res{1} ].
proof.
  proc. simplify.
  seq 1 1 :(adj d{1} d{2} /\ ={r}  /\
           forall d d', (adj d d') => (`|(q2{1} d)-(q2{2} d')| <= 3))
            <[ (3%r* eps) & 0%r ]>. 
   lap 0 3=> //; smt.
   toequiv. rnd{1}. rnd{2}. auto. smt. 
qed.


(*
   EX 4
*)

module M4 = {
  proc f (d:int list , q :int list -> int): int = {
      var r;
      var s;
      r <$ lap (eps/gs%r) (q d);
      s <$ lap eps r;
  return r;
  }
}.


lemma ex4 : aequiv [ [eps & 0%r]
  M4.f ~ M4.f
    :  (adj d{1} d{2} /\ ={q} /\
        forall d d', (adj d d') => (`|(q{1} d)-(q{2} d')| <= gs)
       )  
==> res{2} = res{1} ].
  proc. simplify.
  seq 1 1 :(adj d{1} d{2} /\ ={r} )
            <[ eps & 0%r ]>. 
  lap 0 gs=> //; smt. 
  lap 0 0=> //; smt.
qed.

(*
   Auxiliary definitions and lemmas
*)

pred eq_in_range (ms ns : int list, i j : int) =
  forall (k : int),
  i <= k <= j => nth 0 ms k = nth 0 ns k.

lemma eq_in_range_succ (ms ns : int list, i j : int) :
  eq_in_range ms ns i j => eq_in_range ms ns (i + 1) j.
proof.
move => eir_i.
rewrite /eq_in_range => k le_iplus1_k_j.
rewrite eir_i /#.
qed.
  

pred adjacent_e (ms ns : int list) (i:int)=
  size ms = size ns /\
   0 <= i < size ms /\
   eq_in_range ms ns 0 (i - 1) /\
   nth 0 ms i <> nth 0 ns i /\
   eq_in_range ms ns (i + 1) (size ms - 1).
pred adjacent (ms ns : int list) = exists i, adjacent_e ms ns i. 

lemma size_eq_adjacent (ms ns : int list) :
  adjacent ms ns => size ms = size ns.
proof.
rewrite /adjacent. rewrite /adjacent_e => [#].
smt().
qed.

lemma adjacent_sub_abs_bound (ms ns : int list, i : int) :
  adjacent_e ms ns i => 0 <= i < size ms =>
  nth 0 ms i <> nth 0 ns i.
proof.  
  rewrite /adjacent_e. smt().  
qed.

lemma adjacent_ne_sub_eq_after (ms ns : int list, i : int) :
  adjacent_e ms ns i => 0 <= i < size ms => nth 0 ms i <> nth 0 ns i =>
  eq_in_range ms ns (i + 1) (size ms - 1).
proof.
rewrite /adjacent_e. smt(). 
qed.

lemma bigops_help : forall n r, 
0<=n => r * n%r = bigi predT (fun (_ : int) => r) 0 n.
proof. 
move => n r H. rewrite sumr_const count_predT size_range ler_maxr. 
     smt().
  rewrite intmulr; auto.
qed.

(*
  End auxiliary definitions and lemmas
*)

(*
  Ex 5
*)


module M5 = {
  proc f (ls : int list, k : int) : int = {
    var s : int <- 0;
    var z : int;
    var i : int <- 0;
    while (i < size ls ) {
      if (nth 0 ls i <= k) { 
        s <- s + 1;
      }
      i <- i + 1;
    }
    z <$ lap (2%r*eps) s;
    return z;
  }
}.


lemma ex5 :
  aequiv
  [[ (2%r*eps) & 0%r]
   M5.f ~ M5.f : adjacent ls{1} ls{2} /\ ={k} ==> res{1} = res{2}].
proof.
proc.
seq 2 2: (adjacent ls{1} ls{2} /\ ={i, s,k} /\ i{1} = 0 /\ s{1} = 0).
auto. toequiv. auto.
seq 1 1 : (`|s{1} - s{2}| <= 1).
toequiv.
while
  (adjacent ls{1} ls{2} /\ ={i, k} /\ 0 <= i{1} <= size ls{1} /\
   (! ={s} =>
    `|s{1} - s{2}| <= 1 /\
     eq_in_range ls{1} ls{2} i{1} (size ls{1} - 1))).
auto;smt().
auto; progress.
rewrite size_ge0.
smt(size_eq_adjacent).
smt(size_eq_adjacent).
smt(addzN).
lap 0 1.
qed.



(* 
  Ex  6
*) 


module M6 = {
  proc f (ls : int list, k:int): int list = {
    var s :int <- 0;
    var output : int list;
    var z : int;
    var i :int <- 0;
    output <- [];
    while (i < size ls ) {
          if (nth 0 ls i <= k) { 
             s <- s + 2;
      }
      i <- i + 1;
      z <$ lap eps s;
      output <- z :: output;
    }
    return output;
  }
}.


lemma ex6 n : 0<=n => aequiv [ [  ((2%r*eps) *n%r ) & 0%r]
  M6.f ~ M6.f
 :  (adjacent ls{1} ls{2} /\  ={k} /\ n = size ls{1})
      ==> res{2} = res{1} ].
    proof.
    move => H.
    proc.
    seq 3 3: (adjacent ls{1} ls{2} /\ ={i, s, k, output} /\ i{1} = 0
              /\ s{1} = 0  /\ 0<= n /\ n = size ls{1} ).
      toequiv; auto.
    awhile  [ (fun _ => (2%r*eps)) & (fun _ => 0%r) ] n [n-i-1] 
    (adjacent ls{1} ls{2} /\ ={i, output, k} /\ 0 <= i{1} <= n /\
     (! ={s} => `|s{1} - s{2}| <= 2 /\
     eq_in_range ls{1} ls{2} i{1} (n - 1)) /\
     0<= n /\ n = size ls{1}).
     auto; smt(ge0_eps).
     auto; smt().
     auto; smt().
     apply bigops_help. trivial.
     rewrite /sumr sumr_const intmulr;auto.
     move => v.
     seq 2 2:
     (adjacent ls{1} ls{2} /\ ={i, output, k} /\ 0 <= i{1} <= size ls{1} /\
     (! ={s} => `|s{1} - s{2}| <= 2 /\
     eq_in_range ls{1} ls{2} i{1} (size ls{1} - 1)) /\
     0<= n /\ n = size ls{1} /\ v=n-i{1}).
     toequiv; auto; progress; try smt().
     seq 1 1:
     (adjacent ls{1} ls{2} /\ ={i, output, k} /\ 0 <= i{1} <= size ls{1} /\
     (! ={s} => `|s{1} - s{2}| <= 2 /\
     eq_in_range ls{1} ls{2} i{1} (size ls{1} - 1)) /\
     0<= n /\ n = size ls{1} /\ v=n-i{1} /\ z{1}=z{2}) <[ (2%r*eps) & 0%r ]>.
     lap 0 2. smt().  
     auto; progress.
     toequiv. auto. progress.
     smt().
     smt().
qed.