Section OracleHybrid.
Variable A B State S_A : Set.
(* At one point in the proof, we need to know that B is inhabited *)
Variable b : B.
Hypothesis eqdb : EqDec B.
Hypothesis eqdState : EqDec State.
Hypothesis eqdS_A : EqDec S_A.
(* Two oracles, which we want to show are indistinguishable *)
Variable O1 O2 : State -> A -> Comp (B * State).
Hypothesis O1_wf : forall s a, well_formed_comp (O1 s a).
Hypothesis O2_wf : forall s a, well_formed_comp (O2 s a).
(* The "adversary" that attempts to distinguish the oracles using at most q queries. *)
Variable A0 : Comp (State * S_A).
Variable A1 : S_A -> OracleComp A B bool.
Hypothesis A0_wf : well_formed_comp A0.
Hypothesis A1_wf : forall s_A, well_formed_oc (A1 s_A).
Variable q : nat.
Hypothesis A1_qam :
forall s s_A,
In (s, s_A) (getSupport A0) ->
queries_at_most (A1 s_A) q.
(* This proof will show that G1 and G2 are close *)
Definition G1 :=
[s, s_A] <-$2 A0;
[b, _] <-$2 (A1 s_A) _ _ O1 s;
ret b.
Definition G2 :=
[s, s_A] <-$2 A0;
[b, _] <-$2 (A1 s_A) _ _ O2 s;
ret b.
(* The ith oracle responds to the first (i-1) queries using O1, and the remaining queries using O2. *)
Definition Oi (i : nat)(s : (State * nat))(a : A) : Comp (B * (State * nat)) :=
[s_s, s_i] <-2 s;
let O_c := if (ge_dec s_i i) then O2 else O1 in
[b, s_s'] <-$2 O_c s_s a;
ret (b, (s_s', (S s_i))).
(* The ith game uses the ith oracle.  We will show that G1 is the same as (Gi q) and that G2 is the same as (Gi 0).*)
Definition Gi i :=
[s, s_A] <-$2 A0;
[b, _] <-$2 (A1 s_A) _ _ (Oi i) (s, 0%nat);
ret b.
(* We need an assumption that each adjacent pair of games are distant by at most some constant k. *)
Variable k : Rat.
Hypothesis Gi_Si_close :
forall i,
| Pr[Gi i] - Pr[Gi (S i)] | <= k.
(* Step 1: show that G1 is the same as (Gi q).  This is actually the most complicated part of this proof. *)
(* In order to show that G1 is the same as (Gi q), we need an intermediate game that is like G1 except the oracle also keeps track of how many times it was called.  We will use an "identical until bad" proof, in which the "bad" event is that the number of queries is greater than q.  The probability of this event is 0, so we end up with an equivalence. *)
Definition O1_count (s : (State * nat))(a : A) : Comp (B * (State * nat)) :=
[s_s, s_i] <-2 s;
[b, s_s'] <-$2 O1 s_s a;
ret (b, (s_s', (S s_i))).
Definition G1_count  :=
[s, s_A] <-$2 A0;
[b, s] <-$2 (A1 s_A) _ _ O1_count (s, 0%nat);
ret (b, if (gt_dec (snd s) q) then true else false).
(* We also need a game like Gi that exposes the bad event in the same way as G1_count. Then we can show that G1_count and (Gi_count q) are "identical until bad" and the probability of the "bad" event is 0.  *)
Definition Gi_count i :=
[s, s_A] <-$2 A0;
[b, s] <-$2 (A1 s_A) _ _ (Oi i) (s, 0%nat);
ret (b, if (gt_dec (snd s) q) then true else false).
unfold G1, G1_count.
fcf_inline_first.
fcf_to_prhl_eq.
fcf_skip.
fcf_inline_first.
fcf_skip.
eapply (fcf_oracle_eq (fun x y => x = fst y)); trivial; intuition; subst.
unfold O1_count.
fcf_simp.
fcf_ident_expand_l.
fcf_skip.
fcf_spec_ret; intuition.
simpl in *; intuition; subst.
fcf_simp.
simpl.
eapply comp_spec_eq_refl.
Qed.
intuition.
unfold Gi_count, Gi.
fcf_inline_first.
fcf_skip.
fcf_simp.
fcf_inline_first.
fcf_skip.
fcf_simp.
simpl.
intuition.
Qed.

Theorem O1_count_wf :
forall a b,
well_formed_comp (O1_count a b).