require import AllCore SmtMap List Distr Finite Discrete CV2EC.
require (*--*) PROM.
(*---*) import Distr.MRat.

op b_na : int.
op b_naDDH : int.
op b_naDH9 : int.
op b_nb : int.
op b_nbDDH : int.
op b_nbDH9 : int.

type G.
type Z.

axiom G_fin : finite_type<:G>.
op [lossless] dG : G distr.
axiom Z_fin : finite_type<:Z>.
op [lossless] dZ : Z distr.

op exp : G * Z -> G.
op exp' : G * Z -> G.
op g : G.
op mult : Z * Z -> Z.

axiom multC : commutative mult.

module type Oracles = {
  proc r_ia(_ : int) : unit
  proc r_ib(_ : int) : unit
  proc p_OA(ia : int) : G
  proc p_Oa(ia : int) : Z
  proc p_ODDHa(ia : int, iaDDH : int, m_1 : G, j_1 : int) : bool
  proc p_ODHa9(ia : int, iaDH9 : int, x_2 : Z) : G
  proc p_OB(ib : int) : G
  proc p_Ob(ib : int) : Z
  proc p_ODDHb(ib : int, ibDDH : int, m_2 : G, j_2 : int) : bool
  proc p_ODHb9(ib : int, ibDH9 : int, x_3 : Z) : G
}.

module type Oracles_i = {
  proc init() : unit
  proc unchanged() : (int, Z) fmap * (int, Z) fmap
  include Oracles
}.

(******************************************************************************)
(*                                 LHS module                                 *)
(******************************************************************************)

clone PROM.FullRO as OL_a with
  type in_t    <- int,
  type out_t   <- Z,
  op   dout    <- fun _ => dZ,
  type d_in_t  <- unit,
  type d_out_t <- bool.

clone PROM.FullRO as OL_b with
  type in_t    <- int,
  type out_t   <- Z,
  op   dout    <- fun _ => dZ,
  type d_in_t  <- unit,
  type d_out_t <- bool.

module LHS_ (O_a : OL_a.ROmap) (O_b : OL_b.ROmap) : Oracles_i = {
  var m_r_ib : (int, unit) fmap
  var m_r_ia : (int, unit) fmap
  var m_OA : (int, unit) fmap
  var m_Oa : (int, unit) fmap
  var m_ODDHa : (int * int, G * int) fmap
  var m_ODHa9 : (int * int, Z) fmap
  var m_OB : (int, unit) fmap
  var m_Ob : (int, unit) fmap
  var m_ODDHb : (int * int, G * int) fmap
  var m_ODHb9 : (int * int, Z) fmap

  proc init() = {
    m_r_ib <- empty;
    m_r_ia <- empty;
    m_OA <- empty;
    m_Oa <- empty;
    m_ODDHa <- empty;
    m_ODHa9 <- empty;
    m_OB <- empty;
    m_Ob <- empty;
    m_ODDHb <- empty;
    m_ODHb9 <- empty;
    O_a.init();
    O_b.init();
  }

  proc r_ia(ia : int) = {
    if (1 <= ia <= b_na /\ ia \notin m_r_ia) {
      m_r_ia.[ia] <- ();
      O_a.sample(ia);
    }
  }

  proc r_ib(ib : int) = {
    if (1 <= ib <= b_nb /\ ib \notin m_r_ib) {
      m_r_ib.[ib] <- ();
      O_b.sample(ib);
    }
  }

  proc p_OA(ia : int) = {
    var aout : G <- witness;
    var tmp_a_ia : Z;

    if (1 <= ia <= b_na /\ ia \in m_r_ia /\ ia \notin m_OA) {
      m_OA.[ia] <- ();
      tmp_a_ia <@ O_a.get(ia);
      aout <- exp (g, tmp_a_ia);
    }
    return aout;
  }

  proc p_Oa(ia : int) = {
    var aout : Z <- witness;
    var tmp_a_ia : Z;

    if (1 <= ia <= b_na /\ ia \in m_r_ia /\ ia \notin m_Oa) {
      m_Oa.[ia] <- ();
      tmp_a_ia <@ O_a.get(ia);
      aout <- tmp_a_ia;
    }
    return aout;
  }

  proc p_ODDHa(ia : int, iaDDH : int, m_1 : G, j_1 : int) = {
    var aout : bool <- witness;
    var tmp_a_ia : Z;
    var tmp_b_j_1 : Z;

    if (1 <= ia <= b_na /\ 1 <= iaDDH <= b_naDDH /\ ia \in m_r_ia
        /\ (ia, iaDDH) \notin m_ODDHa /\ 1 <= j_1 <= b_nb /\ j_1 \in m_r_ib) {
      m_ODDHa.[(ia, iaDDH)] <- (m_1, j_1);
      tmp_a_ia <@ O_a.get(ia);
      tmp_b_j_1 <@ O_b.get(j_1);
      aout <- (m_1 = exp (g, mult (tmp_b_j_1, tmp_a_ia)));
    }
    return aout;
  }

  proc p_ODHa9(ia : int, iaDH9 : int, x_2 : Z) = {
    var aout : G <- witness;
    var tmp_a_ia : Z;

    if (1 <= ia <= b_na /\ 1 <= iaDH9 <= b_naDH9 /\ ia \in m_r_ia
        /\ (ia, iaDH9) \notin m_ODHa9) {
      m_ODHa9.[(ia, iaDH9)] <- x_2;
      tmp_a_ia <@ O_a.get(ia);
      aout <- exp (g, mult (tmp_a_ia, x_2));
    }
    return aout;
  }

  proc p_OB(ib : int) = {
    var aout : G <- witness;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ ib \in m_r_ib /\ ib \notin m_OB) {
      m_OB.[ib] <- ();
      tmp_b_ib <@ O_b.get(ib);
      aout <- exp (g, tmp_b_ib);
    }
    return aout;
  }

  proc p_Ob(ib : int) = {
    var aout : Z <- witness;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ ib \in m_r_ib /\ ib \notin m_Ob) {
      m_Ob.[ib] <- ();
      tmp_b_ib <@ O_b.get(ib);
      aout <- tmp_b_ib;
    }
    return aout;
  }

  proc p_ODDHb(ib : int, ibDDH : int, m_2 : G, j_2 : int) = {
    var aout : bool <- witness;
    var tmp_a_j_2 : Z;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ 1 <= ibDDH <= b_nbDDH /\ ib \in m_r_ib
        /\ (ib, ibDDH) \notin m_ODDHb /\ 1 <= j_2 <= b_na /\ j_2 \in m_r_ia) {
      m_ODDHb.[(ib, ibDDH)] <- (m_2, j_2);
      tmp_a_j_2 <@ O_a.get(j_2);
      tmp_b_ib <@ O_b.get(ib);
      aout <- (m_2 = exp (g, mult (tmp_a_j_2, tmp_b_ib)));
    }
    return aout;
  }

  proc p_ODHb9(ib : int, ibDH9 : int, x_3 : Z) = {
    var aout : G <- witness;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ 1 <= ibDH9 <= b_nbDH9 /\ ib \in m_r_ib
        /\ (ib, ibDH9) \notin m_ODHb9) {
      m_ODHb9.[(ib, ibDH9)] <- x_3;
      tmp_b_ib <@ O_b.get(ib);
      aout <- exp (g, mult (tmp_b_ib, x_3));
    }
    return aout;
  }

  proc unchanged() = {
    var tmp_map_a;
    var tmp_map_b;

    tmp_map_a <@ O_a.restrK();
    tmp_map_b <@ O_b.restrK();
    return (tmp_map_a, tmp_map_b);
  }
}.

module LHS = LHS_(OL_a.RO, OL_b.RO).

(******************************************************************************)
(*                                 RHS module                                 *)
(******************************************************************************)

clone PROM.FullRO as OR_a with
  type in_t    <- int,
  type out_t   <- Z,
  op   dout    <- fun _ => dZ,
  type d_in_t  <- unit,
  type d_out_t <- bool.

clone PROM.FullRO as OR_b with
  type in_t    <- int,
  type out_t   <- Z,
  op   dout    <- fun _ => dZ,
  type d_in_t  <- unit,
  type d_out_t <- bool.

module RHS_ (O_a : OR_a.ROmap) (O_b : OR_b.ROmap) : Oracles_i = {
  var v_ka : (int, bool) fmap
  var v_kb : (int, bool) fmap
  var m_r_ib : (int, unit) fmap
  var m_r_ia : (int, unit) fmap
  var m_OA : (int, unit) fmap
  var m_Oa : (int, unit) fmap
  var m_ODDHa : (int * int, G * int) fmap
  var m_ODHa9 : (int * int, Z) fmap
  var m_OB : (int, unit) fmap
  var m_Ob : (int, unit) fmap
  var m_ODDHb : (int * int, G * int) fmap
  var m_ODHb9 : (int * int, Z) fmap

  proc init() = {
    v_ka <- empty;
    v_kb <- empty;
    m_r_ib <- empty;
    m_r_ia <- empty;
    m_OA <- empty;
    m_Oa <- empty;
    m_ODDHa <- empty;
    m_ODHa9 <- empty;
    m_OB <- empty;
    m_Ob <- empty;
    m_ODDHb <- empty;
    m_ODHb9 <- empty;
    O_a.init();
    O_b.init();
  }

  proc r_ia(ia : int) = {
    if (1 <= ia <= b_na /\ ia \notin m_r_ia) {
      m_r_ia.[ia] <- ();
      O_a.sample(ia);
    }
  }

  proc r_ib(ib : int) = {
    if (1 <= ib <= b_nb /\ ib \notin m_r_ib) {
      m_r_ib.[ib] <- ();
      O_b.sample(ib);
    }
  }

  proc p_OA(ia : int) = {
    var aout : G <- witness;
    var tmp_a_ia : Z;

    if (1 <= ia <= b_na /\ ia \in m_r_ia /\ ia \notin m_OA) {
      m_OA.[ia] <- ();
      tmp_a_ia <@ O_a.get(ia);
      aout <- exp' (g, tmp_a_ia);
    }
    return aout;
  }

  proc p_Oa(ia : int) = {
    var aout : Z <- witness;
    var tmp_a_ia : Z;

    if (1 <= ia <= b_na /\ ia \in m_r_ia /\ ia \notin m_Oa) {
      m_Oa.[ia] <- ();
      v_ka.[ia] <- true;
      tmp_a_ia <@ O_a.get(ia);
      aout <- tmp_a_ia;
    }
    return aout;
  }

  proc p_ODDHa(ia : int, iaDDH : int, m_1 : G, j_1 : int) = {
    var aout : bool <- witness;
    var tmp_a_ia : Z;
    var tmp_b_j_1 : Z;

    if (1 <= ia <= b_na /\ 1 <= iaDDH <= b_naDDH /\ ia \in m_r_ia
        /\ (ia, iaDDH) \notin m_ODDHa /\ 1 <= j_1 <= b_nb /\ j_1 \in m_r_ib) {
      m_ODDHa.[(ia, iaDDH)] <- (m_1, j_1);
      if ((j_1 \in v_kb) /\ (true)) {
        tmp_a_ia <@ O_a.get(ia);
        tmp_b_j_1 <@ O_b.get(j_1);
        aout <- (m_1 = exp' (g, mult (tmp_b_j_1, tmp_a_ia)));
      } else {
        if ((ia \in v_ka) /\ (true)) {
          tmp_a_ia <@ O_a.get(ia);
          tmp_b_j_1 <@ O_b.get(j_1);
          aout <- (m_1 = exp' (g, mult (tmp_b_j_1, tmp_a_ia)));
        } else {
          aout <- false;
        }
      }
    }
    return aout;
  }

  proc p_ODHa9(ia : int, iaDH9 : int, x_2 : Z) = {
    var aout : G <- witness;
    var tmp_a_ia : Z;

    if (1 <= ia <= b_na /\ 1 <= iaDH9 <= b_naDH9 /\ ia \in m_r_ia
        /\ (ia, iaDH9) \notin m_ODHa9) {
      m_ODHa9.[(ia, iaDH9)] <- x_2;
      tmp_a_ia <@ O_a.get(ia);
      aout <- exp' (g, mult (tmp_a_ia, x_2));
    }
    return aout;
  }

  proc p_OB(ib : int) = {
    var aout : G <- witness;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ ib \in m_r_ib /\ ib \notin m_OB) {
      m_OB.[ib] <- ();
      tmp_b_ib <@ O_b.get(ib);
      aout <- exp' (g, tmp_b_ib);
    }
    return aout;
  }

  proc p_Ob(ib : int) = {
    var aout : Z <- witness;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ ib \in m_r_ib /\ ib \notin m_Ob) {
      m_Ob.[ib] <- ();
      v_kb.[ib] <- true;
      tmp_b_ib <@ O_b.get(ib);
      aout <- tmp_b_ib;
    }
    return aout;
  }

  proc p_ODDHb(ib : int, ibDDH : int, m_2 : G, j_2 : int) = {
    var aout : bool <- witness;
    var tmp_a_j_2 : Z;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ 1 <= ibDDH <= b_nbDDH /\ ib \in m_r_ib
        /\ (ib, ibDDH) \notin m_ODDHb /\ 1 <= j_2 <= b_na /\ j_2 \in m_r_ia) {
      m_ODDHb.[(ib, ibDDH)] <- (m_2, j_2);
      if ((j_2 \in v_ka) /\ (true)) {
        tmp_a_j_2 <@ O_a.get(j_2);
        tmp_b_ib <@ O_b.get(ib);
        aout <- (m_2 = exp' (g, mult (tmp_a_j_2, tmp_b_ib)));
      } else {
        if ((ib \in v_kb) /\ (true)) {
          tmp_a_j_2 <@ O_a.get(j_2);
          tmp_b_ib <@ O_b.get(ib);
          aout <- (m_2 = exp' (g, mult (tmp_a_j_2, tmp_b_ib)));
        } else {
          aout <- false;
        }
      }
    }
    return aout;
  }

  proc p_ODHb9(ib : int, ibDH9 : int, x_3 : Z) = {
    var aout : G <- witness;
    var tmp_b_ib : Z;

    if (1 <= ib <= b_nb /\ 1 <= ibDH9 <= b_nbDH9 /\ ib \in m_r_ib
        /\ (ib, ibDH9) \notin m_ODHb9) {
      m_ODHb9.[(ib, ibDH9)] <- x_3;
      tmp_b_ib <@ O_b.get(ib);
      aout <- exp' (g, mult (tmp_b_ib, x_3));
    }
    return aout;
  }

  proc unchanged() = {
    var tmp_map_a;
    var tmp_map_b;

    tmp_map_a <@ O_a.restrK();
    tmp_map_b <@ O_b.restrK();
    return (tmp_map_a, tmp_map_b);
  }
}.

module RHS = RHS_(OR_a.RO, OR_b.RO).

(******************************************************************************)
(*                                    Game                                    *)
(******************************************************************************)

module type Adversary (S : Oracles) = {
  proc distinguish() : unit
  proc guess(_ : (int, Z) fmap * (int, Z) fmap) : bool {}
}.

module Game (S : Oracles_i, A : Adversary) = {
  proc main() = {
    var r : bool;
    var rnds;
    S.init();
    A(S).distinguish();
    rnds <@ S.unchanged();
    r <@ A(S).guess(rnds);
    return r;
  }
}.

(*

CV value of the bound:

(#ODDHa + #ODDHb) * (3 * #Oa + 1) * (3 * #Ob + 1) * pCDH(time + (na + nb + (optim-if #Oa = 0 && #Ob = 0 then 0 else #ODDHa + #ODDHb) + 1 + #ODHa9 + #ODHb9) * time(exp)) + (na + nb) * pDistRerandom

Indistinguishability lemma that should be proved
(assuming the current theory has been cloned as CV):

section.

declare module A <: CV.Adversary{-LHS, -RHS}.

declare axiom A_ll : forall (O <: CV.Oracles{-A}),
  islossless O.r_ia =>
  islossless O.r_ib =>
  islossless O.p_OA =>
  islossless O.p_Oa =>
  islossless O.p_ODDHa =>
  islossless O.p_ODHa9 =>
  islossless O.p_OB =>
  islossless O.p_Ob =>
  islossless O.p_ODDHb =>
  islossless O.p_ODHb9 => islossless A(O).distinguish.

declare axiom A_guess_ll : forall (O <: CV.Oracles{-A}),
   islossless A(O).guess.

lemma LHS_RHS &m:
   `| Pr[CV.Game(LHS, A).main() @ &m : res] -
      Pr[CV.Game(RHS, A).main() @ &m : res] | <= <replace with bound>.

end section.
*)