From Coq Require Import ssreflect.
From stdpp Require Import countable gmap list.

Lemma list_to_map_lookup_is_Some {A B} `{Countable A, EqDecision A} (l: list (A * B)) (a: A) :
  is_Some ((list_to_map l : gmap A B) !! a) ↔ a ∈ l.*1.
Proof.
  induction l.
  - cbn. split; by inversion 1.
  - cbn. rewrite lookup_insert_is_Some' elem_of_cons.
    split; intros [HH|HH]; eauto; rewrite -> IHl in *; auto.
Qed.

Lemma zip_app {A B} (l1 l1': list A) (l2 l2' : list B) :
  length l1 = length l2 ->
  zip (l1 ++ l1') (l2 ++ l2') = zip l1 l2 ++ zip l1' l2'.
Proof.
  revert l2. induction l1;intros l2 Hlen.
  - destruct l2;[|inversion Hlen]. done.
  - destruct l2;[inversion Hlen|]. simpl.
    f_equiv. auto.
Qed.

Lemma length_zip_l {A B} (l1: list A) (l2: list B) :
  length l1 ≤ length l2 → length (zip l1 l2) = length l1.
Proof.
  revert l2. induction l1; intros l2 Hl2; auto.
  destruct l2; cbn in Hl2. exfalso; lia.
  cbn. rewrite IHl1; auto. lia.
Qed.

Lemma list_filter_forall { A: Type } (P: A -> Prop) `{ forall x, Decision (P x) } l:
  Forall P l ->
  @list_filter _ P _ l = l.
Proof.
  induction 1; auto.
  simpl. destruct (decide (P x)); rewrite /filter; try congruence.
Qed.

Lemma dom_map_imap_full {K A B}
      `{Countable A, EqDecision A, Countable B, EqDecision B, Countable K, EqDecision K}
      (f: K -> A -> option B) (m: gmap K A):
  (∀ k a, m !! k = Some a → is_Some (f k a)) →
  dom (map_imap f m) = dom m.
Proof.
  intros Hf.
  apply set_eq. intros k.
  rewrite !elem_of_dom map_lookup_imap.
  destruct (m !! k) eqn:Hmk.
  - destruct (Hf k a Hmk) as [? Hfk]. cbn. rewrite Hfk. split; eauto.
  - cbn. split; inversion 1; congruence.
Qed.

Lemma dom_list_to_map_singleton {K V: Type} `{EqDecision K, Countable K} (x:K) (y:V):
  dom (list_to_map [(x, y)] : gmap K V) = list_to_set [x].
Proof. rewrite dom_insert_L /= dom_empty_L. set_solver. Qed.

Lemma list_to_set_disj {A} `{Countable A, EqDecision A} (l1 l2: list A) :
  l1 ## l2 → (list_to_set l1: gset A) ## list_to_set l2.
Proof.
  intros * HH. rewrite elem_of_disjoint. intros x.
  rewrite !elem_of_list_to_set. rewrite elem_of_disjoint in HH |- *. eauto.
Qed.

Lemma map_to_list_fst {A B : Type} `{EqDecision A, Countable A} (m : gmap A B) i :
  i ∈ (map_to_list m).*1 ↔ (∃ x, (i,x) ∈ (map_to_list m)).
Proof.
  split.
  - intros Hi.
    destruct (m !! i) eqn:Hsome.
    + exists b. by apply elem_of_map_to_list.
    + rewrite -(list_to_map_to_list m) in Hsome.
      eapply not_elem_of_list_to_map in Hsome. done.
  - intros [x Hix].
    apply elem_of_list_fmap.
    exists (i,x). auto.
Qed.

Lemma drop_S':
    forall A l n (a: A) l',
      drop n l = a::l' ->
      drop (S n) l = l'.
Proof.
  induction l; intros * HH.
  - rewrite drop_nil in HH. inversion HH.
  - simpl. destruct n.
    + rewrite drop_0 in HH. inversion HH.
      reflexivity.
    + simpl in HH. eapply IHl; eauto.
Qed.

Lemma disjoint_nil_l {A : Type} `{EqDecision A} (a : A) (l2 : list A) :
  [] ## l2.
Proof.
  apply elem_of_disjoint. intros x Hcontr. inversion Hcontr.
Qed.

Lemma disjoint_nil_r {A : Type} `{EqDecision A} (a : A) (l2 : list A) :
  l2 ## [].
Proof.
  apply elem_of_disjoint. intros x Hl Hcontr. inversion Hcontr.
Qed.

Lemma disjoint_cons {A : Type} `{EqDecision A} (a : A) (l1 l2 : list A) :
  a :: l1 ## l2 → a ∉ l2.
Proof.
  rewrite elem_of_disjoint =>Ha.
  assert (a ∈ a :: l1) as Hs; [apply elem_of_cons;auto;apply elem_of_nil|].
  specialize (Ha a Hs). done.
Qed.

Lemma disjoint_weak {A : Type} `{EqDecision A} (a : A) (l1 l2 : list A) :
  a :: l1 ## l2 → l1 ## l2.
Proof.
  rewrite elem_of_disjoint =>Ha a' Hl1 Hl2.
  assert (a' ∈ a :: l1) as Hs; [apply elem_of_cons;auto;apply elem_of_nil|].
  specialize (Ha a' Hs Hl2). done.
Qed.

Lemma disjoint_swap {A : Type} `{EqDecision A} (a : A) (l1 l2 : list A) :
  a ∉ l1 →
  a :: l1 ## l2 -> l1 ## a :: l2.
Proof.
  rewrite elem_of_disjoint =>Hnin Ha a' Hl1 Hl2.
  destruct (decide (a' = a)).
  - subst. contradiction.
  - apply Ha with a'.
    + apply elem_of_cons; by right.
    + by apply elem_of_cons in Hl2 as [Hcontr | Hl2]; [contradiction|].
Qed.

(* delete_list: delete a list of keys in a map *)

Fixpoint delete_list {K V : Type} `{Countable K, EqDecision K}
           (ks : list K) (m : gmap K V) : gmap K V :=
  match ks with
  | k :: ks' => delete k (delete_list ks' m)
  | [] => m
  end.

Lemma delete_list_insert {K V : Type} `{Countable K, EqDecision K}
      (ks : list K) (m : gmap K V) (l : K) (v : V) :
  l ∉ ks →
  delete_list ks (<[l:=v]> m) = <[l:=v]> (delete_list ks m).
Proof.
  intros Hnin.
  induction ks; auto.
  simpl.
  apply not_elem_of_cons in Hnin as [Hneq Hnin].
  rewrite -delete_insert_ne; auto.
  f_equal. by apply IHks.
Qed.

Lemma delete_list_delete {K V : Type} `{Countable K, EqDecision K}
      (ks : list K) (m : gmap K V) (l : K) :
  l ∉ ks →
  delete_list ks (delete l m) = delete l (delete_list ks m).
Proof.
  intros Hnin.
  induction ks; auto.
  simpl.
  apply not_elem_of_cons in Hnin as [Hneq Hnin].
  rewrite -delete_commute; auto.
  f_equal. by apply IHks.
Qed.

Lemma lookup_delete_list_notin {K V : Type} `{Countable K, EqDecision K}
      (ks : list K) (m : gmap K V) (l : K) :
  l ∉ ks →
  (delete_list ks m) !! l = m !! l.
Proof.
  intros HH; induction ks; simpl; auto.
  eapply not_elem_of_cons in HH. destruct HH.
  rewrite lookup_delete_ne; auto.
Qed.

Lemma delete_list_permutation {A B} `{Countable A, EqDecision A}
      (l1 l2: list A) (m: gmap A B):
  l1 ≡ₚ l2 → delete_list l1 m = delete_list l2 m.
Proof.
  induction 1.
  { reflexivity. }
  { cbn. rewrite IHPermutation //. }
  { cbn. rewrite delete_commute //. }
  { rewrite IHPermutation1 //. }
Qed.

Lemma delete_list_swap {A B : Type} `{EqDecision A, Countable A}
      (a a' : A) (l1 l2 : list A) (M : gmap A B) :
  delete a' (delete_list (l1 ++ a :: l2) M) =
  delete a (delete a' (delete_list (l1 ++ l2) M)).
Proof.
  induction l1.
  - apply delete_commute.
  - simpl. repeat rewrite (delete_commute _ _ a0).
    f_equiv. apply IHl1.
Qed.

(* Map difference for heterogeneous maps, and lemmas relating it to delete_list *)

Definition map_difference_het
  {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
  (m1: gmap A B) (m2: gmap A C): gmap A B
:=
  filter (fun '(k, v) => m2 !! k = None) m1.

Notation "m1 ∖∖ m2" := (map_difference_het m1 m2) (at level 40, left associativity).

Lemma map_eq' {A B} `{Countable A, EqDecision A, Countable B, EqDecision B}
  (m1 m2: gmap A B):
  m1 = m2 ↔ (forall k v, m1 !! k = Some v ↔ m2 !! k = Some v).
Proof.
  split. intros ->. done.
  intros Heq. apply map_eq. intro k. destruct (m2 !! k) eqn:HH.
  { by apply Heq. }
  { destruct (m1 !! k) eqn:HHH; auto. apply Heq in HHH. congruence. }
Qed.

Lemma difference_het_lookup_Some
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C) (k: A) (v: B):
  (m1 ∖∖ m2) !! k = Some v ↔ m1 !! k = Some v ∧ m2 !! k = None.
Proof. by rewrite /map_difference_het map_filter_lookup_Some. Qed.

Lemma difference_het_lookup_None
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C) (k: A) (v: B):
  (m1 ∖∖ m2) !! k = None ↔ m1 !! k = None ∨ is_Some (m2 !! k).
Proof.
  rewrite /map_difference_het map_filter_lookup_None.
  split; intros [HH1|HH2]; eauto.
  { destruct (m1 !! k) eqn:?; eauto; right.
    destruct (m2 !! k) eqn:?; eauto. exfalso. eapply HH2; eauto. }
  { destruct (m1 !! k) eqn:?; eauto; right.
    destruct (m2 !! k) eqn:?; eauto. destruct HH2 as [? ?]. congruence. }
Qed.

Lemma difference_het_empty
  {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
  (m: gmap A B):
  m ∖∖ (∅ : gmap A C) = m.
Proof.
  rewrite /map_difference_het map_eq'. intros k v.
  rewrite map_filter_lookup_Some. rewrite lookup_empty. set_solver.
Qed.

 Lemma difference_het_eq_empty
  {A B} `{Countable A, EqDecision A, Countable B, EqDecision B}
  (m: gmap A B):
  m ∖∖ m = (∅ : gmap A B).
Proof.
  rewrite /map_difference_het map_eq'. intros k v.
  rewrite map_filter_lookup_Some. rewrite lookup_empty. set_solver.
Qed.

Lemma difference_het_insert_r
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C) (k: A) (v: C):
  m1 ∖∖ (<[ k := v ]> m2) = delete k (m1 ∖∖ m2).
Proof.
  intros.
  rewrite /map_difference_het map_eq'. intros k' v'.
  rewrite map_filter_lookup_Some lookup_delete_Some.
  rewrite map_filter_lookup_Some lookup_insert_None. set_solver.
Qed.

Lemma difference_het_insert_l
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C) (k: A) (v: B):
  m2 !! k = None ->
  <[ k := v ]> m1 ∖∖ m2 = <[ k := v ]> (m1 ∖∖ m2).
Proof.
  intros.
  rewrite /map_difference_het map_eq'. intros k' v'.
  rewrite map_filter_lookup_Some lookup_insert_Some.
  rewrite -map_filter_insert_True;auto.
    by rewrite map_filter_lookup_Some lookup_insert_Some.
Qed.

Lemma difference_het_delete_assoc
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C) (k: A):
  delete k (m1 ∖∖ m2) = (delete k m1) ∖∖ m2.
Proof.
  intros.
  rewrite /map_difference_het map_eq'. intros k' v'.
  rewrite map_filter_lookup_Some.
  rewrite -map_filter_delete;auto.
  rewrite map_filter_lookup_Some. set_solver.
Qed.

Lemma dom_difference_het
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C):
  dom (m1 ∖∖ m2) = dom m1 ∖ dom m2.
Proof.
  apply (@anti_symm _ _ subseteq).
  typeclasses eauto.
  { rewrite elem_of_subseteq. intro k.
    rewrite elem_of_dom. intros [v Hv].
    rewrite difference_het_lookup_Some in Hv *.
    destruct Hv as [? ?].
    rewrite elem_of_difference !elem_of_dom. split; eauto.
    intros [? ?]. congruence. }
  { rewrite elem_of_subseteq. intro k.
    rewrite elem_of_difference !elem_of_dom. intros [[v ?] Hcontra].
    exists v. rewrite difference_het_lookup_Some. split; eauto.
    destruct (m2 !! k) eqn:?; eauto. exfalso. apply Hcontra. eauto. }
Qed.

Lemma delete_elements_eq_difference_het
    {A B C} `{Countable A, EqDecision A, Countable B, EqDecision B}
    (m1: gmap A B) (m2: gmap A C):
  delete_list (elements (dom m2)) m1 = m1 ∖∖ m2.
Proof.
  set (l := elements (dom m2)).
  assert (l ≡ₚ elements (dom m2)) as Hl by reflexivity.
  clearbody l. revert l Hl. revert m1. pattern m2. revert m2.
  apply map_ind.
  - intros m1 l. rewrite dom_empty_L elements_empty difference_het_empty.
    rewrite Permutation_nil_r. intros ->. reflexivity.
  - intros k v m2 Hm2k HI m1 l Hm1l.
    rewrite difference_het_insert_r.
    rewrite dom_insert in Hm1l *.
    move: Hm1l. rewrite elements_union_singleton.
    rewrite elem_of_dom; intros [? ?]; congruence.
    intros Hm1l.
    transitivity (delete k (delete_list (elements (dom m2)) m1)).
    { erewrite delete_list_permutation. 2: eauto. reflexivity. }
    { rewrite HI//. }
Qed.

(* rtc *)

Lemma rtc_implies {A : Type} (R Q : A → A → Prop) (x y : A) :
  (∀ r q, R r q → Q r q) →
  rtc R x y → rtc Q x y.
Proof.
  intros Himpl HR.
  induction HR.
  - done.
  - apply Himpl in H.
    apply rtc_once in H.
    apply rtc_transitive with y; auto.
Qed.

Lemma rtc_or_intro {A : Type} (R Q : A → A → Prop) (x y : A) :
  rtc (λ a b, R a b) x y →
  rtc (λ a b, R a b ∨ Q a b) x y.
Proof.
  intros HR. induction HR.
  - done.
  - apply rtc_transitive with y; auto.
    apply rtc_once. by left.
Qed.

Lemma rtc_or_intro_l {A : Type} (R Q : A → A → Prop) (x y : A) :
    rtc (λ a b, R a b) x y →
    rtc (λ a b, Q a b ∨ R a b) x y.
Proof.
  intros HR. induction HR.
  - done.
  - apply rtc_transitive with y; auto.
    apply rtc_once. by right.
Qed.

 (* Helper lemmas on list differences *)

Lemma not_elem_of_list {A : Type} `{EqDecision A} (a : A) (l x : list A) :
  a ∈ x → a ∉ list_difference l x.
Proof.
  intros Hax.
  rewrite /not.
  intros Hal.
  by apply elem_of_list_difference in Hal as [Ha' Hax_not].
Qed.

Lemma list_difference_nil {A : Type} `{EqDecision A} (l : list A) :
  list_difference l [] = l.
Proof.
  induction l; auto.
  simpl. f_equal.
  apply IHl.
Qed.

Lemma list_difference_length_cons {A : Type} `{EqDecision A}
      (l2 : list A) (a : A) :
  list_difference [a] (a :: l2) = [].
Proof.
  simpl.
  assert (a ∈ a :: l2); first apply elem_of_list_here.
  destruct (decide_rel elem_of a (a :: l2)); auto; last contradiction.
Qed.

Lemma list_difference_skip {A : Type} `{EqDecision A}
      (l1 l2 : list A) (b : A) :
  ¬ (b ∈ l1) →
  list_difference l1 (b :: l2) = list_difference l1 l2.
Proof.
  intros Hnin.
  induction l1; auto.
  apply not_elem_of_cons in Hnin.
  destruct Hnin as [Hne Hl1].
  simpl.
  destruct (decide_rel elem_of a (b :: l2)).
  - apply elem_of_cons in e.
    destruct e as [Hcontr | Hl2]; first congruence.
    destruct (decide_rel elem_of a l2); last contradiction.
      by apply IHl1.
  - apply not_elem_of_cons in n.
    destruct n as [Hne' Hl2].
    destruct (decide_rel elem_of a l2); first contradiction.
    f_equal.
      by apply IHl1.
Qed.

Lemma list_difference_nested {A : Type} `{EqDecision A}
      (l1 l1' l2 : list A) (b : A) :
  ¬ (b ∈ (l1 ++ l1')) →
  list_difference (l1 ++ b :: l1') (b :: l2) = list_difference (l1 ++ l1') l2.
Proof.
  intros Hnotin.
  induction l1.
  - simpl.
    assert (b ∈ (b :: l2)); first apply elem_of_list_here.
    destruct (decide_rel elem_of b (b :: l2)); last contradiction.
    rewrite list_difference_skip; auto.
  - simpl in *.
    apply not_elem_of_cons in Hnotin.
    destruct Hnotin as [Hne Hnotin].
    destruct (decide_rel elem_of a (b :: l2)).
    + apply elem_of_cons in e.
      destruct e as [Hcontr | Hl2]; first congruence.
      destruct (decide_rel elem_of a l2); last contradiction.
        by apply IHl1.
    + apply not_elem_of_cons in n.
      destruct n as [Hne' Hnotin'].
      destruct (decide_rel elem_of a l2); first contradiction.
      f_equal.
        by apply IHl1.
Qed.

Lemma list_difference_length_ni  {A : Type} `{EqDecision A}
      (l1 : list A) (b : A) :
  ¬ (b ∈ l1) →
  length (list_difference l1 [b]) = length l1.
Proof.
  intros Hna.
  destruct l1; auto.
  simpl.
  apply not_elem_of_cons in Hna.
  destruct Hna as [Hne Hna].
  destruct (decide_rel elem_of a [b]).
  - apply elem_of_list_singleton in e. congruence.
  - simpl. rewrite list_difference_skip; auto.
      by rewrite list_difference_nil.
Qed.

Lemma list_difference_single_length {A : Type} `{EqDecision A}
      (l1 : list A) (b : A) :
  b ∈ l1 →
  NoDup l1 →
  length (list_difference l1 [b]) =
  length l1 - 1.
Proof.
  intros Ha Hndup.
  induction l1; auto.
  destruct (decide (b = a)).
  - subst.
    assert (a ∈ a :: l1); first apply elem_of_list_here.
    rewrite ->NoDup_cons in Hndup. destruct Hndup as [Hni Hndup].
    assert (¬ (a ∈ l1)) as Hni'.
    { rewrite /not. intros Hin. contradiction. }
    simpl.
    assert (a ∈ [a]); first apply elem_of_list_here.
    destruct (decide_rel elem_of a [a]); last contradiction.
    rewrite Nat.sub_0_r.
    apply list_difference_length_ni; auto.
  - simpl.
    assert (¬ (a ∈ [b])).
    { rewrite /not. intros Hin. apply elem_of_list_singleton in Hin. congruence. }
    destruct (decide_rel elem_of a [b]); first contradiction.
    rewrite Nat.sub_0_r /=.
    inversion Hndup; subst.
    apply elem_of_cons in Ha.
    destruct Ha as [Hcontr | Ha]; first congruence.
    apply IHl1 in Ha as Heq; auto.
    rewrite Heq.
    destruct l1; first inversion Ha.
    simpl. lia.
Qed.

Lemma list_difference_app {A : Type} `{EqDecision A}
      (l1 l2 l2' : list A) :
  list_difference l1 (l2 ++ l2') = list_difference (list_difference l1 l2) l2'.
Proof.
  induction l1; auto.
  simpl. destruct (decide_rel elem_of a (l2 ++ l2')).
  - apply elem_of_app in e as [Hl2 | Hl2'].
    + destruct (decide_rel elem_of a l2); last contradiction.
      apply IHl1.
    + destruct (decide_rel elem_of a l2); first by apply IHl1.
      simpl.
      destruct (decide_rel elem_of a l2'); last contradiction.
      apply IHl1.
  - apply not_elem_of_app in n as [Hl2 Hl2'].
    destruct (decide_rel elem_of a l2); first contradiction.
    simpl.
    destruct (decide_rel elem_of a l2'); first contradiction.
    f_equal. apply IHl1.
Qed.

Lemma list_difference_Permutation {A : Type} `{EqDecision A} (l l1 l2 : list A) :
  l1 ≡ₚ l2 -> list_difference l l1 = list_difference l l2.
Proof.
  intros Hl.
  induction l; auto.
  simpl. rewrite IHl.
  destruct (decide_rel elem_of a l1).
  - apply elem_of_list_In in e.
    apply Permutation_in with _ _ _ a in Hl; auto.
    apply elem_of_list_In in Hl.
    destruct (decide_rel elem_of a l2);[|contradiction].
    done.
  - revert n; rewrite elem_of_list_In; intros n.
    assert (¬ In a l2) as Hnin.
    { intros Hcontr. apply Permutation_sym in Hl.
      apply Permutation_in with _ _ _ a in Hl; auto. }
    revert Hnin; rewrite -elem_of_list_In; intro Hnin.
    destruct (decide_rel elem_of a l2);[contradiction|].
    done.
Qed.

Lemma list_difference_length {A} `{EqDecision A} (l1 : list A) :
  forall l2, NoDup l1 -> NoDup l2 -> l2 ⊆+ l1 ->
  length (list_difference l1 l2) = length (l1) - length l2.
Proof.
  induction l1; intros l2 Hdup1 Hdup2 Hsub.
  - simpl. done.
  - simpl. destruct (decide_rel elem_of a l2).
    + apply submseteq_cons_r in Hsub as [Hcontr | [k [Hperm Hk] ] ].
      { apply elem_of_submseteq with (x:=a) in Hcontr;auto. apply NoDup_cons in Hdup1 as [Hnin ?].
        exfalso. by apply Hnin. }
      apply list_difference_Permutation with (l:=l1) in Hperm as Heq. rewrite Heq.
      apply NoDup_cons in Hdup1 as [Hnin ?].
      rewrite list_difference_skip; [intros Hcontr;by apply Hnin|].
      rewrite IHl1;auto.
      revert Hdup2. rewrite Hperm =>Hdup2. by apply NoDup_cons in Hdup2 as [? ?].
      rewrite Hperm /=. auto.
    + simpl. apply submseteq_cons_r in Hsub as [Hsub | Hcontr].
      rewrite IHl1;auto. assert (length l2 ≤ length l1).
      { apply submseteq_length. auto. }
      by apply NoDup_cons in Hdup1 as [? ?]; auto.
      by apply submseteq_length in Hsub; lia.
      destruct Hcontr as [l' [Hperm Hl'] ].
      exfalso. apply n. rewrite Hperm. constructor.
Qed.

Lemma list_to_set_difference A {_: EqDecision A} {_: Countable A} (l1 l2: list A):
  (list_to_set (list_difference l1 l2): gset A) = (list_to_set l1: gset A) ∖ (list_to_set l2: gset A).
Proof.
  revert l2. induction l1.
  - intro. cbn [list_difference list_to_set]. set_solver.
  - intros l2. cbn [list_difference list_to_set]. destruct (decide_rel elem_of a l2); set_solver.
Qed.

(* creates a gmap with domain from the list, all pointing to a default value *)
Fixpoint create_gmap_default {K V : Type} `{Countable K}
         (l : list K) (d : V) : gmap K V :=
  match l with
  | [] => ∅
  | k :: tl => <[k:=d]> (create_gmap_default tl d)
  end.

Lemma create_gmap_default_lookup {K V : Type} `{Countable K}
      (l : list K) (d : V) (k : K) :
  k ∈ l ↔ (create_gmap_default l d) !! k = Some d.
Proof.
  split.
  - intros Hk.
    induction l; inversion Hk.
    + by rewrite lookup_insert.
    + destruct (decide (a = k)); [subst; by rewrite lookup_insert|].
      rewrite lookup_insert_ne; auto.
  - intros Hl.
    induction l; inversion Hl.
    destruct (decide (a = k)); [subst;apply elem_of_list_here|].
    apply elem_of_cons. right.
    apply IHl. simplify_map_eq. auto.
Qed.

Lemma create_gmap_default_lookup_is_Some {K V} `{EqDecision K, Countable K} (l: list K) (d: V) x v:
  create_gmap_default l d !! x = Some v → x ∈ l ∧ v = d.
Proof.
  revert x v d. induction l as [| a l]; cbn.
  - done.
  - intros x v d. destruct (decide (a = x)) as [->|].
    + rewrite lookup_insert. intros; simplify_eq. repeat constructor.
    + rewrite lookup_insert_ne //. intros [? ?]%IHl. subst. repeat constructor; auto.
Qed.

Lemma create_gmap_default_dom {K V} `{EqDecision K, Countable K} (l: list K) (d: V):
  dom (create_gmap_default l d) = list_to_set l.
Proof.
  induction l as [| a l].
  - cbn. rewrite dom_empty_L //.
  - cbn [create_gmap_default list_to_set]. rewrite dom_insert_L // IHl //.
Qed.

Lemma fst_zip_prefix A B (l : list A) (k : list B) :
  (zip l k).*1 `prefix_of` l.
Proof.
  revert k. induction l; cbn; auto.
  destruct k; cbn.
  - apply prefix_nil.
  - apply prefix_cons; auto.
Qed.

Lemma prefix_of_nil A (l : list A) :
  l `prefix_of` [] →
  l = [].
Proof. destruct l; auto. by intros ?%prefix_nil_not. Qed.

Lemma in_prefix A (l1 l2 : list A) x :
  l1 `prefix_of` l2 →
  x ∈ l1 →
  x ∈ l2.
Proof.
  unfold prefix. intros [? ->] ?.
  apply elem_of_app. eauto.
Qed.

Lemma NoDup_prefix A (l1 l2 : list A) :
  NoDup l2 →
  l1 `prefix_of` l2 →
  NoDup l1.
Proof.
  intros H. revert l1. induction H.
  - intros * ->%prefix_of_nil. constructor.
  - intros l1. destruct l1.
    + intros _. constructor.
    + intros HH. rewrite (prefix_cons_inv_1 _ _ _ _ HH).
      apply prefix_cons_inv_2 in HH. constructor; eauto.
      intro Hx. pose proof (in_prefix _ _ _ _ HH Hx). done.
Qed.

Lemma take_lookup_Some_inv A (l : list A) (n i : nat) x :
  take n l !! i = Some x →
  i < n ∧ l !! i = Some x.
Proof.
  revert l i x. induction n; cbn.
  { intros *. inversion 1. }
  { intros *. destruct l; cbn. by inversion 1. destruct i; cbn.
    - intros; simplify_eq. split; auto. lia.
    - intros [? ?]%IHn. split. lia. auto. }
Qed.

Lemma NoDup_fst {A B : Type} (l : list (A*B)) :
  NoDup l.*1 -> NoDup l.
Proof.
  intros Hdup.
  induction l.
  - by apply NoDup_nil.
  - destruct a. simpl in Hdup. apply NoDup_cons in Hdup as [Hin Hdup].
    apply NoDup_cons. split;auto.
    intros Hcontr. apply Hin. apply elem_of_list_fmap.
    exists (a,b). simpl. split;auto.
Qed.

Lemma fst_elem_of_cons {A B} `{EqDecision A} (l : list A) (x : A) (l': list B) :
  x ∈ (zip l l').*1 →
  x ∈ l.
Proof. intros H. eapply in_prefix. eapply fst_zip_prefix. done. Qed.

Lemma length_fst_snd {A B} `{Countable A} (m : gmap A B) :
  length (map_to_list m).*1 = length (map_to_list m).*2.
Proof.
  induction m using map_ind.
  - rewrite map_to_list_empty. auto.
  - rewrite map_to_list_insert;auto. simpl. auto.
Qed.

Lemma map_to_list_delete {A B} `{Countable A} `{EqDecision A} (m : gmap A B) (i : A) (x : B) :
  ∀ l, (i,x) :: l ≡ₚmap_to_list m ->
       NoDup (i :: l.*1) →
       (map_to_list (delete i m)) ≡ₚl.
Proof.
  intros l Hl Hdup.
  assert ((i,x) ∈ map_to_list m) as Hin.
  { rewrite -Hl. constructor. }
  assert (m !! i = Some x) as Hsome.
  { apply elem_of_map_to_list; auto. }
  apply NoDup_Permutation;auto.
  by apply NoDup_map_to_list.
  apply NoDup_fst. apply NoDup_cons in Hdup as [? ?]. by auto.
  intros [i0 x0]. split.
  - intros Hinx%elem_of_map_to_list.
    assert (i ≠ i0) as Hne;[intros Hcontr;subst;simplify_map_eq|simplify_map_eq].
    assert ((i0, x0) ∈ (i, x) :: l) as Hin'.
    { rewrite Hl. apply elem_of_map_to_list. auto. }
    apply elem_of_cons in Hin' as [Hcontr | Hin'];auto.
    simplify_eq.
  - intros Hinx. apply elem_of_map_to_list.
    assert (i ≠ i0) as Hne;[|simplify_map_eq].
    { intros Hcontr;subst.
      assert (NoDup ((i0, x) :: l)) as Hdup'.
      { rewrite Hl. apply NoDup_map_to_list. }
      assert (i0 ∈ l.*1) as HWInt.
      { apply elem_of_list_fmap. exists (i0,x0). simpl. split;auto. }
      apply NoDup_cons in Hdup as [Hcontr ?]. by apply Hcontr.
    }
    assert ((i0, x0) ∈ (i, x) :: l) as Hin'.
    { constructor. auto. }
    revert Hin'. rewrite Hl =>Hin'. apply elem_of_map_to_list in Hin'.
    auto.
Qed.

Lemma NoDup_map_to_list_fst (A B : Type) `{EqDecision A} `{Countable A}
       (m : gmap A B):
  NoDup (map_to_list m).*1.
Proof.
  induction m as [|i x m] using map_ind.
  - rewrite map_to_list_empty. simpl. by apply NoDup_nil.
  - rewrite map_to_list_insert;auto.
    simpl. rewrite NoDup_cons. split.
    + intros Hcontr%elem_of_list_fmap.
      destruct Hcontr as [ab [Heqab Hcontr] ].
      destruct ab as [a b]. subst. simpl in *.
      apply elem_of_map_to_list in Hcontr. rewrite Hcontr in H0. inversion H0.
    + auto.
Qed.

Lemma map_to_list_delete_fst {A B} `{Countable A} (m : gmap A B) (i : A) (x : B) :
  ∀ l, i :: l ≡ₚ(map_to_list m).*1 ->
       NoDup (i :: l) →
       (map_to_list (delete i m)).*1 ≡ₚl.
Proof.
  intros l Hl Hdup.
  assert (i ∈ (map_to_list m).*1) as Hin.
  { rewrite -Hl. constructor. }
  apply NoDup_cons in Hdup as [Hnin Hdup].
  apply NoDup_Permutation;auto.
  apply NoDup_map_to_list_fst. done.
  set l' := zip l (repeat x (length l)).
  assert (l = l'.*1) as Heq;[rewrite fst_zip;auto;rewrite repeat_length;lia|].
  intros i0. split.
  - intros Hinx%elem_of_list_fmap.
    destruct Hinx as [ [? ?] [? Hinx] ]. simpl in *. subst a.
    apply elem_of_map_to_list in Hinx.
    destruct (decide (i = i0));[subst i;rewrite lookup_delete in Hinx;inversion Hinx|].
    rewrite lookup_delete_ne in Hinx;auto.
    apply elem_of_map_to_list in Hinx.
    assert (i0 ∈ (map_to_list m).*1) as Hinx'.
    { apply elem_of_list_fmap. exists (i0,b). split;auto. }
    revert Hinx'. rewrite -Hl =>Hinx'.
    by apply elem_of_cons in Hinx' as [Hcontr | Hinx'];[congruence|].
  - intros Hinx. assert (i ≠ i0) as Hne;[congruence|simplify_map_eq].
    assert (i0 ∈ i :: l) as Hin'.
    { constructor. auto. }
    revert Hin'. rewrite Hl =>Hin'.
    apply map_to_list_fst in Hin' as [x' Hx].
    apply elem_of_map_to_list in Hx.
    apply map_to_list_fst. exists x'. apply elem_of_map_to_list.
    rewrite lookup_delete_ne;auto.
Qed.

Lemma submseteq_list_difference {A} `{EqDecision A} (l1 l2 l3 : list A) :
  NoDup l1 → (∀ a, a ∈ l3 → a ∉ l1) → l1 ⊆+ l2 → l1 ⊆+ list_difference l2 l3.
Proof.
  intros Hdup Hnin Hsub.
  apply NoDup_submseteq;auto.
  intros x Hx. apply elem_of_list_difference.
  split.
  - eapply elem_of_submseteq;eauto.
  - intros Hcontr. apply Hnin in Hcontr. done.
Qed.

Lemma list_difference_cons {A} `{EqDecision A} (l1 l2 : list A) (a : A) :
  NoDup l1 → a ∈ l1 → a ∉ l2 → list_difference l1 l2 ≡ₚ a :: list_difference l1 (a :: l2).
Proof.
  revert l2 a. induction l1;intros l2 a' Hdup Hin1 Hin2.
  - inversion Hin1.
  - simpl. destruct (decide_rel elem_of a l2).
    + assert (a ≠ a') as Hne; [intros Hcontr;subst;contradiction|].
      rewrite decide_True. { apply elem_of_cons. right;auto. }
      apply IHl1;auto. apply NoDup_cons in Hdup as [? ?];auto.
      apply elem_of_cons in Hin1 as [? | ?];[congruence|auto].
    + destruct (decide (a = a'));subst.
      * apply NoDup_cons in Hdup as [Hnin Hdup].
        f_equiv. rewrite decide_True;[constructor|].
        rewrite list_difference_skip;auto.
      * apply NoDup_cons in Hdup as [Hnin Hdup].
        apply elem_of_cons in Hin1 as [? | ?];[congruence|auto].
        erewrite IHl1;eauto. rewrite decide_False.
        apply not_elem_of_cons;auto. apply Permutation_swap.
Qed.

Lemma list_to_set_map_to_list {K V : Type} `{EqDecision K} `{Countable K}
      (m : gmap K V) :
  list_to_set (map_to_list m).*1 = dom m.
Proof.
  induction m using map_ind.
  - rewrite map_to_list_empty dom_empty_L. auto.
  - rewrite map_to_list_insert// dom_insert_L. simpl. rewrite IHm. auto.
Qed.

(* The last element of a list is the same as a list where we drop fewer elements than the list *)
Lemma last_drop_lt {A : Type} (l : list A) (i : nat) (a : A) :
  i < (length l) → list.last l = Some a → list.last (drop i l) = Some a.
Proof.
  generalize i. induction l.
  - intros i' Hlen Hlast. inversion Hlast.
  - intros i' Hlen Hlast. destruct i'.
    + simpl. apply Hlast.
    + simpl; simpl in Hlen. apply IHl; first lia.
      assert (0 < length l) as Hl; first lia.
      destruct l; simpl in Hl; first by apply Nat.lt_irrefl in Hl. auto.
Qed.

Lemma last_lookup {A : Type} (l : list A) :
  list.last l = l !! (length l - 1).
Proof.
  induction l.
  - done.
  - simpl; rewrite {1}/last -/last.
    destruct l; auto.
    rewrite IHl. simpl. rewrite PeanoNat.Nat.sub_0_r. done.
Qed.

Lemma last_app_iff {A : Type} (l1 l2 : list A) a :
  list.last l2 = Some a <-> length l2 > 0 ∧ list.last (l1 ++ l2) = Some a.
Proof.
  split.
  - intros Hl2.
    induction l1.
    + destruct l2; inversion Hl2. simpl. split; auto. lia.
    + destruct IHl1 as [Hlt Hlast]. split; auto. simpl; rewrite {1}/last -/last. rewrite Hlast.
      destruct (l1 ++ l2); auto.
      inversion Hlast.
  - generalize l1. induction l2; intros l1' [Hlen Hl].
    + inversion Hlen.
    + destruct l2;[rewrite last_snoc in Hl; inversion Hl; done|].
      rewrite -(IHl2 (l1' ++ [a0])); auto.
      simpl. split;[lia|]. rewrite -app_assoc -cons_middle. done.
Qed.

Lemma last_app_eq {A : Type} (l1 l2 : list A) :
  length l2 > 0 ->
  list.last l2 = list.last (l1 ++ l2).
Proof.
  revert l1. induction l2;intros l1 Hlen.
  - inversion Hlen.
  - destruct l2.
    + rewrite last_snoc. done.
    + rewrite cons_middle app_assoc -(IHl2 (l1 ++ [a]));[simpl;lia|auto].
Qed.

Lemma rev_nil_inv {A} (l : list A) :
  rev l = [] -> l = [].
Proof.
  destruct l;auto.
  simpl. intros Hrev. exfalso.
  apply app_eq_nil in Hrev as [Hrev1 Hrev2].
  inversion Hrev2.
Qed.

Lemma rev_singleton_inv {A} (l : list A) (a : A) :
  rev l = [a] -> l = [a].
Proof.
  destruct l;auto.
  simpl. intros Hrev.
  destruct l.
  - simpl in Hrev. inversion Hrev. auto.
  - exfalso. simpl in Hrev.
    apply app_singleton in Hrev. destruct Hrev as [ [Hrev1 Hrev2] | [Hrev1 Hrev2] ].
    + destruct (rev l);inversion Hrev1.
    + inversion Hrev2.
Qed.

Lemma rev_lookup {A} (l : list A) (a : A) :
  rev l !! 0 = Some a <-> l !! (length l - 1) = Some a.
Proof.
  split; intros Hl.
  - rewrite -last_lookup.
    induction l.
    + inversion Hl.
    + simpl in Hl. simpl. destruct l.
      { simpl in Hl. inversion Hl. auto. }
      { apply IHl. rewrite lookup_app_l in Hl;[simpl;rewrite app_length /=;lia|]. auto. }
  - rewrite -last_lookup in Hl.
    induction l.
    + inversion Hl.
    + simpl. destruct l.
      { simpl. inversion Hl. auto. }
      { rewrite lookup_app_l;[simpl;rewrite app_length /=;lia|]. apply IHl. auto. }
Qed.

Lemma rev_cons_inv {A} (l l' : list A) (a : A) :
  rev l = a :: l' ->
  ∃ l'', l = l'' ++ [a].
Proof.
  intros Hrel.
  destruct l;inversion Hrel.
  assert ((a0 :: l) !! (length l) = Some a) as Hsome.
  { assert (length l = length (a0 :: l) - 1) as ->;[simpl;lia|]. apply rev_lookup. rewrite Hrel. constructor. }
  apply take_S_r in Hsome.
  exists (take (length l) (rev (rev l ++ [a0]))).
    simpl. rewrite rev_unit. rewrite rev_involutive. rewrite -Hsome /=.
    f_equiv. rewrite firstn_all. auto.
Qed.


Require Import Coq.ZArith.ZArith.
Require Import Coq.setoid_ring.Ring_tac.
Require Import Crypto.Util.Prod.
Require Import Crypto.Util.SideConditions.CorePackages.
Require Export Crypto.Util.FixCoqMistakes.

Definition eq_by_Zring_prod_package (P : Prop) := P.

Ltac auto_split_prod_step_early special_ring_intros_tac :=
  match goal with
  | _ => progress hnf
  | _ => progress special_ring_intros_tac ()
  | [ H : prod _ _ |- _ ] => destruct H
  | [ |- forall a, _ ] => let a := fresh in intro a; compute in a
  end.

Ltac auto_split_prod_step special_ring_intros_tac :=
  match goal with
  | _ => auto_split_prod_step_early special_ring_intros_tac
  | [ |- pair _ _ = pair _ _ ] => apply f_equal2
  | [ |- @eq (prod _ _) _ _ ] => apply path_prod
  end.

Ltac Zring_prod_eq_tac_gen special_ring_intros_tac :=
  repeat auto_split_prod_step_early special_ring_intros_tac;
  cbv -[Z.add Z.sub Z.mul Z.div Z.pow Z.opp Z.log2 Z.land Z.lor Z.log2_up Z.abs];
  repeat match goal with
         | _ => auto_split_prod_step special_ring_intros_tac
         | [ |- @eq Z _ _ ] => ring
         end.

Ltac default_ring_intros_tac _ := fail.

Ltac Zring_prod_eq_tac _ := Zring_prod_eq_tac_gen default_ring_intros_tac.

Ltac autosolve_gen_intros special_ring_intros_tac else_tac :=
  lazymatch goal with
  | [ |- eq_by_Zring_prod_package _ ]
    => abstract Zring_prod_eq_tac_gen special_ring_intros_tac
  | _ => else_tac ()
  end.

Ltac autosolve else_tac := autosolve_gen_intros default_ring_intros_tac else_tac.


(* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
SPDX-License-Identifier: MIT-0 *)

Set Implicit Arguments.

Require Import FCF.
Require Import FCF.CompFold.
Require Import FCF.PRF.
Require Import Permutation.

Require Import KeyExchange.

Local Open Scope list_scope.

Theorem maybeBind_right_ident : forall (A : Type)(a : option A),
  (x <-? a; Some x) = a.

  intros.
  destruct a; reflexivity.

Qed.

Theorem maybeBind_assoc : forall (A B C: Type) (a : option A)(f : A -> option B)(g : B -> option C),
  Comp.maybeBind (Comp.maybeBind a f) g = Comp.maybeBind a (fun x => Comp.maybeBind (f x) g).

  intros.
  destruct a; reflexivity.

Qed.

Fixpoint stripOpt(A : Type)(ls : list (option A)) : option (list A) :=
  match ls with
  | nil => Some nil
  | a :: z =>
    a' <-? a;
    z' <-? stripOpt z;
    Some (a' :: z')
  end.

Theorem maybeBind_seq_eq : forall (A B : Type)(a : option A)(f1 f2 : A -> option B),
  (forall x, (f1 x) = (f2 x)) ->
  Comp.maybeBind a f1 = Comp.maybeBind a f2.

  intros. destruct a; simpl. eauto.
  reflexivity.

Qed.

Theorem stripOpt_app : forall (A : Type)(ls1 ls2 : list (option A)),
  stripOpt (ls1 ++ ls2) = (x <-? stripOpt ls1; y <-? stripOpt ls2; Some (x ++ y)).

  induction ls1; intuition; simpl in *.
  symmetry.
  eapply maybeBind_right_ident.
  rewrite maybeBind_assoc.
  eapply maybeBind_seq_eq; intros.
  rewrite IHls1.
  rewrite maybeBind_assoc.
  rewrite maybeBind_assoc.
  eapply maybeBind_seq_eq; intros.
  rewrite maybeBind_assoc.
  reflexivity.

Qed.

Section EntropySource.

  Variable Secret SecretAux : Set.
  Hypothesis Secret_EqDec : EqDec Secret.
  Hypothesis SecretAux_EqDec : EqDec SecretAux.

  Variable KeyGen : Comp (Secret * SecretAux).

  Definition EntropySource_G s a :=
     [s', _] <-$2 Repeat KeyGen (fun p => eqb (snd p) a);
     ret (eqb s s').

  Definition EntropySource m :=
    forall s a,
    In (s, a) (getSupport KeyGen) ->
    Pr[EntropySource_G s a] <= m.

End EntropySource.

Section KDF.

  Variable Secret Salt Context SecretAux OKM : Set.

  Hypothesis Secret_EqDec : EqDec Secret.
  Hypothesis SecretAux_EqDec : EqDec SecretAux.
  Hypothesis OKM_EqDec : EqDec OKM.
  Hypothesis Context_EqDec : EqDec Context.

  Definition KDF := Secret -> Salt -> Context -> nat -> OKM.
  Variable kdf : KDF.

  Variable RndOKM : nat -> Comp OKM.
  Variable defOKM : OKM.

  Section KDF_WeakSecure.

    Variable A_State : Set.
    Variable A1 : SecretAux -> Salt -> Comp (Context * nat * A_State).
    Variable A2 : A_State -> OKM -> Comp bool.
    Variable SaltGen : Comp Salt.

    Section KDF_SecureExact.
      Variable KeyGen : Comp (option (Secret * SecretAux)).

      (* This definition is weaker than usual---adversary is not given a
          KDF oracle. *)
      Definition KDF_WeakG0 :=
        sa_opt <-$ KeyGen;
        match sa_opt with
        | None => ret false
        | Some (s, a) =>
          r <-$ SaltGen;
          [c, l, s_A] <-$3 A1 a r;
          okm <- kdf s r c l;
          A2 s_A okm
        end.

      Definition KDF_WeakG1 :=
        sa_opt <-$ KeyGen;
        match sa_opt with
        | None => ret false
        | Some (s, a) =>
          r <-$ SaltGen;
          [c, l, s_A] <-$3 A1 a r;
          okm <-$ RndOKM l;
          A2 s_A okm
        end.

      Definition WeakKDF_Advantage :=
        | Pr[KDF_WeakG0] - Pr[KDF_WeakG1] |.

    End KDF_SecureExact.

  End KDF_WeakSecure.

  Section KDF_Secure.

    Definition KDF_Oracle s r (c : option Context) (_ : unit) cl : Comp (OKM * unit) :=
      res <- kdf s r (fst cl) (snd cl);
      match c with
      | None => ret (res, tt)
      | Some c' => ret (if (eqb c' (fst cl)) then (defOKM, tt) else (res, tt))
      end.

    Variable A_State : Set.
    Variable A1 : SecretAux -> Salt -> OracleComp (Context * nat) OKM (Context * nat * A_State).
    Variable A2 : A_State -> OKM -> OracleComp (Context * nat) OKM bool.

    Section KDF_SecureExact.
      Variable KeyGen : Comp (Secret * SecretAux).
      Variable SaltGen : Comp Salt.

      Definition KDF_G0 :=
        [s, a] <-$2 KeyGen;
        r <-$ SaltGen;
        [r_a, _] <-$2 A1 a r _ _ (KDF_Oracle s r None) tt;
        [c, l, s_A] <-3 r_a;
        okm <- kdf s r c l;
        [r_a, _] <-$2 A2 s_A okm _ _ (KDF_Oracle s r (Some c)) tt;
        ret r_a.

      Definition KDF_G1 :=
        [s, a] <-$2 KeyGen;
        r <-$ SaltGen;
        [r_a, _] <-$2 A1 a r _ _ (KDF_Oracle s r None) tt;
        [c, l, s_A] <-3 r_a;
        okm <-$ RndOKM l;
        [r_a, _] <-$2 A2 s_A okm _ _ (KDF_Oracle s r (Some c)) tt;
        ret r_a.

    End KDF_SecureExact.

    Definition EntropySecureKDF m eps :=
      forall KeyGen SaltGen,
      EntropySource _ _ KeyGen m ->
      | Pr[KDF_G0 KeyGen SaltGen] - Pr[KDF_G1 KeyGen SaltGen]| <= eps.
  End KDF_Secure.

  Section KDF_CR.

    Definition kdfProd p :=
      kdf (fst p) (fst (fst (snd p))) (snd (fst (snd p))) (snd (snd p)).

    Variable A : Comp ((Secret * Secret) * (Salt * Context * nat)).

    Definition KDF_CR :=
      [s1, s2, x] <-$3 A;
      ret (if (s1 ?= s2) then false else
        ((kdfProd (s1, x)) ?= (kdfProd (s2, x)))).

    Definition KDF_CR_Advantage :=
      Pr[KDF_CR].

  End KDF_CR.

End KDF.

Definition Octet := Bvector 8.
Definition OctetString := list Octet.
Definition OctetStringKDF := KDF OctetString OctetString OctetString OctetString.

Ltac OctetString_EqDec := repeat
  match goal with
  | [|- EqDec (prod _ _) ] => apply pair_EqDec
  | [|- EqDec (list _) ] => apply list_EqDec
  | [|- EqDec OctetString ] => apply list_EqDec
  | [|- EqDec Octet ] => apply Bvector_EqDec
  | [|- EqDec (Bvector _) ] => apply Bvector_EqDec
  end.

Definition rndOctet := {0,1}^8.

Definition rndOctetString strLength :=
  compMap _ (fun _ => rndOctet) (forNats strLength).

Section CtKDF.

  Variable kdf : OctetStringKDF.
  Variable format : OctetString -> list OctetString -> list OctetString -> OctetString.
  Variable keyLength : nat.
  Variable psk : OctetString.
  Variable context : OctetString.
  Variable label : OctetString.

  Definition ctKDF k MA MB length :=
    let secret := psk ++ flatten k in
    let kdf_context := format context MA MB in
    kdf secret label kdf_context length.

  Definition OctetStringKE := (@KeyExchange OctetString OctetString OctetString OctetString _ _ ).
  Variable KE : list OctetStringKE.
  Variable defKE : OctetStringKE.

  Definition ctKeyGen (KE : list OctetStringKE) :=
    ls <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE;
      ret (fst (split ls), snd (split ls)).

  Definition ctResponseFunc (KE : list OctetStringKE)(Ps : list OctetString) :=
    ls_opt <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE Ps);
    ret (
      ls <-? stripOpt ls_opt;
      Rs <- snd (split ls);
      k <- ctKDF (fst (split ls)) Ps Rs keyLength;
      Some (k, Rs)
    ).

  Instance CtKE : (@KeyExchange (list OctetString) (list OctetString) (list OctetString) OctetString _ _ ) := {
    keyGen := ctKeyGen KE;
    responseFunc := ctResponseFunc KE
    (*;receiveFunc := ctReceiveFunc KE *)
  }.

  Hypothesis keyGen_wf : forall ke,
    In ke KE ->
    well_formed_comp (@keyGen _ _ _ _ _ _ ke).
  Hypothesis responseFunc_wf : forall ke x,
    In ke KE ->
    well_formed_comp (@responseFunc _ _ _ _ _ _ ke x).

  (* assume an arbitrary "strong" KE scheme *)
  Variable strongKE : nat.
  Hypothesis strongKE_small : strongKE < length KE.

  (* Prove that we can split the list of key exchange operations at an arbitrary
    position and reorder them so that some chosen key exchange happens first. *)
  Section CtKDF_SplitReorder.

    Variable A : (list (OctetString * OctetString)) -> (list (option (OctetString * OctetString))) -> Comp bool.

    Definition CtKDF_SplitReorder_G0 :=
      lsSK <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE;
      lsR <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE (snd (split lsSK)));
      (A lsSK lsR).

    Definition CtKDF_SplitReorder_G1 :=
      SK_strong <-$ @keyGen _ _ _ _ _ _ (nth strongKE KE defKE);
      R_strong <-$ (@responseFunc _ _ _ _ _ _ (nth strongKE KE defKE) (snd SK_strong) );
      lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (firstn strongKE KE);
      lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (skipn (S strongKE) KE);
      lsR1 <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (firstn strongKE KE) (snd (split lsSK1)));
      lsR2 <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (skipn (S strongKE) KE) (snd (split lsSK2)));
      (A (lsSK1 ++ (SK_strong :: lsSK2)) (lsR1 ++ (R_strong :: lsR2))).

    (* first split map operations *)
    Definition CtKDF_SplitReorder_G0a :=
      lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (firstn strongKE KE);
      SK_strong <-$ @keyGen _ _ _ _ _ _ (nth strongKE KE defKE);
      lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (skipn (S strongKE) KE);
      lsR1 <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (firstn strongKE KE) (snd (split lsSK1)));
      R_strong <-$ (@responseFunc _ _ _ _ _ _ (nth strongKE KE defKE) (snd SK_strong) );
      lsR2 <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (skipn (S strongKE) KE) (snd (split lsSK2)));
      (A (lsSK1 ++ (SK_strong :: lsSK2)) (lsR1 ++ (R_strong :: lsR2))).

    Theorem CtKDF_SplitReorder_G0_G0a_equiv :
       Pr[CtKDF_SplitReorder_G0] == Pr[CtKDF_SplitReorder_G0a].

      unfold CtKDF_SplitReorder_G0, CtKDF_SplitReorder_G0a.
      rewrite (@lsAppCons _ KE strongKE defKE) at 1; trivial.
      eapply eqRat_trans.
      eapply evalDist_seq_eq_trans.
      intros.
      eapply compMap_app.
      fcf_inline_first.
      fcf_skip.
      reflexivity.
      fcf_inline_first.
      fcf_skip.
      fcf_inline_first.
      fcf_skip.
      reflexivity.

      rewrite (@lsAppCons _ (combine KE (snd (split (x ++ x0 :: x1)))) strongKE (defKE, nil)) at 1; trivial.
      eapply eqRat_trans.
      eapply evalDist_seq_eq_trans.
      intros.
      eapply compMap_app.
      fcf_inline_first.
      fcf_skip.
      rewrite firstn_combine.
      rewrite firstn_snd_split_app_gen.
      reflexivity.
      assert (length x = length (firstn strongKE KE)).
      eapply compMap_length; eauto.
      rewrite H3.
      symmetry.
      apply firstn_length_le.
      omega.

      fcf_inline_first.
      fcf_skip.
      rewrite nth_combine. simpl.
      rewrite nth_snd_split_app_gen.
      rewrite nth_0_snd_split_cons.
      reflexivity.
      assert (length x = length (firstn strongKE KE)).
      eapply compMap_length; eauto.
      rewrite H4.
      symmetry.
      apply firstn_length_le.
      omega.

      rewrite snd_split_length.
      apply compMap_length in H2.
      apply compMap_length in H0.
      rewrite app_length. simpl.
      rewrite H0.
      rewrite H2.
      rewrite firstn_length_le.
      rewrite skipn_length.
      omega.
      omega.

      fcf_inline_first.
      fcf_skip.
      rewrite skipn_combine.
      rewrite skipn_snd_split_app_cons_gen.
      reflexivity.
      f_equal.
      assert (length x = length (firstn strongKE KE)).
      eapply compMap_length; eauto.
      rewrite H5.
      symmetry.
      apply firstn_length_le.
      omega.

      rewrite combine_length.
      rewrite min_l.
      trivial.
      apply compMap_length in H0.
      apply compMap_length in H2.

      rewrite snd_split_app.
      rewrite snd_split_cons.
      repeat rewrite app_length.
      simpl.
      repeat rewrite snd_split_length.
      rewrite H0.
      rewrite H2.
      rewrite firstn_length_le.
      rewrite skipn_length.
      omega.
      omega.
    Qed.

    (* reorder *)
    Theorem CtKDF_SplitReorder_G0a_G1_equiv :
     Pr[CtKDF_SplitReorder_G0a] == Pr[CtKDF_SplitReorder_G1].

      unfold CtKDF_SplitReorder_G0a, CtKDF_SplitReorder_G1.
      fcf_swap leftc.
      fcf_skip.
      fcf_swap rightc.
      fcf_skip.
      fcf_swap rightc.
      fcf_skip.
      fcf_swap rightc.
      reflexivity.
    Qed.

    Theorem CtKDF_SplitReorder_equiv :
      Pr[CtKDF_SplitReorder_G0] == Pr[CtKDF_SplitReorder_G1].

      rewrite CtKDF_SplitReorder_G0_G0a_equiv.
      apply CtKDF_SplitReorder_G0a_G1_equiv.

    Qed.

  End CtKDF_SplitReorder.

  (* Assume a distribution on octet strings that has a certain amount of entropy *)
  Variable rndSharedSecret : Comp OctetString.
  Hypothesis rndSharedSecret_wf : well_formed_comp rndSharedSecret.
  Variable rndSecretEps : Rat.
  Hypothesis rndSharedSecret_prob_le : forall b,
    evalDist rndSharedSecret b <= rndSecretEps.

  Section CtKDF_IND_CPA.

  Variable A : (list OctetString) -> (list OctetString) -> OctetString -> Comp bool.

  (* constructed adversary against the strong KE scheme*)
  Definition StrongKE_A (P_strong R_strong k_strong: OctetString) : Comp bool:=
    let KE1 := firstn strongKE KE in
    let KE2 := skipn (S strongKE) KE in
    lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE1;
    sk1 <- fst (split lsSK1);
    P1 <- snd (split lsSK1);
    lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE2;
    sk2 <- fst (split lsSK2);
    P2 <- snd (split lsSK2);

    lsR1_opt <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE1 P1);
    lsR2_opt <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE2 P2);
    match (
      lsR1 <-? stripOpt lsR1_opt;
      R1 <- snd (split lsR1);
      lsK1 <- fst (split lsR1);
      lsR2 <-? stripOpt lsR2_opt;
      R2 <- snd (split lsR2);
      lsK2 <- fst (split lsR2);
      k <- ctKDF (lsK1 ++ (k_strong :: lsK2)) (P1 ++ P_strong :: P2) (R1 ++ R_strong :: R2) keyLength;
      (Some (k, (R1 ++ R_strong :: R2)))
    ) with
    | None => ret false
    | Some (k, R) => A (P1 ++ P_strong :: P2) R k
    end.

  (* constructed adversary against the KDF *)
  Definition KDF_A1 a (r : OctetString) :=
    ret ((format context (fst a) (snd a)), keyLength, a).

  Definition KDF_A2 s_A okm :=
    (A (fst s_A) (snd s_A) okm).

  (* We cannot assume that the KDF is a generic extractor, because there is no salt. *)
  (* Instead, we assume something weaker, that the KDF is an extractor for any string containing
    a random shared secret from a given distribution. *)
  (* In this assumption, an adversary determines the rest of the source distribution. *)
  Definition StrongKE_KeyGen (A : Comp (option (OctetString * OctetString * (list OctetString * list OctetString)))) :=
    k_strong <-$ rndSharedSecret;
    xOpt <-$ A;
    ret (
      x <-? xOpt;
      [s1, s2, aux] <-3 x;
      Some (s1 ++ k_strong ++ s2, aux)
    ).

  (* The constructed adversary that determines the rest of the source distribution *)
  Definition KDF_KeyGen_A :=

    let KE_strong := nth strongKE KE defKE in

    [_, P_strong] <-$2 @keyGen _ _ _ _ _ _ KE_strong;
    krStrong_opt <-$ @responseFunc _ _ _ _ _ _ KE_strong P_strong;
    let KE1 := firstn strongKE KE in
    let KE2 := skipn (S strongKE) KE in

    lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE1;
    P1 <- snd (split lsSK1);

    lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE2;
    P2 <- snd (split lsSK2);

    lsR1_opt <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE1 P1);
    lsR2_opt <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE2 P2);
    ret (
      krStrong <-? krStrong_opt;
      R_strong <- snd krStrong;
      lsR1 <-? stripOpt lsR1_opt;
      R1 <- snd (split lsR1);
      lsK1 <- fst (split lsR1);
      lsR2 <-? stripOpt lsR2_opt;
      R2 <- snd (split lsR2);
      lsK2 <- fst (split lsR2);
      Some ((psk ++ (flatten lsK1), (flatten lsK2)), (P1 ++ P_strong :: P2, R1 ++R_strong :: R2))
    ).

  Definition kdfAdvantage :=
    @WeakKDF_Advantage _ _ _ _ _ kdf rndOctetString _ KDF_A1 KDF_A2 (ret label) (StrongKE_KeyGen KDF_KeyGen_A).

  Definition strongKE_Advantage := (@KeyExchangeCPA_Advantage _ _ _ _ _ _
      (nth strongKE KE defKE) rndSharedSecret StrongKE_A).

  (* inline and simplify *)
  Definition CtKDF_G1 :=
    lsSK <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE;
    sk <- fst (split lsSK);
    P <- snd (split lsSK);
    lsR_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE P);
    lsR_opt <- stripOpt lsR_opts;
    kr_opt <- (
      lsR <-? lsR_opt;
      R <- snd (split lsR);
      lsK <- fst (split lsR);
      k <- ctKDF lsK P R keyLength;
      Some (k, R)
    );
    match kr_opt with
    | None => ret false
    | Some (k, R) => A P R k
    end.

  Theorem CtKDF_G1_equiv :
    Pr[(@KeyExchangeCPA_G0 _ _ _ _ _ _ CtKE A)] == Pr[CtKDF_G1].

    unfold KeyExchangeCPA_G0, keyGen, responseFunc, CtKE,
      ctKeyGen, ctResponseFunc, CtKDF_G1.
    repeat (fcf_inline_first; fcf_skip).
  Qed.

  (* Split list at strong ke, use constructed adversary *)
  Definition CtKDF_G2 :=

    let KE_strong := nth strongKE KE defKE in

    [_, P_strong] <-$2 @keyGen _ _ _ _ _ _ KE_strong;
    krStrong_opt <-$ @responseFunc _ _ _ _ _ _ KE_strong P_strong;
    match krStrong_opt with
    | None => ret false
    | Some (k, R) => StrongKE_A P_strong R k
    end.

  (* remove let bindings *)
  Definition CtKDF_G1a :=
    lsSK <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE;
    lsR_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE (snd (split lsSK)));
    match (stripOpt lsR_opts) with
    | None => ret false
    | Some lsR =>
      (A (snd (split lsSK)) (snd (split lsR)) (ctKDF (fst (split lsR)) (snd (split lsSK)) (snd (split lsR)) keyLength))
    end.

  Theorem CtKDF_G1_G1a_equiv :
    Pr[CtKDF_G1] == Pr[CtKDF_G1a].

    unfold CtKDF_G1, CtKDF_G1a.
    repeat (fcf_simp; fcf_skip).
    unfold Comp.maybeBind.
    symmetry.
    match goal with
    | [|- context[match ?a with |Some _ => _ | None => _ end] ] =>
      destruct a
    end;
    reflexivity.

  Qed.

  (* split and reorder *)
  Definition CtKDF_G1b :=
    SK_strong <-$ @keyGen _ _ _ _ _ _ (nth strongKE KE defKE);
    R_strong_opt <-$ (@responseFunc _ _ _ _ _ _ (nth strongKE KE defKE) (snd SK_strong) );
    lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (firstn strongKE KE);
    lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (skipn (S strongKE) KE);
    lsR1_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (firstn strongKE KE) (snd (split lsSK1)));
    lsR2_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (skipn (S strongKE) KE) (snd (split lsSK2)));
    match (stripOpt (lsR1_opts ++ R_strong_opt :: lsR2_opts)) with
    | None => ret false
    | Some R =>
    (A ((snd (split lsSK1))++(snd SK_strong)::(snd (split lsSK2))) (snd (split R))
      (ctKDF
        (fst (split R))
        ((snd (split lsSK1))++(snd SK_strong)::(snd (split lsSK2)))
        (snd (split R))
        keyLength))
    end.

  Definition CtKDF_G1a_A (lsSK : list (OctetString * OctetString)) (lsR : list (OctetString * OctetString)):=
    (A (snd (split lsSK)) (snd (split lsR))
      (ctKDF (fst (split lsR)) (snd (split lsSK)) (snd (split lsR)) keyLength)).

  Theorem CtKDF_G1a_G1b_equiv :
     Pr[CtKDF_G1a] == Pr[CtKDF_G1b].

    unfold CtKDF_G1a, CtKDF_G1b.
    eapply eqRat_trans.
    apply (CtKDF_SplitReorder_equiv).
    unfold CtKDF_SplitReorder_G1.
    repeat (fcf_skip).
    repeat rewrite snd_split_app.
    repeat rewrite fst_split_app.
    repeat rewrite snd_split_cons.
    repeat rewrite fst_split_cons.
    reflexivity.
  Qed.

  (* put let bindings back in *)
  Theorem CtKDF_G1b_G2_equiv :
     Pr[CtKDF_G1b] == Pr[CtKDF_G2].

    unfold CtKDF_G1b, CtKDF_G2, StrongKE_A.
    repeat (fcf_simp; fcf_skip).
    reflexivity.
    case_eq x; intros.
    destruct p.
    do 4 fcf_skip.
    rewrite stripOpt_app.
    destruct (stripOpt x1).
    + simpl.
    destruct (stripOpt x2).
    * simpl.
    repeat rewrite snd_split_app.
    repeat rewrite fst_split_app.
    repeat rewrite snd_split_cons.
    repeat rewrite fst_split_cons.
    reflexivity.
    * simpl.
    reflexivity.
    + simpl.
    reflexivity.

    + subst.
    fcf_irr_l.
    eapply compMap_wf; intuition.
    eapply keyGen_wf.
    eapply firstn_In; eauto.
    fcf_irr_l.
    eapply compMap_wf; intuition.
    eapply keyGen_wf.
    destruct KE; trivial.
    simpl. right.
    eapply skipn_In; eauto.
    fcf_irr_l.
    eapply compMap_wf; intuition.
    eapply responseFunc_wf.
    destruct a; simpl in *.
    eapply firstn_In.
    eapply in_combine_l; eauto.
    fcf_irr_l.
    eapply compMap_wf; intuition.
    eapply responseFunc_wf.
    destruct KE; trivial.
    simpl. right.
    destruct a; simpl in *.
    eapply skipn_In.
    eapply in_combine_l; eauto.
    rewrite stripOpt_app.
    destruct (stripOpt x1); simpl;  reflexivity.

  Qed.

  Theorem CtKDF_G2_equiv :
     Pr[CtKDF_G1] == Pr[CtKDF_G2].

    rewrite CtKDF_G1_G1a_equiv.
    rewrite CtKDF_G1a_G1b_equiv.
    apply CtKDF_G1b_G2_equiv.
  Qed.

  (* apply security assumption for strong KE *)
  Definition CtKDF_G3 :=

    let KE_strong := nth strongKE KE defKE in

    [sk_strong, P_strong] <-$2 @keyGen _ _ _ _ _ _ KE_strong;
    R_opt <-$ @responseFunc _ _ _ _ _ _ KE_strong P_strong;
    match R_opt with
    | None => ret false
    | Some (_, R_strong) =>
        k_strong <-$ rndSharedSecret;
        StrongKE_A P_strong R_strong k_strong
    end.

  Theorem CtKDF_G3_close :
     | Pr[CtKDF_G2] - Pr[CtKDF_G3] | == strongKE_Advantage.

    reflexivity.

  Qed.

  (* inline KE constructed adversary and put game in correct form to use KDF assumption*)
  Definition KDF_KeyGen :=

    let KE_strong := nth strongKE KE defKE in

    [_, P_strong] <-$2 @keyGen _ _ _ _ _ _ KE_strong;
    R_opt <-$ @responseFunc _ _ _ _ _ _ KE_strong P_strong;
    k_strong <-$ rndSharedSecret;

    let KE1 := firstn strongKE KE in
    let KE2 := skipn (S strongKE) KE in

    lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE1;
    P1 <- snd (split lsSK1);

    lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE2;
    P2 <- snd (split lsSK2);

    lsR1_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE1 P1);
    lsR2_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE2 P2);
    ret (
      lsR1 <-? stripOpt lsR1_opts;
      lsR2 <-? stripOpt lsR2_opts;
      R <-? R_opt;

      Some (psk ++ flatten ((fst (split lsR1)) ++ k_strong :: (fst (split lsR2))), (P1 ++ P_strong :: P2, (snd (split lsR1))++(snd R) :: (snd (split lsR2))))
    ).

  Definition CtKDF_G4 :=
    r <-$ ret label;
    x_opt <-$ KDF_KeyGen;
    match x_opt with
    | None => ret false
    | Some (s, a) =>
      [c, l, s_A] <-$3 KDF_A1 a r;
      okm <- kdf s r c l;
      KDF_A2 s_A okm
    end.

  Theorem CtKDF_G4_equiv :
    Pr[CtKDF_G3] == Pr[CtKDF_G4].

    Local Opaque evalDist.

    unfold CtKDF_G3, CtKDF_G4.
    unfold StrongKE_A, KDF_KeyGen, KDF_A1, KDF_A2, ctKDF.
    fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    destruct x.
    destruct p.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    case_eq (stripOpt x2); intros.
    simpl.
    case_eq (stripOpt x3); intros.
    simpl.
    fcf_simp.
    simpl.
    reflexivity.
    reflexivity.
    reflexivity.

    fcf_inline_first.
    fcf_irr_r.
    fcf_inline_first.
    fcf_irr_r.
    eapply compMap_wf; intros.
    eapply keyGen_wf.
    eapply firstn_In; eauto.
    fcf_inline_first.
    fcf_irr_r.
    eapply compMap_wf; intros.
    destruct KE; simpl in *; intuition idtac.
    eapply keyGen_wf.
    right.
    eapply skipn_In; eauto.
    fcf_inline_first.
    fcf_irr_r.
    eapply compMap_wf.
    intros.
    eapply responseFunc_wf.
    destruct a. simpl in *.
    eapply firstn_In.
    eapply in_combine_l; eauto.
    fcf_inline_first.
    fcf_irr_r.
    eapply compMap_wf.
    intros.
    eapply responseFunc_wf.
    destruct a. simpl in *.
    destruct KE; simpl in *; intuition idtac.
    right.
    eapply skipn_In.
    eapply in_combine_l; eauto.

    fcf_simp.
    destruct (stripOpt x2).
    simpl.
    destruct (stripOpt x3).
    simpl.
    reflexivity.
    reflexivity.
    reflexivity.

    Local Transparent evalDist.

  Qed.

  (* apply KDF definition to replace okm with random *)
  Definition CtKDF_G5 :=
    r <-$ ret label;
    x_opt <-$ KDF_KeyGen;
    match x_opt with
    | None => ret false
    | Some (s, a) =>
      [c, l, s_A] <-$3 KDF_A1 a r;
      okm <-$ rndOctetString l;
      KDF_A2 s_A okm
    end.

  (* Need to prove that KDF source distribution is of the correct form.
    Transform KDF_KeyGen to make this easier. *)
  (* move secret generation to the end so we can factor it out of the repeat *)
  Theorem KDF_KeyGen_A_wf :
    well_formed_comp KDF_KeyGen_A.

    unfold KDF_KeyGen_A.
    wftac.
    eapply keyGen_wf.
    apply nth_In.
    trivial.
    apply responseFunc_wf.
    apply nth_In.
    trivial.
    apply compMap_wf.
    intros.
    apply keyGen_wf.
    eapply firstn_In; eauto.
    apply compMap_wf.
    intros.
    apply keyGen_wf.
    eapply skipn_In.
    eauto.
    apply compMap_wf.
    intros.
    apply responseFunc_wf.
    destruct a. simpl in *.
    eapply in_combine_l in H4.
    eapply firstn_In.
    eauto.
    apply compMap_wf.
    intros.
    apply responseFunc_wf.
    destruct a. simpl.
    apply in_combine_l in H5.
    eapply skipn_In.
    eauto.

  Qed.

  Theorem KDF_KeyGen_strong_equiv :
    forall x,
    evalDist KDF_KeyGen x == evalDist (StrongKE_KeyGen KDF_KeyGen_A) x.

    intuition.
    unfold StrongKE_KeyGen, KDF_KeyGen, KDF_KeyGen_A.
    fcf_at fcf_inline rightc 1%nat.
    fcf_swap rightc.
    fcf_skip.
    fcf_simp.
    fcf_inline_first; fcf_swap leftc; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip.
    fcf_inline_first; fcf_skip.
    fcf_inline_first; fcf_skip.
    fcf_inline_first; fcf_skip.
    fcf_inline_first; fcf_skip.
    fcf_simp.
    eapply evalDist_ret_eq.
    destruct x1; simpl.
    destruct (stripOpt x4); simpl.
    destruct (stripOpt x5); simpl.
    repeat rewrite flatten_app.
    rewrite app_assoc.
    reflexivity.
    trivial.
    trivial.
    destruct (stripOpt x4); simpl.
    destruct (stripOpt x5); simpl.
    trivial.
    trivial.
    trivial.
  Qed.

  Definition CtKDF_G4a :=
    r <-$ ret label;
    x_opt <-$ (StrongKE_KeyGen KDF_KeyGen_A);
    match x_opt with
    | None => ret false
    | Some (s, a) =>
      [c, l, s_A] <-$3 KDF_A1 a r;
      okm <- kdf s r c l;
      KDF_A2 s_A okm
    end.

  Theorem CtKDF_G4a_equiv :
    Pr[CtKDF_G4] == Pr[CtKDF_G4a].

    unfold CtKDF_G4, CtKDF_G4a.
    fcf_simp.
    eapply evalDist_seq_eq.
    intros.
    eapply KDF_KeyGen_strong_equiv.
    reflexivity.

  Qed.

  Definition CtKDF_G4b :=
    r <-$ ret label;
    x_opt <-$ (StrongKE_KeyGen KDF_KeyGen_A);
    match x_opt with
    | None => ret false
    | Some (s, a) =>
      [c, l, s_A] <-$3 KDF_A1 a r;
      okm <-$ rndOctetString l;
      KDF_A2 s_A okm
    end.

  Theorem CtKDF_G4b_G5_equiv :
    Pr[CtKDF_G4b] == Pr[CtKDF_G5].

    unfold CtKDF_G4b, CtKDF_G5.
    fcf_simp.
    eapply evalDist_seq_eq.
    intros.
    symmetry.
    eapply KDF_KeyGen_strong_equiv.
    intuition.

  Qed.

  Theorem CtKDF_G5_close:
    | Pr[CtKDF_G4] - Pr[CtKDF_G5] | <= kdfAdvantage.

    rewrite <- CtKDF_G4b_G5_equiv.
    rewrite CtKDF_G4a_equiv.

    unfold kdfAdvantage, CtKDF_G4a, CtKDF_G4b, WeakKDF_Advantage, KDF_WeakG0, KDF_WeakG1.
    eapply eqRat_impl_leRat.
    eapply ratDistance_eqRat_compat.
    +fcf_simp.
    fcf_skip.
    destruct x. destruct p.
    fcf_simp.
    reflexivity.
    reflexivity.

    +fcf_simp.
    fcf_skip.
    destruct x. destruct p.
    fcf_simp.
    reflexivity.
    reflexivity.

  Qed.

  (* put everything back in the correct order*)
  Definition CtKDF_G6 :=
    lsSK <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) KE;
    sk <- fst (split lsSK);
    P <- snd (split lsSK);
    lsR_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine KE P);
    lsR_opt <- stripOpt lsR_opts;
    match lsR_opt with
    | None => ret false
    | Some lsR =>
      R <- snd (split lsR);
      lsK <- fst (split lsR);
      k <-$ rndOctetString keyLength;
      A P R k
    end.

  (* remove let bindings and remove unnecessary shared secret *)
  Definition CtKDF_G5a :=
  SK_strong <-$ @keyGen _ _ _ _ _ _ (nth strongKE KE defKE);
  R_opt <-$ (@responseFunc _ _ _ _ _ _ (nth strongKE KE defKE) (snd SK_strong) );
  lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (firstn strongKE KE);
  lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (skipn (S strongKE) KE);
  lsR1_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (firstn strongKE KE) (snd (split lsSK1)));
  lsR2_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (skipn (S strongKE) KE) (snd (split lsSK2)));
  okm <-$ rndOctetString keyLength;
  match (stripOpt (lsR1_opts ++ R_opt :: lsR2_opts)) with
  | None => ret false
  | Some lsR =>
    (A ((snd (split lsSK1))++(snd SK_strong)::(snd (split lsSK2))) (snd (split lsR))
      okm)
  end.

  Theorem CtKDF_G5_G5a_equiv :
    Pr[CtKDF_G5] == Pr[CtKDF_G5a].

    Local Opaque evalDist.

    unfold CtKDF_G5, CtKDF_G5a, KDF_KeyGen, KDF_A1, KDF_A2.
    fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    reflexivity.
    fcf_inline_first.
    fcf_irr_l.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    fcf_inline_first; fcf_skip; fcf_simp.
    rewrite stripOpt_app.
    destruct (stripOpt x3).
    simpl.
    destruct (stripOpt x4).
    simpl.
    destruct x.
    simpl.
    fcf_simp.
    fcf_skip.
    simpl.
    rewrite snd_split_app.
    rewrite snd_split_cons.
    reflexivity.
    simpl.
    fcf_irr_r.
    eapply compMap_wf.
    intros.
    unfold rndOctet. wftac.
    simpl.
    fcf_irr_r.
    eapply compMap_wf.
    intros.
    unfold rndOctet. wftac.
    destruct x.
    simpl.
    reflexivity.
    simpl.
    reflexivity.
    simpl.
    fcf_irr_r.
    eapply compMap_wf.
    intros.
    unfold rndOctet. wftac.

    Local Transparent evalDist.
  Qed.

  (* reorder *)
  Definition CtKDF_G5b :=
    lsSK1 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (firstn strongKE KE);
    SK_strong <-$ @keyGen _ _ _ _ _ _ (nth strongKE KE defKE);
    lsSK2 <-$ compMap _ (fun ke => @keyGen _ _ _ _ _ _ ke) (skipn (S strongKE) KE);

    lsR1_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (firstn strongKE KE) (snd (split lsSK1)));
    R_strong_opt <-$ (@responseFunc _ _ _ _ _ _ (nth strongKE KE defKE) (snd SK_strong) );
    lsR2_opts <-$ compMap _ (fun p => @responseFunc _ _ _ _ _ _ (fst p) (snd p)) (combine (skipn (S strongKE) KE) (snd (split lsSK2)));

    okm <-$ rndOctetString keyLength;
    match (stripOpt (lsR1_opts ++ R_strong_opt :: lsR2_opts)) with
    | None => ret false
    | Some lsR =>
      (A ((snd (split lsSK1))++(snd SK_strong)::(snd (split lsSK2))) (snd (split lsR))
        okm)
    end.

  Theorem CtKDF_G5a_G5b_equiv :
    Pr[CtKDF_G5a] == Pr[CtKDF_G5b].

    unfold CtKDF_G5a, CtKDF_G5b.
    fcf_swap rightc.
    fcf_skip.
    fcf_swap leftc.
    fcf_skip.
    fcf_swap leftc.
    fcf_skip.
    fcf_swap leftc.
    reflexivity.

  Qed.

  (* combine map operations *)
  Theorem CtKDF_G5b_G6_equiv :
     Pr[CtKDF_G5b] == Pr[CtKDF_G6].

    unfold CtKDF_G5b, CtKDF_G6.
    symmetry.
    rewrite (@lsAppCons _ KE strongKE defKE) at 1; trivial.
    eapply eqRat_trans.
    eapply evalDist_seq_eq_trans.
    intros.
    eapply compMap_app.
    fcf_inline_first.
    fcf_skip.
    reflexivity.
    fcf_inline_first.
    fcf_skip.
    fcf_inline_first.
    fcf_skip.
    reflexivity.

    rewrite (@lsAppCons _ (combine KE (snd (split (x ++ x0 :: x1)))) strongKE (defKE, nil)) at 1; trivial.
    eapply eqRat_trans.
    eapply evalDist_seq_eq_trans.
    intros.
    eapply compMap_app.
    fcf_inline_first.
    fcf_skip.
    rewrite firstn_combine.
    rewrite firstn_snd_split_app_gen.
    reflexivity.
    assert (length x = length (firstn strongKE KE)).
    eapply compMap_length; eauto.
    rewrite H3.
    symmetry.
    apply firstn_length_le.
    omega.

    fcf_inline_first.
    fcf_skip.
    rewrite nth_combine. simpl.
    rewrite nth_snd_split_app_gen.
    rewrite nth_0_snd_split_cons.
    reflexivity.
    assert (length x = length (firstn strongKE KE)).
    eapply compMap_length; eauto.
    rewrite H4.
    symmetry.
    apply firstn_length_le.
    omega.

    rewrite snd_split_length.
    apply compMap_length in H2.
    apply compMap_length in H0.
    rewrite app_length. simpl.
    rewrite H0.
    rewrite H2.
    rewrite firstn_length_le.
    rewrite skipn_length.
    omega.
    omega.

    fcf_inline_first.
    fcf_skip.
    rewrite skipn_combine.
    rewrite skipn_snd_split_app_cons_gen.
    reflexivity.
    f_equal.
    assert (length x = length (firstn strongKE KE)).
    eapply compMap_length; eauto.
    rewrite H5.
    symmetry.
    apply firstn_length_le.
    omega.

    repeat rewrite snd_split_app.
    repeat rewrite fst_split_app.
    repeat rewrite snd_split_cons.
    repeat rewrite fst_split_cons.
    destruct (stripOpt (x2 ++ x3 :: x4)).
    reflexivity.
    fcf_irr_r.
    eapply compMap_wf.
    intros.
    unfold rndOctet. wftac.

    rewrite combine_length.
    rewrite min_l.
    trivial.
    apply compMap_length in H0.
    apply compMap_length in H2.

    rewrite snd_split_app.
    rewrite snd_split_cons.
    repeat rewrite app_length.
    simpl.
    repeat rewrite snd_split_length.
    rewrite H0.
    rewrite H2.
    rewrite firstn_length_le.
    rewrite skipn_length.
    omega.
    omega.
  Qed.

  Theorem CtKDF_G6_equiv :
     Pr[CtKDF_G5] == Pr[CtKDF_G6].

    rewrite CtKDF_G5_G5a_equiv.
    rewrite CtKDF_G5a_G5b_equiv.
    apply CtKDF_G5b_G6_equiv.
  Qed.

  (* the last game is a simplified form of the second game in the security definition *)
  Theorem CtKDF_G6_CPA_equiv :
     Pr[CtKDF_G6] == Pr[(@KeyExchangeCPA_G1 _ _ _ _ _ _ CtKE (rndOctetString keyLength) A)].

    unfold CtKDF_G6, KeyExchangeCPA_G1.
    unfold CtKE, keyGen, ctKeyGen.
    fcf_inline_first; fcf_skip; fcf_simp.
    reflexivity.
    fcf_inline_first; fcf_skip; fcf_simp.
    destruct (stripOpt x0);
    reflexivity.
  Qed.

  Theorem CtKE_secure :
    @KeyExchangeCPA_Advantage _ _ _ _ _ _ CtKE (rndOctetString keyLength) A <=
    strongKE_Advantage + kdfAdvantage.

    unfold KeyExchangeCPA_Advantage.
    rewrite CtKDF_G1_equiv.
    rewrite CtKDF_G2_equiv.
    eapply ratDistance_le_trans.
    apply eqRat_impl_leRat.
    apply CtKDF_G3_close.
    rewrite CtKDF_G4_equiv.
    rewrite <- CtKDF_G6_CPA_equiv.
    rewrite <- CtKDF_G6_equiv.
    apply CtKDF_G5_close.
  Qed.

End CtKDF_IND_CPA.

End CtKDF.


From Cat Require Import Imports Category Preorder Monoid Poset Isomorphism
Terminal Dual Initial Product CoProduct Exponential CCC IPL STLC Functor
NaturalTransformation Adjunction ParamKleisliTriple Monad Presheaf Yoneda.
Require Import FunctionalExtensionality.
Require Import ProofIrrelevance.
Require Import JMeq.
Set Universe Polymorphism.

Arguments fmap {_} {_} _ {_} {_} _.
Arguments fobj {_} {_} _ _.
Arguments Compose_Functors {_} {_} {_} _ _.
Arguments NaturalTransformation {_} {_} _ _.
Arguments trans {_} {_} {_} {_} _ _.
Arguments Compose_NaturalTransformations {_ _ _ _ _ } _ _.
Arguments Compose_NaturalTransformations_H {_ _ _ _ _ } _ _.

Definition CategoryInd (I: @obj TypeCat): Category.
Proof. unshelve econstructor.
       - exact (forall i: I,  @obj TypeCat).
       - simpl. intros X Y.
         exact (forall i, X i -> Y i).
       - simpl. intros X i.
         intro a.
         exact a.
       - simpl. intros X Y Z f g i.
         intro z.
         exact (g i (f i z)).
       - repeat intro. now subst.
       - simpl.
         intros X Y Z W f g h.
         easy.
       - simpl. intros X Y f.
         easy.
       - simpl. intros X Y f.
         easy.
Defined.

Definition FunctorInd (I J: @obj TypeCat) (p: @arrow TypeCat J I):
  Functor (CategoryInd J) (CategoryInd I).
Proof. unshelve econstructor.
       - simpl in *.
         intros f i.
         exact (f (p i)).
       - simpl in *.
         intros a b f i d.
         exact (f (p i) d).
       - repeat intro. now subst.
       - simpl in *.
         intro a.
         easy.
       - simpl in *.
         intros a b c g f.
         easy.
Defined.

Lemma TypeCatProd: forall A B: @obj TypeCat, Product TypeCat A B (prod A B).
Proof. intros.
       unshelve econstructor.
       - simpl. intros (a, b). exact a.
       - simpl. intros (a, b). exact b.
       - simpl. intros Z f g z. exact (f z, g z).
       - simpl. intros Z f g. easy.
       - simpl. intros Z f g. easy.
       - simpl. intros Z f g p1 H1 H2.
         pose proof (fun x => eq_ind_r (fun f => f x = _ x ) eq_refl H1) as Hax1';
         cbv beta in Hax1'.
         pose proof (fun x => eq_ind_r (fun f => f x = _ x ) eq_refl H2) as Hax2';
         cbv beta in Hax2'.
(*          pose proof (fun x => eq_ind_r (fun f => f x = _ x ) eq_refl H3) as Hax3';
         cbv beta in Hax3'.
         pose proof (fun x => eq_ind_r (fun f => f x = _ x ) eq_refl H4) as Hax4';
         cbv beta in Hax4'. *)
         apply functional_extensionality.
         intro z.
         specialize (Hax1' z).
         specialize (Hax2' z).
         destruct (p1 z).
         rewrite <- Hax1', <- Hax2'.
         reflexivity.
Defined.

Lemma TypeCathasProd: hasProducts TypeCat.
Proof. unshelve econstructor.
       - intros a b.
         exact (prod a b).
       - intros. exact (TypeCatProd A B).
Defined.

Definition pi1Ind (I J: @obj TypeCat):
  Functor (CategoryInd I) (CategoryInd (@pobj TypeCat TypeCathasProd I J)).
Proof. unshelve econstructor.
       - simpl in *.
         intros f (i, j).
         exact (f i).
       - simpl in *.
         intros a b f (i, j) p.
         exact (f i p).
       - repeat intro. now subst.
       - simpl in *.
         intros a.
         apply functional_extensionality_dep.
         intros (i, j).
         easy.
       - simpl in *.
         intros a b c f g.
         apply functional_extensionality_dep.
         intros (i, j).
         easy.
Defined.

Lemma pi1IndLA: forall (I J: @obj TypeCat),
  { G: Functor (CategoryInd (@pobj TypeCat TypeCathasProd I J)) (CategoryInd I) & HomAdjunction (pi1Ind I J) G }.
Proof. intros.
       specialize (UP (CategoryInd I)
                      (CategoryInd (@pobj TypeCat TypeCathasProd I J))
                      (pi1Ind I J)
                  ); intros H.
       destruct H as (HL, HR).
       apply HL.
       clear HL HR.
       intro E.
       simpl in *.
       unshelve econstructor.
       - intro i.
         exact {j: J & E (i, j)}.
       - simpl.
         unshelve econstructor.
         + intros (i, j) y.
           exists j.
           exact y.
         + simpl in *.
           intros X g.
           unshelve econstructor.
           ++ intros i (j, Hj).
              exact (g (i, j) Hj).
           ++ cbn in *.
              split.
              * apply functional_extensionality_dep.
                intros (i, j).
                easy.
              * intros.
                apply functional_extensionality_dep.
                intro i.
                apply functional_extensionality_dep.
                intros (j, Hj).
                rewrite <- H.
                easy.
Qed.


From Classes Require Import Foldable.
From Equations Require Import Equations.
From stdpp Require Import numbers list.

Module Values.

  Open Scope N_scope.

  Definition values X := list (N * X).

  Fixpoint lookup {X} (v : values X) (pos : N) : option X :=
    match v with
    | [] => None
    | (k,v) :: t => if k =? pos then Some v else lookup t pos
    end.

  Fixpoint set {X} (vs : values X) (pos : N) (value : X) : values X :=
    match vs with
    | [] => (pos, value) :: []
    | (k, v) :: t =>
        match N.eq_dec k pos with
        | left _ => (k, value) :: t
        | right _ => (k, v) :: set t pos value
        end
    end.

  Fixpoint size {X} (vs : values X) : N :=
    match vs with
    | [] => 0
    | (k, v) :: t =>
        N.succ k `max` (size t)
    end.

  Lemma set_size : forall A (vs : values A) pos v,
      size (set vs pos v) = size vs `max` (N.succ pos).
  Proof.
    induction vs.
    - intros. simpl. apply N.max_comm.
    - intros pos v. simpl. destruct a.
      destruct (N.eq_dec n pos).
      + simpl. subst. rewrite <- N.max_comm. rewrite <- N.max_assoc.
        rewrite N.max_id. reflexivity.
      + simpl. rewrite <- N.max_assoc. f_equal. apply IHvs.
  Defined.

  Global Instance Foldable_Values : Foldable values.
  econstructor. intros. destruct sg, H.
  eapply list.foldr. intros c b. eapply (f b (X c.2)). eapply mempty. eapply X0.
  Defined.

  Lemma values_size : forall X (l : Values.values X),
      0 < Values.size l ->
      exists v, Values.lookup l (Values.size l - 1) = Some v.
  Proof.
    induction l; simpl; intros.
    - lia.
    - destruct a.
      destruct (n =? N.succ n `max` Values.size l - 1) eqn:?. exists x. reflexivity.
      eapply N.eqb_neq in Heqb. revert H Heqb. eapply N.max_case.
      + intros P IH. rewrite <- N.pred_sub in IH. rewrite N.pred_succ in IH.
        contradiction IH. auto.
      + intros P IH. apply IHl in P. exact P.
  Qed.

  Lemma size_add : forall A vs (pos : N) (v : A),
      size vs <= pos ->
      Values.set vs pos v = vs ++ [(pos, v)].
  Proof.
    induction vs; simpl in *.
    - intros. reflexivity.
    - intros pos v LE. destruct a.
      destruct (N.eq_dec n pos).
      + subst. lia.
      + rewrite IHvs. reflexivity. lia.
  Qed.

  Close Scope N_scope.

End Values.

Import Values.

Open Scope N_scope.

Definition lt_dec (n0 n1 : N) : {n0 < n1} + {n1 <= n0}.
  destruct (N.ltb n0 n1) eqn:P.
  - apply N.ltb_lt in P. left. exact P.
  - apply N.ltb_ge in P. right. exact P.
Defined.

Record vector A :=
  mk_vector {
      capacity : N;
      size : N;
      values : values A
    }.

Arguments capacity {A}.
Arguments size {A}.
Arguments values {A}.

Structure Vector {A} (vec : vector A) :=
  properties_vector {
      Size_wf : Values.size (values vec) = size vec;
      Capacity_wf : size vec <= capacity vec
    }.

Definition VECTOR A := {arr : vector A | Vector arr}.

Equations make (A : Type) (n : N) : VECTOR A :=
  make A n := exist _ (mk_vector A n N0 []) _.
Next Obligation.
  econstructor; simpl in *.
  reflexivity.
  eapply N.le_0_l.
Defined.

Equations set {A} (vec : VECTOR A) (pos : N) (v : A) : VECTOR A by wf (N.to_nat (N.succ pos - capacity (`vec))) lt :=
  set vec pos v :=
    let cap := capacity (`vec) in
    match lt_dec pos cap with
    | right _ =>
        let arr := exist _ (mk_vector A (cap + ((cap / 2) `max` 1)) (size (`vec)) (values (`vec))) _ in
        set arr pos v
    | left _ =>
        let size := N.succ pos `max` size (`vec) in
        exist _ (mk_vector A (capacity (`vec)) size (Values.set (values (`vec)) pos v)) _
    end.
Next Obligation.
  destruct vec as [vec Vec]. destruct Vec.
  constructor; simpl in *.
  - rewrite <- Size_wf0. clear Size_wf0. induction (values vec). simpl in *. reflexivity.
    simpl. destruct a. destruct (N.eq_dec n pos).
    + subst. simpl. rewrite N.max_assoc. rewrite N.max_id. reflexivity.
    + simpl in * . rewrite IHv0. do 2 (rewrite N.max_assoc). f_equal.
      rewrite N.max_comm. reflexivity.
  - eapply N.max_lub_iff. split. eapply N.le_succ_l. exact l. eapply Capacity_wf0.
Defined.
Next Obligation.
  destruct vec as [vec Vec]. destruct Vec.
  econstructor; simpl in *.
  - exact Size_wf0.
  - etransitivity. eapply Capacity_wf0. eapply N.le_add_r.
Defined.
Next Obligation.
  intros. simpl. lia.
Defined.

Global Transparent make.
Global Transparent set.

Definition test_perf0 := make N 10.

(* Compute test_perf0. *)

Definition test_perf_test :=
  let vec := make N 10 in
  set vec 10000 10.

(* Compute test_perf_test. *)

Definition add {A} (vec : VECTOR A) (v : A) : VECTOR A := set vec (size (`vec)) v.

Definition test_perf_test_suite :=
  add (add (add test_perf_test 10) 8) 0.

(* Compute test_perf_test_suite. *)

Definition get {A} (vec : VECTOR A) (pos : N) : option A := lookup (values (`vec)) pos.

Fixpoint list_to_vector_aux {X} (l : list X) (vec : VECTOR X) :=
  match l with
  | [] => vec
  | h :: t => list_to_vector_aux t (add vec h)
  end.

Definition list_to_vector {X} (l : list X) :=
  list_to_vector_aux l (Vector.make X 10).

Definition vector_to_list {X} (vec : VECTOR X) : list (N * X) := Vector.values (`vec).

Notation "'[]↓' vec" := (vector_to_list vec)(at level 20).

Definition eq_dec_option {X} (comp : ∀ x y : X, {x = y} + {x ≠ y})
           (x y : option X) : {x = y} + {x ≠ y}.
Proof.
  destruct x,y.
  - destruct (comp x x0) as [eq|neq].
    + left. f_equal. exact eq.
    + right. intro some. eapply neq. eapply Some_inj in some. exact some.
  - right. intro neq. eapply Some_ne_None. exact neq.
  - right. intro neq. eapply None_ne_Some. exact neq.
  - left. reflexivity.
Defined.

Definition eqb_option {X} (comp : ∀ x y : X, {x = y} + {x ≠ y})
           (x y : option X) : bool :=
  proj1_sig (Sumbool.bool_of_sumbool (eq_dec_option comp x y)).

Equations ext_eq_aux {X} (comp : ∀ x y : X, {x = y} + {x ≠ y})
          (a1 a2 : Values.values X) (n : N) : bool by wf (N.to_nat n) lt :=
  ext_eq_aux comp a1 a2 0 :=
    eqb_option comp (lookup a1 0) (lookup a2 0);
  ext_eq_aux comp a1 a2 n :=
    andb (eqb_option comp (lookup a1 n) (lookup a2 n)) (ext_eq_aux comp a1 a2 (N.pred n)).
Next Obligation.
  lia.
Defined.

Definition ext_eq {X} (comp : ∀ x y : X, {x = y} + {x ≠ y}) (a1 a2 : VECTOR X) : bool :=
  if size (`a1) =? size (`a2)
  then
    ext_eq_aux comp (values (`a1)) (values (`a2)) (size (`a1))
  else
    false.

Global Instance Foldable_vector : Foldable VECTOR.
econstructor.
intros M sg m A f vec. destruct sg. destruct m.
eapply foldr. intros a b. eapply (f0 (f a.2) b). eapply mempty. eapply (vector_to_list vec).
Defined.

Global Instance Foldable_Vector `{Foldable M}: Foldable (fun X => VECTOR (M X)).
econstructor.
intros M0 sg m0 A f vec. destruct H. destruct sg.
destruct Foldable_vector.
eapply foldMap0. eapply m0.
intros ma. eapply foldMap. eapply m0. eapply f. eapply ma. eapply vec.
Defined.

Close Scope N_scope.

Lemma set_vector_list : forall (pos : N) A (vec : VECTOR A) v,
  ∃ n : N, []↓Vector.set vec pos v = match Vector.lt_dec pos (Vector.size (`vec)) with
                                     | left _ => []↓Vector.set vec pos v
                                     | right _ => []↓vec ++  [(n,v)]
                                     end.
Proof.
  intros pos A vec v. destruct (Vector.lt_dec pos (size (`vec))).
  exists pos. reflexivity.
  exists pos. revert l. eapply Vector.FunctionalElimination_set. clear pos A v vec.
  intros A vec pos v IH size_pos P0 P1. simpl in *.
  destruct (Vector.lt_dec pos (Vector.capacity (`vec))).
  - destruct vec as [vec Vec]. simpl in *. destruct Vec. unfold vector_to_list.
    eapply size_add. simpl in *. lia.
  - eapply P0; auto.
Qed.

Lemma add_vector_list : forall A (v :A) (vec : VECTOR A),
  exists n, []↓add vec v = []↓vec ++ [(n,v)].
Proof.
  intros A v vec. edestruct (set_vector_list (Vector.size (`vec)) A vec v).
  exists x. unfold Vector.add. rewrite H.
  destruct (Vector.lt_dec (size (`vec)) (size (`vec))). lia.
  reflexivity.
Qed.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas64_2e511m187_10limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition carry :
  { carry : feBW_loose -> feBW_tight
  | forall a, phiBW_tight (carry a) = (phiBW_loose a) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_carry ().
  Show Ltac Profile.
Time Defined.

Print Assumptions carry.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas32_2e222m117_10limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition square :
  { square : feBW_loose -> feBW_tight
  | forall a, phiBW_tight (square a) = F.mul (phiBW_loose a) (phiBW_loose a) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_square ().
  Show Ltac Profile.
Time Defined.

Print Assumptions square.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.montgomery32_2e256m189_8limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition mul :
  { mul : feBW_small -> feBW_small -> feBW_small
  | forall a b, phiM_small (mul a b) = F.mul (phiM_small a) (phiM_small b) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_mul ().
  Show Ltac Profile.
Time Defined.

Print Assumptions mul.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas32_2e511m481_23limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition square :
  { square : feBW_loose -> feBW_tight
  | forall a, phiBW_tight (square a) = F.mul (phiBW_loose a) (phiBW_loose a) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_square ().
  Show Ltac Profile.
Time Defined.

Print Assumptions square.


Load LFindLoad.
From lfind Require Import LFind.
Unset Printing Notations.
Set Printing Implicit.

Require Import Nat.

Inductive lst : Type :=
  | Nil : lst
  | Cons : nat -> lst -> lst.

Inductive queue : Type :=
  | Queue : lst -> lst -> queue.

Fixpoint len (l : lst) : nat :=
match l with
  | Nil => 0
  | Cons a l1 => 1 + (len l1)
end.

Definition qlen (q : queue) : nat :=
match q with
  | Queue l1 l2 => (len l1) + (len l2)
end.

Fixpoint app (l : lst) (m: lst): lst :=
match l with
  | Nil => m
  | Cons a l1 => Cons a (app l1 m)
end.

Fixpoint rev (l: lst): lst :=
match l with
  | Nil => Nil
  | Cons a l1 => app (rev l1) (Cons a Nil)
end.

Fixpoint leb (n m : nat) : bool :=
match (n, m) with
  | (0, _) => true
  | (S n', S m') => leb n' m'
  | _ => false
end.

Definition amortizeQueue (l1 l2 : lst) : queue :=
  if leb (len l2)  (len l1) then Queue l1 l2
  else Queue (app l1 (rev l2)) Nil.

Definition qpush (q : queue) (n : nat) : queue :=
match q with
  | Queue l1 l2 => amortizeQueue l1 (Cons n l2)
end.

Lemma len_app : forall l1 l2, len (app l1 l2) = (len l1) + (len l2).
Proof.
  intros. induction l1.
  - reflexivity.
  - simpl. rewrite IHl1. reflexivity.
Qed.

Lemma plus_comm: forall m n, m + n = n + m.
Proof.
induction m.
- intros. simpl. rewrite <- plus_n_O. reflexivity.
- intros. simpl. rewrite IHm. rewrite plus_n_Sm. reflexivity.
Qed.

Lemma len_rev : forall l, len (rev l) = len l.
Proof.
induction l.
- reflexivity.
- simpl. rewrite len_app. simpl. rewrite plus_comm. simpl. rewrite IHl. reflexivity.
Qed.

Theorem queue_push : forall q n, qlen (qpush q n) = 1 + (qlen q).
Proof.
  intros. simpl. destruct q. unfold qpush. unfold qlen. unfold amortizeQueue.
  destruct (leb (len (Cons n l0)) (len l)).
  - simpl. rewrite plus_n_Sm. reflexivity.
  - simpl. rewrite <- plus_n_O.
    lfind. Admitted.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas64_2e389m21_7limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition carry :
  { carry : feBW_loose -> feBW_tight
  | forall a, phiBW_tight (carry a) = (phiBW_loose a) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_carry ().
  Show Ltac Profile.
Time Defined.

Print Assumptions carry.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.montgomery64_2e235m15_4limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition sub :
  { sub : feBW_small -> feBW_small -> feBW_small
  | forall a b, phiM_small (sub a b) = F.sub (phiM_small a) (phiM_small b) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_sub ().
  Show Ltac Profile.
Time Defined.

Print Assumptions sub.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.solinas64_2e468m17_9limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition carry :
  { carry : feBW_loose -> feBW_tight
  | forall a, phiBW_tight (carry a) = (phiBW_loose a) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_carry ().
  Show Ltac Profile.
Time Defined.

Print Assumptions carry.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.montgomery32_2e171m19_6limbs.Synthesis.

(* TODO : change this to field once field isomorphism happens *)
Definition sub :
  { sub : feBW_small -> feBW_small -> feBW_small
  | forall a b, phiM_small (sub a b) = F.sub (phiM_small a) (phiM_small b) }.
Proof.
  Set Ltac Profiling.
  Time synthesize_sub ().
  Show Ltac Profile.
Time Defined.

Print Assumptions sub.


Require Import
  ssreflect ssrfun ssrbool eqtype ssrnat seq fintype div bigop prime binomial.

Set Implicit Arguments.
Unset Strict Implicit.
Import Prenex Implicits.

Lemma primeP' p :
  reflect (1 < p /\ forall p', p' < p -> prime p' -> ~ p' %| p) (prime p).
Proof.
  case/boolP: (prime p) => H; constructor; [split | case].
  - by move: p H; do 2 case => //.
  - case/primeP: H => H H0 p' H1 H2; case/(H0 p')/orP.
    + by move: p' H2 {H1}; do 2 case => //.
    + by rewrite (ltn_eqF H1).
  - move => H0 H1; case/primePn: H; first by move: p H0 {H1}; do 2 case => //.
    case => x; case/andP => H2 H3 H4; apply (H1 (pdiv x)).
    + by apply leq_ltn_trans with x => //; apply pdiv_leq, ltnW.
    + by apply pdiv_prime.
    + by apply dvdn_trans with x => //; apply pdiv_dvd.
Qed.

Fixpoint search_prime n m :=
  if prime n
    then n
    else if m is m.+1
      then search_prime n.+1 m
      else 0.

Definition next_prime n := search_prime n n`!.+1.

Lemma next_prime_is_prime n : prime (next_prime n).
Proof.
  suff: next_prime n != 0 by
    rewrite /next_prime; move: n`!.+1 => m;
      elim: m n => /= [| m IH] n; case: ifP => //= _; apply IH.
  rewrite /next_prime; apply/negP => H.
  have {H} H p : n <= p <= n`!.+1 + n -> ~~ prime p.
    move: H; move: n`!.+1 => m; elim: m n => /= [| m IH] n.
    - by rewrite add0n -eqn_leq => H; move/eqP => <-;
        apply/negP => H0; move: H; rewrite H0 /=; case: n H0.
    - case: ifP; first by case: n.
      move => H; move/IH => {IH} H0; rewrite leq_eqVlt; case/andP; case/orP.
      + by move/eqP => <-; rewrite H.
      + by rewrite addSnnS => H1 H2; apply/H0/andP.
  suff: prime n`!.+1
    by apply/negP/H; rewrite leq_addr andbT {H}; apply leqW;
      case: n => //= n; rewrite factS -{1}(muln1 n.+1) leq_mul2l fact_gt0.
  apply/primeP'; split; first by rewrite ltnS fact_gt0.
  move => p'; rewrite ltnS => H0 H1 H2; case (leqP p' n) => H3;
    last by move: (H p'); rewrite (ltnW H3);
      move/(_ (leq_trans (leqW H0) (leq_addr _ _)))/negP.
  move: H2; rewrite -addn1 dvdn_addr;
    first by rewrite dvdn1; move/eqP => ?; subst p'; move: H1.
  by rewrite -(ffact_fact (leq_subr p' n)) subKn //; apply dvdn_mull;
    case: p' H1 {H H0 H3} => // p' _; rewrite factS; apply dvdn_mulr.
Qed.

Lemma next_prime_smallest n p : n <= p < next_prime n -> ~~ prime p.
Proof.
  move: (next_prime_is_prime n).
  rewrite /next_prime; move: n`!.+1 => m.
  elim: m n => //= [| m IH] n; case: ifP => //=;
    try by move => _ _; case/andP;
           rewrite -ltnS => H; move/(leq_trans H); rewrite ltnn.
  move => H; move/IH => H0; rewrite leq_eqVlt; case/andP; case/orP.
  - by move/eqP => <- _; rewrite H.
  - by move => H1 H2; apply H0; rewrite H1.
Qed.

Lemma next_prime_spec3 n : n <= next_prime n.
Proof.
  by move: (next_prime_is_prime n); rewrite /next_prime;
    move: _.+1 => m; elim: m n => /= [| m IH] n;
    case: ifP => // H; move/IH; rewrite (leq_eqVlt n) orbC => ->.
Qed.


Set Warnings "-notation-overridden,-parsing,-deprecated-hint-without-locality".
From Coq Require Import Lia.
From LF Require Import Maps.
From LF Require Import Imp2.

Example auto_example_1' : forall (P Q R: Prop),
  (P -> Q) -> (Q -> R) -> P -> R.
Proof.
  (* auto *) info_auto.
Qed.

Lemma le_antisym : forall n m: nat, (n <= m /\ m <= n) -> n = m.
Proof. lia. Qed.

Example auto_example_6 : forall n m p : nat,
  (n <= p -> (n <= m /\ m <= n)) ->
  (n <= p -> n = m).
Proof.
  (* auto cant solve it without additional hypothesis *)
  auto using le_antisym.
Qed.

(*
  Hints for auto command

  Hint Resolve T : core.
  Hint Constructors c : core.
  Hint Unfold d : core.
*)

Hint Resolve le_antisym : core.

Example auto_example_6' : forall n m p : nat,
  (n <= p -> (n <= m /\ m <= n)) ->
  (n <= p -> n = m).
Proof. auto. Qed.

Import Imp2.AExp.
Import Maps.

Theorem ceval_deterministic'_alt: forall c st st1 st2,
  st =[ c ]=> st1 ->
  st =[ c ]=> st2 ->
  st1 = st2.
Proof with auto.                 (* adds ... shortcut for auto *)
  intros c st st1 st2 E1 E2;
  generalize dependent st2;
  induction E1;
           intros st2 E2; inversion E2; subst...
  - (* E_Seq *)
    rewrite (IHE1_1 st'0 H1) in *...
  - (* E_IfTrue *)
    + (* b evaluates to false (contradiction) *)
      rewrite H in H5. discriminate.
  - (* E_IfFalse *)
    + (* b evaluates to true (contradiction) *)
      rewrite H in H5. discriminate.
  - (* E_WhileFalse *)
    + (* b evaluates to true (contradiction) *)
      rewrite H in H2. discriminate.
  (* E_WhileTrue *)
  - (* b evaluates to false (contradiction) *)
    rewrite H in H4. discriminate.
  - (* b evaluates to true *)
    rewrite (IHE1_1 st'0 H3) in *...
Qed.

Ltac rwd H1 H2 := rewrite H1 in H2; discriminate.

Theorem ceval_deterministic'': forall c st st1 st2,
  st =[ c ]=> st1 ->
  st =[ c ]=> st2 ->
  st1 = st2.
Proof.
  intros c st st1 st2 E1 E2.
  generalize dependent st2;
  induction E1; intros st2 E2; inversion E2; subst; auto.
  - (* E_Seq *)
    rewrite (IHE1_1 st'0 H1) in *.
    auto.
  - (* E_IfTrue *)
    + (* b evaluates to false (contradiction) *)
      rwd H H5.
  - (* E_IfFalse *)
    + (* b evaluates to true (contradiction) *)
      rwd H H5.
  - (* E_WhileFalse *)
    + (* b evaluates to true (contradiction) *)
      rwd H H2.
  (* E_WhileTrue *)
  - (* b evaluates to false (contradiction) *)
    rwd H H4.
  - (* b evaluates to true *)
    rewrite (IHE1_1 st'0 H3) in *.
    auto. Qed.

Ltac find_rwd :=
  match goal with
    H1: ?E = true,
    H2: ?E = false
    |- _ => rwd H1 H2
  end.

Theorem ceval_deterministic''': forall c st st1 st2,
  st =[ c ]=> st1 ->
  st =[ c ]=> st2 ->
  st1 = st2.
Proof.
  intros c st st1 st2 E1 E2.
  generalize dependent st2;
  induction E1; intros st2 E2; inversion E2; subst; try find_rwd; auto.
  - (* E_Seq *)
    rewrite (IHE1_1 st'0 H1) in *.
    auto.
  - (* E_WhileTrue *)
    + (* b evaluates to true *)
      rewrite (IHE1_1 st'0 H3) in *.
      auto. Qed.

Ltac find_eqn :=
  match goal with
    H1: forall x, ?P x -> ?L = ?R,
    H2: ?P ?X
    |- _ => rewrite (H1 X H2) in *
  end.

Theorem ceval_deterministic'''': forall c st st1 st2,
  st =[ c ]=> st1 ->
  st =[ c ]=> st2 ->
  st1 = st2.
Proof.
  intros c st st1 st2 E1 E2.
  generalize dependent st2;
  induction E1; intros st2 E2; inversion E2; subst; try find_rwd;
    try find_eqn; auto.
Qed.

Example ceval'_example1:
  empty_st =[
    X := 2;
    if (X <= 1)
      then Y := 3
      else Z := 4
    end
  ]=> (Z !-> 4 ; X !-> 2).
Proof.
  eapply E_Seq. (* 1 *)
  - apply E_Asgn. (* 2 *)
    reflexivity. (* 3 *)
  - (* 4 *) apply E_IfFalse. reflexivity. apply E_Asgn. reflexivity.
Qed.


Global Set Implicit Arguments.
Global Unset Strict Implicit.
Global Unset Printing Records.
Global Unset Printing Implicit Defensive.
Global Set Regular Subst Tactic.

Set Implicit Arguments.

Hint Extern 4 => exact _.

Require Export
        Arith_base
        List
        Tactics.Tactics
        Setoid
        Morphisms.

(** basic stuff *)

Lemma neg_or A B: ~ (A \/ B) -> ~ A /\ ~ B.
Proof.
  tauto.
Qed.

Lemma DM_not_exists X (p : X -> Prop) :
  ~ (exists x, p x) <-> forall x, ~ p x.
Proof.
  firstorder.
Qed.

Ltac sum_solver :=
  try solve [left ; congruence] ;
  try solve [right ; congruence].

(** definition of a type class for decidable
    predicates *)

Definition dec (P : Prop) : Type := {P} + {~ P}.

Notation "'eq_dec' A" := (forall x y : A, dec (x = y))(at level 70).

Existing Class dec.

Definition decision (P : Prop){D : dec P} : dec P := D.
Arguments decision P {D}.

Ltac dec := repeat (destruct decision).

Ltac deq x := destruct (decision (x=x)) as [[]  | isnotequal] ;
              [> | contradict isnotequal; reflexivity] .

Tactic Notation "decide" constr(p) :=
  destruct (decision p).
Tactic Notation "decide" constr(p) "as" simple_intropattern(i) :=
  destruct (decision p) as i.

(** coercion from Prop to bool *)

Coercion toBool (P: Prop) (D: dec P) := if decision P then true else false.
Arguments toBool P {D}.

Definition injective (A B : Type) (f : A -> B):= forall x y, f x = f y -> x = y.
Definition surjective (A B : Type) (f : A -> B) : Prop := forall y, exists x, f x = y.
Definition bijective (A B : Type) (f : A -> B): Prop := injective f /\ surjective f.

Lemma injective_apply (A B : Type) (f : A -> B) (x y : A) :
  injective f -> f x = f y -> x = y.
Proof.
  auto.
Qed.

Lemma surjective_apply (A B : Type) (f : A -> B) (x y : A) :
  surjective f -> forall y, exists x, f x = y.
Proof.
  auto.
Qed.

Lemma loop (A B : Type) (f : A -> B) (g : B -> A)
  : (forall x, g (f x) = x) -> injective f.
Proof.
  intro H. unfold injective. intros x  y G.
  pose proof (H x) as H'. specialize (H y).
  rewrite G in H'.
     subst x. exact H.
Qed.

(** improving automation *)

Hint Extern 4 =>
    match goal with
    | [ |- dec ?p ] => exact (decision p)
    end.

Hint Extern 4 =>
    match goal with
    | [ |- dec ((fun _ => _) _) ] => simpl
    end : typeclass_instances.

(** basic instances for decidabitity *)

Fact dec_trans P Q :
  dec P -> P <-> Q -> dec Q.
Proof.
  unfold dec. tauto.
Qed.

Instance True_dec :
  dec True.
Proof.
  unfold dec ; auto.
Qed.

Instance False_dec :
  dec False.
Proof.
  unfold dec ; auto.
Qed.

Instance impl_dec (X Y : Prop) :
  dec X -> dec Y -> dec (X -> Y).
Proof.
  unfold dec ; tauto.
Defined.

Instance and_dec (X Y : Prop) :
  dec X -> dec Y -> dec (X /\ Y).
Proof.
  unfold dec ; tauto.
Defined.

Instance or_dec (X Y : Prop) :
  dec X -> dec Y -> dec (X \/ Y).
Proof.
  unfold dec ; tauto.
Defined.

Instance not_dec (X : Prop) :
  dec X -> dec (~ X).
Proof.
  unfold not. auto.
Defined.

Instance iff_dec (X Y : Prop) :
  dec X -> dec Y -> dec (X <-> Y).
Proof.
  unfold iff. auto.
Qed.

Lemma dec_DN X :
  dec X -> ~~ X -> X.
Proof.
  unfold dec ; tauto.
Qed.

Lemma dec_DM_and X Y :
  dec X -> dec Y -> ~ (X /\ Y) -> ~ X \/ ~ Y.
Proof.
  unfold dec ; tauto.
Qed.

Lemma dec_DM_impl X Y :
  dec X -> dec Y -> ~ (X -> Y) -> X /\ ~ Y.
Proof.
  unfold dec ; tauto.
Qed.

Lemma dec_prop_iff (X Y : Prop) :
  (X <-> Y) -> dec X -> dec Y.
Proof.
  unfold dec ; tauto.
Defined.

Instance unit_eq_dec : eq_dec unit.
Proof.
  intros x y. hnf. decide equality.
Defined.

Instance Empty_set_eq_dec:
  eq_dec Empty_set.
Proof.
  tauto.
Qed.

Instance True_eq_dec:
  eq_dec True.
Proof.
  intros x y. destruct x, y. left*.
Qed.

Instance False_eq_dec:
  eq_dec False.
Proof.
  intros x  ; destruct x.
Qed.

Instance bool_eq_dec :
  eq_dec bool.
Proof.
  intros x y. hnf. decide equality.
Defined.

Instance nat_eq_dec :
  eq_dec nat.
Proof.
  intros x y. hnf. decide equality.
Defined.

Instance list_eq_dec X :
  eq_dec X -> eq_dec (list X).
Proof.
  intros D. apply list_eq_dec. exact D.
Qed.

Instance nat_le_dec (x y : nat) : dec (x <= y) :=
  le_dec x y.

Instance prod_eq_dec (T1 T2 : Type) (E1 : eq_dec T1) (E2 : eq_dec T2):
  eq_dec (T1 * T2).
Proof.
  intros. unfold dec.  destruct x, y.
  decide (t = t1) ; decide (t0 = t2) ; sum_solver.
Qed.

Instance option_eq_dec (X : Type) (e : eq_dec X) :
  eq_dec (option X).
Proof.
  intros. destruct x, y ; sum_solver.
  +
    decide (x = x0) ; sum_solver.
Qed.

Instance sum_eq_dec (X : Type) (Y : Type) (ex : eq_dec X) (ey : eq_dec Y):
  eq_dec (sum X Y).
Proof.
  intros. destruct x, y ; sum_solver.
  -
    decide (x = x0) ; sum_solver.
  -
    decide (y0 = y) ; sum_solver.
Qed.

Structure discType
  := DiscType {
        disctype :> Type
     ;  decide_eq : eq_dec disctype
     }.

Arguments DiscType disctype {decide_eq}.
Existing Instance decide_eq.

Definition toDiscType (A : Type) {D : eq_dec A} : discType := DiscType A.
Arguments toDiscType A {D}.

Lemma toDiscTypeCorrect (T : discType) : T = toDiscType T.
Proof.
  destruct* T.
Qed.

Lemma Hedberg (A : discType) (x y : A) (E E' : x = y) : E = E'.
Proof.
  pose (f := fun y : A => fun E: x = y => match (decision (x = y)) with
                                   | left b => b
                                   | _ => E end).
  apply (@injective_apply _ _(f y) E E').
  -
    pose (join := fun (y z : A) (e : x = y) => eq_trans (z := z) (eq_sym e)).
    apply  (@loop _ _ (f y)  (join x y (f x (eq_refl x)))).
    intros []. unfold f. now deq x.
  -
    unfold f.  decide (x = y); tauto.
Qed.

Require Import Logic.EqdepFacts.

Theorem Streicher_K (A : discType) (x : A) (P : x = x -> Prop)
  : P (eq_refl x) ->  forall p, P p.
Proof.
  intros H p. now rewrite (Hedberg p (eq_refl x)).
Qed.

Lemma dec_refl (A : discType) (x : A)
  : decision (x = x) = left eq_refl.
Proof.
  dec.
  -  f_equal.  now apply (@Streicher_K _ _  (fun e => e = eq_refl x) ).
  - congruence.
Qed.

Hint Resolve dec_refl.

Lemma eq_funTrans (X Y : Type) (f: X -> Y) (x x': X): x = x' -> f x = f x'.
Proof.
  crush.
Qed.

Canonical Structure DiscUnit := DiscType unit.
Canonical Structure DiscBool := DiscType bool.
Canonical Structure DiscNat := DiscType nat.
Canonical Structure DiscOption (T : discType)  := DiscType (option T).
Canonical Structure DiscProd (T1 T2 : discType) := DiscType (T1 * T2).
Canonical Structure DiscEmpty_set := DiscType Empty_set.
Canonical Structure DiscSum (X Y : discType) := DiscType (X + Y).
Canonical Structure DiscList (X : discType) := DiscType (list X).
Canonical Structure DiscFalse := DiscType False.
Canonical Structure DiscTrue := DiscType True.

Reserved Notation "T1 ** T2" (at level 69, no associativity).
Bind Scope DiscTypeScope with discType.
Notation "T1 ** T2" := (DiscProd T1 T2) : DiscTypeScope.
Notation "T1 ** T2" := (DiscProd (toDiscType T1) (toDiscType T2)).
Notation " ?? F" := (DiscOption F) (at level 65).

(** decidable predicates and relations *)

Structure dec_pred A
  := DecPred {
         predicate :> A -> Prop;
         decide_pred x : dec (predicate x)
     }.

Arguments DecPred {A} predicate {decide_pred}.

Existing Instance decide_pred.

Instance pred_dec A (p : dec_pred A) x :
  dec (p x).
Proof.
  apply decide_pred.
Qed.

Structure dec_relation A
  := DecRel {
         relation :> A -> A -> Prop;
         decide_rel x y : dec (relation x y)
     }.

Arguments DecRel {A} relation {decide_rel}.

Instance dec_relation_dec A (R : dec_relation A) x y :
  dec (R x y).
Proof.
  apply decide_rel.
Qed.

Lemma dec_DM_and' X Y :
  dec Y -> ~ (X /\ Y) -> ~ X \/ ~ Y.
Proof.
  unfold dec; tauto.
Qed.

Lemma dec_DM_all X (p : X -> Prop) (D : forall x, dec (p x)):
  (forall x, p x) <-> ~ exists x, ~ p x.
Proof.
  firstorder.
Qed.