Require Import ct01.

From Hammer Require Import Hammer.

Lemma sort_prog :
  forall (l : list nat), {l' : list nat | sorted l' /\ permutation l' l}.
Proof. (* try hammer. *) Abort.

Lemma sort_prog_base : {l' : list nat | sorted l' /\ permutation l' []}.
Proof. try hammer. Defined.

Lemma sort_prog_IH : forall (a : nat) (l x : list nat),
  sorted x -> permutation x l
  -> {l' : list nat | sorted l' /\ permutation l' (a :: l)}.
Proof. (* try hammer. *) Abort.


Require Import HoTT.Basics HoTT.Types.Bool HoTT.Types.Paths.
Require Import Colimits.Diagram Colimits.Colimit.

Local Open Scope path.

Section MappingTelescope.

  Definition mappingtelescope_graph : graph.
    simple refine (Build_graph _ _).
    - exact nat.
    - intros n m; exact (S n = m).
  Defined.

  Definition equiv_mappingtelescope_diag (D1 D2: diagram mappingtelescope_graph)
             (H0: (D1 0) <~> (D2 0))
             (Hn: forall n (e: (D1 n) <~> (D2 n)),
                   {e' : (D1 n.+1) <~> (D2 n.+1) & (D2 _f 1) o e == e' o (D1 _f 1)})
  : D1 ≃ D2.
    simple refine (Build_diagram_equiv (Build_diagram_map _ _) _).
    - intros n. apply equiv_fun. induction n. apply H0. exact (Hn n IHn).1.
    - intros n m q; destruct q. induction n; simpl. exact (Hn 0 H0).2.
      simple refine (Hn n.+1 _).2.
    - intros n; simpl. induction n; simpl. apply H0. apply (Hn n _ ).1.
  Defined.

End MappingTelescope.


Require Import Coq.ZArith.ZArith.
Require Import Coq.micromega.Lia.
Require Import coqutil.Word.Interface coqutil.Word.Properties.
Require Import coqutil.Z.BitOps coqutil.Z.ZLib.
Require Import coqutil.Tactics.destr.

Local Open Scope Z_scope.

Module word.

  Section RiscvWord.
    Context {width: Z} {word: word.word width}.
    Implicit Types x y z : word.

    (* TODO maybe we can put more fundamental axioms here, and turn the axioms below into lemmas *)
    Class riscv_ok: Prop := {
      sru_ignores_hibits: forall y z,
          word.sru y (word.of_Z (word.unsigned z mod 2 ^ Z.log2 width)) = word.sru y z;
      slu_ignores_hibits: forall y z,
          word.slu y (word.of_Z (word.unsigned z mod 2 ^ Z.log2 width)) = word.slu y z;
      srs_ignores_hibits: forall y z,
          word.srs y (word.of_Z (word.unsigned z mod 2 ^ Z.log2 width)) = word.srs y z;

      divu0_simpl: forall y z,
          (if word.eqb z (word.of_Z 0) then word.of_Z (2 ^ width - 1) else word.divu y z) =
          word.divu y z;
      modu0_simpl: forall y z,
          (if word.eqb z (word.of_Z 0) then y else word.modu y z) =
          word.modu y z;
    }.

    Lemma mulhuu_simpl{ok: word.ok word}{rok: riscv_ok}: forall y z,
        word.of_Z (bitSlice (word.unsigned y * word.unsigned z) width (2 * width)) =
        word.mulhuu y z.
    Proof.
      intros. unfold bitSlice.
      eapply word.unsigned_inj.
      rewrite word.unsigned_mulhuu. unfold word.wrap.
      replace (2 * width - width) with width by lia.
      pose proof word.width_pos.
      rewrite Z.shiftl_mul_pow2 by lia.
      rewrite (Z.mul_comm (-1)).
      rewrite <- Z.opp_eq_mul_m1.
      replace (Z.lnot (- 2 ^ width)) with (Z.pred (2 ^ width)). 2: {
        pose proof (Z.add_lnot_diag (- 2 ^ width)). lia.
      }
      rewrite <- Z.ones_equiv.
      rewrite Z.land_ones by lia.
      rewrite Z.shiftr_div_pow2 by lia.
      apply word.unsigned_of_Z_nowrap.
      apply Z.mod_pos_bound.
      apply Z.pow2_pos.
      lia.
    Qed.

  End RiscvWord.

  Arguments riscv_ok {_} _.

End word.


Require Import Setoid.
Require Import Morphisms.
Require Import Coq.Program.Basics.
Require Import NArith.
Require Import List.
Import ListNotations.
Open Scope list.
Require Import Lia.

Unset Printing Records.

From Ordinal Require Import Defs.
From Ordinal Require Import Operators.
From Ordinal Require Import Arith.
From Ordinal Require Import Cantor.
From Ordinal Require Import Fixpoints.
From Ordinal Require Import Reflection.
From Ordinal Require Import VeblenDefs.
From Ordinal Require Import VeblenCon.
From Ordinal Require Import VeblenFacts.

From Ordinal Require Import Notation.CantorDecomposition.

Open Scope ord_scope.

Inductive VForm : Set :=
| Z : VForm
| V : VForm -> VForm -> VForm.

Fixpoint VF_denote (x:VForm) : Ord :=
  match x with
  | Z => 0
  | V a x => veblen (addOrd 1) (VF_denote a) (VF_denote x)
  end.

Canonical Structure VF := ord VForm VF_denote.

Fixpoint VF_normal (x:VForm) : Prop :=
  match x with
  | Z => True
  | V a x => VF_normal a /\ VF_normal x /\
             match x with
             | Z => True
             | V b y => VF_denote a >= VF_denote b
             end
  end.

Lemma VF_complete (x:VForm) : complete (VF_denote x).
Proof.
  induction x.
  apply zero_complete.
  simpl.
  apply veblen_complete; auto.
  apply onePlus_normal.
  intros; apply addOrd_complete; auto.
  apply succ_complete; apply zero_complete.
Qed.

Lemma veblen_onePlus_complete a x :
  complete a -> complete x -> complete (veblen (addOrd 1) a x).
Proof.
  intros; apply veblen_complete; auto.
  apply onePlus_normal.
  intros; apply addOrd_complete; auto.
  apply succ_complete. apply zero_complete.
Qed.

Local Hint Resolve VF_complete onePlus_normal veblen_onePlus_complete
      succ_complete zero_complete : core.

Lemma VF_denote_shrink1 : forall a x,
  VF_denote a < VF_denote (V a x).
Proof.
  simpl; intros.
  apply ord_lt_le_trans with (veblen (addOrd 1) (VF_denote a) 0).
  - clear x.
    induction a; simpl; intuition.
    { apply veblen_nonzero; auto. }
    apply veblen_shrink_lemma; auto.
  - apply veblen_monotone; auto with ord.
    intros; apply addOrd_monotone; auto with ord.
Qed.

Lemma VF_denote_le2 : forall x a,
  VF_denote x <= VF_denote (V a x).
Proof.
  induction x; simpl; intuition.
  apply veblen_inflationary; auto.
Qed.

Lemma VF_denote_shrink2 : forall x a,
  VF_normal (V a x) ->
  VF_denote x < VF_denote (V a x).
Proof.
  induction x; simpl; intuition.
  { apply veblen_nonzero; auto. }
  apply veblen_increasing'; auto.
  apply IHx2. simpl; intuition.
Qed.

Lemma VNotation_below_ε₀ x :
  VF_denote x < ε 0.
Proof.
  induction x; simpl VF_denote.
  - rewrite ε_fixpoint.
    apply expOrd_nonzero.
  - rewrite veblen_onePlus; auto.
    apply epslion_additively_closed; auto.
    rewrite ε_fixpoint.
    apply expOrd_increasing; auto.
    apply (index_lt _ 1%nat).
Qed.

Fixpoint VF_compare (x:VForm) : VForm -> ordering :=
  fix inner (y:VForm) : ordering :=
    match x, y with
    | Z, Z  => EQ
    | Z, V _ _ => LT
    | V _ _, Z => GT
    | V a x' , V b y' =>
      match VF_compare a b with
      | LT => VF_compare x' y
      | EQ => VF_compare x' y'
      | GT => inner y'
      end
    end.

Lemma VF_compare_correct : forall x y,
    match VF_compare x y with
    | LT => VF_denote x < VF_denote y
    | EQ => VF_denote x ≈ VF_denote y
    | GT => VF_denote x > VF_denote y
    end.
Proof.
  induction x as [|a IHa x IHx].
  { destruct y as [|b y]; simpl; auto with ord.
    apply veblen_nonzero.
    apply onePlus_normal. }
  induction y as [|b IHb y IHy]; simpl.
  { apply veblen_nonzero.  apply onePlus_normal. }
  generalize (IHa b).
  destruct (VF_compare a b); intro Hab.
  - generalize (IHx (V b y)).
    destruct (VF_compare x (V b y)); intro Hxsub.
    + apply ord_lt_le_trans with (veblen (addOrd 1) (VF_denote a) (veblen (addOrd 1) (VF_denote b) (VF_denote y))).
      apply veblen_increasing; auto.
      apply veblen_fixpoints; auto.

    + transitivity (veblen (addOrd 1) (VF_denote a) (veblen (addOrd 1) (VF_denote b) (VF_denote y))).
      { split; (apply veblen_monotone; [ intros; apply addOrd_monotone; auto with ord | apply Hxsub ]). }
      apply veblen_fixpoints; auto.

    + apply ord_lt_le_trans with (VF_denote x); auto.
      apply veblen_inflationary; auto.

  - generalize (IHx y).
    destruct (VF_compare x y); intro Hxy.
    + eapply ord_lt_le_trans. apply veblen_increasing; auto. apply Hxy.
      apply veblen_monotone_first; auto.
      intros; apply addOrd_monotone; auto with ord.
      apply Hab.
    + split.
      etransitivity.
      apply veblen_monotone_first; auto.
      intros; apply addOrd_monotone; auto with ord.
      apply Hab.
      apply veblen_monotone; auto.
      intros; apply addOrd_monotone; auto with ord.
      apply Hxy.
      etransitivity.
      apply veblen_monotone_first; auto.
      intros; apply addOrd_monotone; auto with ord.
      apply Hab.
      apply veblen_monotone; auto.
      intros; apply addOrd_monotone; auto with ord.
      apply Hxy.
    + eapply ord_lt_le_trans. apply veblen_increasing; auto. apply Hxy.
      apply veblen_monotone_first; auto.
      intros; apply addOrd_monotone; auto with ord.
      apply Hab.

  - change (match (VF_compare (V a x) y) with
            | LT =>
              veblen (addOrd 1) (VF_denote a) (VF_denote x) <
              veblen (addOrd 1) (VF_denote b) (VF_denote y)
            | EQ =>
              veblen (addOrd 1) (VF_denote a) (VF_denote x)
                     ≈ veblen (addOrd 1) (VF_denote b) (VF_denote y)
            | GT =>
              veblen (addOrd 1) (VF_denote b) (VF_denote y) <
              veblen (addOrd 1) (VF_denote a) (VF_denote x)
            end).
    destruct (VF_compare (V a x) y).
    + apply ord_lt_le_trans with (VF_denote y); auto.
      apply veblen_inflationary; auto.
    + symmetry.
      transitivity
        (veblen (addOrd 1) (VF_denote b) (veblen (addOrd 1) (VF_denote a) (VF_denote x))).
      { split; (apply veblen_monotone; [ intros; apply addOrd_monotone; auto with ord | apply IHy ]). }
      apply veblen_fixpoints; auto.
    + apply ord_lt_le_trans with (veblen (addOrd 1) (VF_denote b) (veblen (addOrd 1) (VF_denote a) (VF_denote x))).
      apply veblen_increasing; auto.
      apply veblen_fixpoints; auto.
Qed.

Theorem V_decide_order (x y:VF) : {x < y} + {x ≥ y}.
Proof.
  simpl sz.
  generalize (VF_compare_correct x y).
  destruct (VF_compare x y); intros.
  - left; assumption.
  - right; apply H.
  - right; auto with ord.
Qed.

Theorem V_normal_forms_unique : forall x y,
  VF_normal x ->
  VF_normal y ->
  VF_denote x ≈ VF_denote y ->
  x = y.
Proof.
  induction x as [|a Ha x Hx].
  - simpl. intro y; destruct y; simpl; auto.
    intros.
    elim (ord_lt_irreflexive 0).
    rewrite H1 at 2.
    apply veblen_nonzero; auto.
  - destruct y as [|b y].
    + simpl; intros.
      elim (ord_lt_irreflexive 0).
      rewrite <- H1 at 2.
      apply veblen_nonzero; auto.
    + simpl; intuition.

      cut (VF_denote a ≈ VF_denote b /\ VF_denote x ≈ VF_denote y).
      { intros [??]. f_equal; auto. }
      clear Ha Hx.

      assert (VF_denote a ≈ VF_denote b).
      { generalize (VF_compare_correct a b).
        destruct (VF_compare a b); intros; auto.
        * elim (ord_lt_irreflexive (veblen (addOrd 1) (VF_denote a) (VF_denote x))).
          rewrite H1 at 2.
          rewrite <- (veblen_fixpoints _ onePlus_normal (VF_denote a) (VF_denote b) (VF_denote y)); auto.
          apply veblen_increasing; auto.
          rewrite <- H1.
          apply VF_denote_shrink2. simpl; intuition.
        * elim (ord_lt_irreflexive (veblen (addOrd 1) (VF_denote b) (VF_denote y))).
          rewrite <- H1 at 2.
          rewrite <- (veblen_fixpoints _ onePlus_normal (VF_denote b) (VF_denote a) (VF_denote x)); auto.
          apply veblen_increasing; auto.
          rewrite H1.
          apply VF_denote_shrink2. simpl; intuition. }
      split; auto.

      generalize (VF_compare_correct x y).
      destruct (VF_compare x y); intros; auto.
      * elim (ord_lt_irreflexive (veblen (addOrd 1) (VF_denote a) (VF_denote x))).
        rewrite H1 at 2.
        apply ord_le_lt_trans with (veblen (addOrd 1) (VF_denote b) (VF_denote x)).
        apply veblen_monotone_first; auto.
        intros; apply addOrd_monotone; auto with ord.
        apply H4.
        apply veblen_increasing; auto.
      * elim (ord_lt_irreflexive (veblen (addOrd 1) (VF_denote b) (VF_denote y))).
        rewrite <- H1 at 2.
        apply ord_le_lt_trans with (veblen (addOrd 1) (VF_denote a) (VF_denote y)).
        apply veblen_monotone_first; auto.
        intros; apply addOrd_monotone; auto with ord.
        apply H4.
        apply veblen_increasing; auto.
Qed.

Definition Vnorm (a x:VForm) :=
  match x with
  | Z => V a Z
  | V b y =>
    match VF_compare a b with
    | GT => V a (V b y)
    | EQ => V a (V b y)
    | LT => V b y
    end
  end.

Lemma Vnorm_normal a x :
  VF_normal a -> VF_normal x -> VF_normal (Vnorm a x).
Proof.
  unfold Vnorm; simpl; intros.
  destruct x; simpl in *; intuition.
  generalize (VF_compare_correct a x1).
  destruct (VF_compare a x1); simpl; intuition.
  rewrite H2. reflexivity.
Qed.

Lemma Vnorm_V a x : VF_denote (V a x) ≈ VF_denote (Vnorm a x).
Proof.
  simpl. unfold Vnorm.
  destruct x; simpl; try reflexivity.
  generalize (VF_compare_correct a x1).
  destruct (VF_compare a x1); simpl; intuition.
  apply veblen_fixpoints; auto.
Qed.

Fixpoint Vnormalize (x:VForm) : VForm :=
  match x with
  | Z => Z
  | V a y => Vnorm (Vnormalize a) (Vnormalize y)
  end.

Lemma Vnormalize_normal : forall x, VF_normal (Vnormalize x).
Proof.
  induction x; simpl; intuition.
  apply Vnorm_normal; auto.
Qed.

Lemma Vnormalize_equal : forall x, VF_denote (Vnormalize x) ≈ VF_denote x.
Proof.
  induction x; simpl; intuition.
  rewrite <- Vnorm_V. simpl.
  transitivity (veblen (addOrd 1) (VF_denote x1) (VF_denote (Vnormalize x2))).
  - split; apply veblen_monotone_first; auto with ord.
    intros; apply addOrd_monotone; auto with ord.
    apply IHx1.
    intros; apply addOrd_monotone; auto with ord.
    apply IHx1.
  - split; apply veblen_monotone; auto with ord.
    intros; apply addOrd_monotone; auto with ord.
    apply IHx2.
    intros; apply addOrd_monotone; auto with ord.
    apply IHx2.
Qed.

Fixpoint VF_decompose (x:VForm) : list VForm :=
  match x with
  | Z => []
  | V a b => a :: VF_decompose b
  end.

Fixpoint VF_recompose (xs:list VForm) : VForm :=
  match xs with
  | [] => Z
  | x::xs' => V x (VF_recompose xs')
  end.

Lemma VF_decompose_correct:
  forall x,
    VF_normal x ->
    each VF_normal (VF_decompose x) /\
    cantor_ordered VF_denote (VF_decompose x) /\
    VF_denote x ≈ cantor_denote VF_denote (VF_decompose x).
Proof.
  induction x; simpl; intuition.
  destruct x2; simpl; auto.
  rewrite veblen_onePlus; auto.
  rewrite H8. reflexivity.
Qed.

Lemma VF_recompose_correct:
  forall xs,
    each VF_normal xs ->
    cantor_ordered VF_denote xs ->
    VF_normal (VF_recompose xs) /\ cantor_denote VF_denote xs ≈ VF_denote (VF_recompose xs).
Proof.
  induction xs; simpl in *; intuition.
  destruct xs; simpl; intuition.
  rewrite veblen_onePlus; auto with ord.
  rewrite H5.
  reflexivity.
Qed.

Definition VF_has_cantor_decomposition : has_cantor_decomposition VF_denote VF_normal :=
  Build_has_cantor_decomposition
    VF_denote
    VF_normal
    VF_compare
    VF_decompose
    VF_recompose
    VF_complete
    (fun x y Hx Hy => VF_compare_correct x y)
    VF_decompose_correct
    VF_recompose_correct.

Definition VF_succ := cantor_succ VF_has_cantor_decomposition.
Definition VF_add := cantor_add VF_has_cantor_decomposition.
Definition VF_mul := cantor_mul VF_has_cantor_decomposition.
Definition VF_exp := cantor_exp VF_has_cantor_decomposition.

Definition VF_one := V Z Z.
Definition VF_expOmega x := Vnorm x Z.

Theorem VF_reflects_zero : reflects VForm VF_denote VF_normal ORD 0 Z.
Proof.
  simpl; auto with ord.
Qed.

Theorem VF_reflects_one : reflects VForm VF_denote VF_normal ORD 1 VF_one.
Proof.
  simpl.
  rewrite veblen_onePlus; auto.
  rewrite addOrd_zero_r.
  rewrite expOrd_zero.
  intuition.
Qed.

Theorem VF_succ_reflects: reflects VForm VF_denote VF_normal (ORD ==> ORD) succOrd VF_succ.
Proof.
  apply cantor_succ_reflects.
Qed.

Theorem VF_add_reflects: reflects VForm VF_denote VF_normal (ORD ==> ORD ==> ORD) addOrd VF_add.
Proof.
  apply cantor_add_reflects.
Qed.

Theorem VF_mul_reflects: reflects VForm VF_denote VF_normal (ORD ==> ORD ==> ORD) mulOrd VF_mul.
Proof.
  apply cantor_mul_reflects.
Qed.

Theorem VF_exp_reflects: reflects VForm VF_denote VF_normal (ORD ==> ORD ==> ORD) expOrd VF_exp.
Proof.
  apply cantor_exp_reflects.
Qed.

Theorem VF_reflects_expOmega : reflects VForm VF_denote VF_normal (ORD ==> ORD) (expOrd ω) VF_expOmega.
Proof.
  simpl; intuition.
  rewrite veblen_onePlus; auto.
  rewrite addOrd_zero_r.
  rewrite H0.
  reflexivity.
Qed.

Theorem VF_ε₀ : VF ≈ ε 0.
Proof.
  split.
  - rewrite ord_le_unfold; intro x.
    apply VNotation_below_ε₀.
  - apply ε0_least_exp_closed with VF_normal Z VF_succ VF_expOmega.
    intro x. exists (Vnormalize x).
    split. symmetry. apply Vnormalize_equal.
    apply Vnormalize_normal.
    apply VF_reflects_zero.
    apply VF_succ_reflects.
    apply VF_reflects_expOmega.
Qed.

Require Import ClassicalFacts.
From Ordinal Require Import Classical.

Theorem VF_has_enough_notations (EM:excluded_middle) :
  forall x:Ord, x < ε 0 -> exists c:VF, x ≈ c.
Proof.
  induction x as [x Hx] using ordinal_induction. intro H.
  destruct (classical.ordinal_discriminate EM x) as [Hzero|[Hsucc|Hlimit]].
  - (* Zero ordinal, exhibit Z *)
    exists Z. simpl.
    apply ord_isZero in Hzero. auto.

  - (* Successor ordinal *)
    apply ord_isSucc in Hsucc.
    destruct Hsucc as [o Ho].

    (* invoke the induction hypothesis *)
    destruct (Hx o) as [vo Hvo].
    rewrite Ho. apply succ_lt.
    transitivity x; auto.
    rewrite Ho. apply succ_lt.

    (* exhibit the successor V form and wrap up *)
    exists (VF_succ (Vnormalize vo)).
    rewrite Ho.
    apply VF_succ_reflects; auto.
    simpl. rewrite Hvo.
    split. symmetry. apply Vnormalize_equal.
    apply Vnormalize_normal.

  - (* x is a limit, it must be a fixpoint of (addOrd 1) *)
    assert (Hlimit' : 1 + x <= x).
    { apply limit_onePlus; auto. }

    (* x cannot be an ε number, as it would be too large *)
    assert (Hepsilon: x < veblen (addOrd 1) x 0).
    { destruct (classical.order_total EM (veblen (addOrd 1) x 0) x); auto.
      elim (ord_lt_irreflexive (ε 0)).
      apply ord_le_lt_trans with x; auto.
      unfold ε. simpl.
      apply fixOrd_least; auto.
      apply normal_monotone. apply powOmega_normal.
      rewrite ord_le_unfold; intros [].
      rewrite veblen_onePlus in H0; auto.
      rewrite addOrd_zero_r in H0. auto.
      apply classical.ord_complete; auto.
    }

    (* decompose the ordinal *)
    destruct (veblen_decompose EM (addOrd 1) (onePlus_normal) x) as [a [b [Hab[_[Ha Hb]]]]]; auto.

    (* invoke the induction hypotheses *)
    destruct (Hx a) as [va Hva]; auto.
    transitivity x; auto.
    destruct (Hx b) as [vb Hvb]; auto.
    transitivity x; auto.

    (* exhibit the V form and wrap up *)
    exists (V va vb).
    rewrite Hab. simpl; symmetry.
    transitivity (veblen (addOrd 1) a (VF_denote vb)).
    split; apply veblen_monotone_first; auto.
    apply normal_monotone; auto.
    apply Hva.
    apply normal_monotone; auto.
    apply Hva.
    split; apply veblen_monotone; auto.
    apply normal_monotone; auto.
    apply Hvb.
    apply normal_monotone; auto.
    apply Hvb.
Qed.


Require Export Frap.Frap.
From Coq Require Export NArith.
From Coq Require Export Program.Equality.

Arguments N.add : simpl nomatch.
Arguments N.sub : simpl nomatch.
Arguments N.mul : simpl nomatch.
Arguments N.div : simpl nomatch.
Arguments N.shiftl : simpl nomatch.
Arguments N.shiftr : simpl nomatch.

(*|
Totals below don't sum up to 100%!  This is because this pset is a
chose-your-own-adventure assignment, so you can pick what to do.  Here is the
complete rubric, which choice points indicated by `==`::

    1    Axiom opt_binop_fold_test1
    8    Axiom opt_binop_fold_sound
         = Arith (23 points)
           == Precompute
   17         Axiom opt_binop_precompute_sound
    3         Axiom opt_arith_precompute_test1
    3         Axiom opt_arith_precompute_test2
           == Log
   17         Axiom opt_binop_log2_sound
    3         Axiom opt_arith_log2_test1
    3         Axiom opt_arith_log2_test2
           == Bitwise
   17         Axiom opt_binop_bitwise_sound
    3         Axiom opt_arith_bitwise_test1
    3         Axiom opt_arith_bitwise_test2
    4    Axiom opt_arith_fold_test1
   10    Axiom opt_arith_sound
    3    Axiom opt_unskip_test1
    3    Axiom opt_unskip_test2
   16    Axiom opt_unskip_sound
         = Eval (32 points)
           == ConstProp
    8         Axiom opt_arith_constprop_sound
    3         Axiom opt_constprop_test1
    3         Axiom opt_constprop_test2
   18         Axiom opt_constprop_sound
           == Unroll
    6         Parameter eval_read_only
    3         Axiom opt_unroll_test1
    3         Axiom opt_unroll_test2
   20         Axiom opt_unroll_sound
         -- Total 100
|*)

Module Type S.
  Inductive BinopName : Set :=
    LogAnd : BinopName
  | Eq : BinopName
  | ShiftLeft : BinopName
  | ShiftRight : BinopName
  | Times : BinopName
  | Divide : BinopName
  | Plus : BinopName
  | Minus : BinopName
  | Modulo : BinopName.

  Inductive expr : Set :=
    Const : nat -> expr
  | Var : var -> expr
  | Binop : BinopName -> expr -> expr -> expr.

  Inductive cmd : Set :=
    Skip : cmd
  | Assign : var -> expr -> cmd
  | AssignCall : var -> var -> expr -> expr -> cmd
  | Sequence : cmd -> cmd -> cmd
  | If : expr -> cmd -> cmd -> cmd
  | While : expr -> cmd -> cmd.

  Declare Scope expr.
  Delimit Scope expr with expr.

  Coercion Const : nat >-> expr.
  Coercion Var : var >-> expr.

  Infix "&" := (Binop LogAnd) (at level 80) : expr.
  Infix "==" := (Binop Eq) (at level 70) : expr.
  Infix ">>" := (Binop ShiftRight) (at level 60) : expr.
  Infix "<<" := (Binop ShiftLeft) (at level 60) : expr.
  Infix "+" := (Binop Plus) (at level 50, left associativity) : expr.
  Infix "-" := (Binop Minus) (at level 50, left associativity) : expr.
  Infix "*" := (Binop Times) (at level 40, left associativity) : expr.
  Infix "/" := (Binop Divide) (at level 40, left associativity) : expr.
  Infix "mod" := (Binop Modulo) (at level 40) : expr.

  Notation "x <- e" :=
    (Assign x e%expr)
      (at level 75).
  Notation "x <- 'call1' f e1" :=
    (AssignCall x f e1%expr (Const 0))
      (at level 75, f at level 0, e1 at level 0).
  Notation "x <- 'call2' f e1 e2" :=
    (AssignCall x f e1%expr e2%expr)
      (at level 75, f at level 0, e1 at level 0, e2 at level 0).
  Infix ";;" :=
    Sequence (at level 76).
  Notation "'when' e 'then' then_ 'else' else_ 'done'" :=
    (If e%expr then_ else_)
      (at level 75, e at level 0).
  Notation "'while' e 'loop' body 'done'" :=
    (While e%expr body)
      (at level 75).

  Example Times3Plus1Body :=
    ("n" <- "n" * 3 + 1).

  Example FactBody :=
    ("f" <- 1;;
     while "n" loop
       "f" <- "f" * "n";;
       "n" <- "n" - 1
     done).

  Example FactRecBody :=
    (when "n" == 1 then
       "f" <- 1
     else
       "f" <- call1 "fact_r" ("n" - 1);;
       "f" <- "f" * "n"
     done).

  Example FactTailRecBody :=
    (when "n" == 1 then
       "f" <- "acc"
     else
       "f" <- call2 "fact_tr" ("n" - 1) ("acc" * "n")
     done).

  Example CollatzBody :=
    (when ("start" == 1) then
       Skip
     else when ("start" mod 2 == 0) then
            "start" <- "start" / 2
          else
            (* `call1 f arg` is short for `call2 f arg 0` *)
            "start" <- call1 "times3plus1" ("start" + 0)
          done;;
          "flight" <- call2 "collatz" "start" ("flight" + 1)
     done).

  Notation TimesThreePlus1_signature := (("n", ""), "n", Times3Plus1Body).
  Notation Fact_signature := (("n", ""), "f", FactBody).
  Notation FactRec_signature := (("n", ""), "f", FactRecBody).
  Notation FactTailRec_signature := (("n", "acc"), "f", FactTailRecBody).
  Notation Collatz_signature := (("start", "flight"), "flight", CollatzBody).

  Notation Collatz_env :=
  ($0
    $+ ("collatz", Collatz_signature)
    $+ ("times3plus1", TimesThreePlus1_signature)).

  Notation Fact_env :=
    ($0
      $+ ("fact", Fact_signature)
      $+ ("fact_r", FactRec_signature)
      $+ ("fact_tr", FactTailRec_signature)).

  Parameter interp_binop : BinopName -> nat -> nat -> nat.

  Definition valuation := fmap var nat.
  Parameter interp_arith : expr -> valuation -> nat.

  Definition environment := fmap string ((var * var) * var * cmd).

  Inductive eval (phi: environment): valuation -> cmd -> valuation -> Prop :=
  | EvalSkip: forall v,
      eval phi v Skip v
  | EvalAssign: forall v x e a,
      interp_arith e v = a ->
      eval phi v (Assign x e) (v $+ (x, a))
  | EvalAssignCall: forall x f e1 e2 x1 x2 y body a1 a2 a v v',
      phi $? f = Some ((x1, x2), y, body) ->
      interp_arith e1 v = a1 ->
      interp_arith e2 v = a2 ->
      eval phi ($0 $+ (x1, a1) $+ (x2, a2)) body v' ->
      v' $? y = Some a ->
      eval phi v (AssignCall x f e1 e2) (v $+ (x, a))
  | EvalSequence: forall v c1 v1 c2 v2,
      eval phi v c1 v1 ->
      eval phi v1 c2 v2 ->
      eval phi v (Sequence c1 c2) v2
  | EvalIfTrue: forall v e thn els v' c,
      interp_arith e v = c ->
      c <> 0 ->
      eval phi v thn v' ->
      eval phi v (If e thn els) v'
  | EvalIfFalse: forall v e thn els v',
      interp_arith e v = 0 ->
      eval phi v els v' ->
      eval phi v (If e thn els) v'
  | EvalWhileTrue: forall v e body v' v'' c,
      interp_arith e v = c ->
      c <> 0 ->
      eval phi v body v' ->
      eval phi v' (While e body) v'' ->
      eval phi v (While e body) v''
  | EvalWhileFalse: forall v e body,
      interp_arith e v = 0 ->
      eval phi v (While e body) v.

  Definition eval_returns phi v cmd outVar result :=
    exists v', eval phi v cmd v' /\ v' $? outVar = Some result.

  Axiom TwoPlusTwoIsFour :
    eval_returns $0 $0 ("out" <- 2 + 2) "out" 4.
  Axiom EvalVars :
    eval_returns $0 $0 ("x" <- 1 + 1;; "x" <- "x" + "x" + "y") "x" 4.
  Axiom EvalSimpleArith :
    eval_returns $0 $0 ("out" <- (((14 >> 1) + 8 / 4 / 2) * (7 - 2) << 1) + 2 == 82) "out" 1.
  Axiom EvalTimes3Plus1 : forall n : nat,
      eval_returns $0 ($0 $+ ("n", n)) Times3Plus1Body "n" (3 * n + 1).
  Axiom EvalFact6 : exists v : valuation,
      eval $0 ($0 $+ ("n", 3)) FactBody v /\ v $? "f" = Some 6.
  Axiom EvalFactRec6 : exists v : valuation,
      eval Fact_env ($0 $+ ("n", 3)) FactRecBody v /\ v $? "f" = Some 6.
  Axiom EvalFactTailRec6 : exists v : valuation,
      eval Fact_env ($0 $+ ("n", 3) $+ ("acc", 1)) FactTailRecBody v /\
      v $? "f" = Some 6.
  Axiom collatz_result : exists v : valuation,
      eval Collatz_env ($0 $+ ("flight", 0) $+ ("start", 10)) CollatzBody v /\
      v $? "flight" = Some 6.

  Parameter opt_binop_fold : BinopName -> expr -> expr -> expr.
  (*[1%]*) Axiom opt_binop_fold_test1 : opt_binop_fold Plus "x" 0 = "x".

  (*[8%]*) Axiom opt_binop_fold_sound :
    forall (b : BinopName) (e1 e2 : expr) (v : valuation),
      interp_arith (opt_binop_fold b e1 e2) v =
      interp_binop b (interp_arith e1 v) (interp_arith e2 v).

  Parameter opt_binop_precompute : BinopName -> expr -> expr -> expr.
  (*[Arith/Precompute:17%]*) Axiom opt_binop_precompute_sound :
    forall (b : BinopName) (e1 e2 : expr) (v : valuation),
      interp_arith (opt_binop_precompute b e1 e2) v =
      interp_binop b (interp_arith e1 v) (interp_arith e2 v).

  Parameter opt_binop_log2 : BinopName -> expr -> expr -> expr.
  (*[Arith/Log:17%]*) Axiom opt_binop_log2_sound :
    forall (b : BinopName) (e1 e2 : expr) (v : valuation),
      interp_arith (opt_binop_log2 b e1 e2) v =
      interp_binop b (interp_arith e1 v) (interp_arith e2 v).

  Parameter opt_binop_bitwise : BinopName -> expr -> expr -> expr.
  (*[Arith/Bitwise:17%]*) Axiom opt_binop_bitwise_sound :
    forall (b : BinopName) (e1 e2 : expr) (v : valuation),
      interp_arith (opt_binop_bitwise b e1 e2) v =
      interp_binop b (interp_arith e1 v) (interp_arith e2 v).

  Parameter opt_arith : expr -> expr.

  (*[4%]*) Axiom opt_arith_fold_test1 :
    opt_arith (1 + "z" * ("y" * ("x" * (0 + 0 / 1))))%expr =
    (1)%expr.
  (*[Arith/Precompute:3%]*) Axiom opt_arith_precompute_test1:
    opt_arith (("x" + (3 - 3)) / (0 + 1) + ("y" + "y" * 0))%expr =
    ("x" + "y")%expr.
  (*[Arith/Precompute:3%]*) Axiom opt_arith_precompute_test2 :
    opt_arith ((("y" / ("x" * 0 + 7 / 1)) mod (12 - 5)) / (2 * 3))%expr =
    (("y" / 7) mod 7 / 6)%expr.
  (*[Arith/Log:3%]*) Axiom opt_arith_log2_test1 :
    opt_arith (("y" * 8) mod 8 / 4)%expr =
    (("y" << 3 & 7) >> 2)%expr.
  (*[Arith/Log:3%]*) Axiom opt_arith_log2_test2 :
    opt_arith (("y" * 1 + (4 + 0)) mod 9 / 3)%expr =
    (("y" + 4) mod 9 / 3)%expr.
  (*[Arith/Bitwise:3%]*) Axiom opt_arith_bitwise_test1 :
    opt_arith ("y" * 13)%expr =
    ("y" + (("y" << 2) + ("y" << 3)))%expr.
  (*[Arith/Bitwise:3%]*) Axiom opt_arith_bitwise_test2 :
    opt_arith ("y" * (3 + 0))%expr =
    ("y" + ("y" << 1))%expr.

  (*[10%]*) Axiom opt_arith_sound :
    forall (e : expr) (v : valuation),
      interp_arith (opt_arith e) v = interp_arith e v.

  Parameter opt_unskip : cmd -> cmd.

  (*[3%]*) Axiom opt_unskip_test1 :
    opt_unskip (Skip;; (Skip;; Skip);; (Skip;; Skip;; Skip)) =
    Skip.
  (*[3%]*) Axiom opt_unskip_test2 :
    opt_unskip (when 0 then (Skip;; Skip) else Skip done;;
                while 0 loop Skip;; Skip done;; Skip) =
    (when 0 then Skip else Skip done;; while 0 loop Skip done).

  (*[16%]*) Axiom opt_unskip_sound :
    forall (phi : environment) (c : cmd) (v v' : valuation),
      eval phi v c v' -> eval phi v (opt_unskip c) v'.

  Parameter opt_arith_constprop : expr -> valuation -> expr.

  (*[Eval/ConstProp:8%]*) Axiom opt_arith_constprop_sound :
    forall (e : expr) (v consts : fmap var nat),
      consts $<= v ->
      interp_arith (opt_arith_constprop e consts) v = interp_arith e v.

  Parameter opt_constprop : cmd -> cmd.

  (*[Eval/ConstProp:3%]*) Axiom opt_constprop_test1 :
    opt_constprop FactBody = FactBody.
  (*[Eval/ConstProp:3%]*) Axiom opt_constprop_test2 :
    opt_constprop ("x" <- 7;; "y" <- "x";;
                   when "x" mod "w" then
                     "z" <- "x";; "t" <- "z";; while "t" loop "t" <- "t" - 1 done
                   else
                     "z" <- "u" + "x";; "t" <- "z"
                   done;;
                   "r" <- "z") =
   ("x" <- 7;; "y" <- 7;;
    when 7 mod "w" then
      "z" <- 7;; "t" <- 7;; while "t" loop "t" <- "t" - 1 done
    else
      "z" <- "u" + 7;; "t" <- "z"
    done;;
    "r" <- "z").

  (*[Eval/ConstProp:18%]*) Axiom opt_constprop_sound :
    forall (phi : environment) (c : cmd) (v v' : valuation),
      eval phi v c v' -> eval phi v (opt_constprop c) v'.

  Parameter read_only : forall (c: cmd) (x0: var), bool.

  (*[Eval/Unroll:6%]*)
  Parameter eval_read_only: forall {phi v v' x c},
      eval phi v c v' ->
      read_only c x = true ->
      v' $? x = v $? x.

  Parameter opt_unroll : cmd -> cmd.

  (*[Eval/Unroll:3%]*) Axiom opt_unroll_test1 :
    opt_unroll CollatzBody = CollatzBody.
  (*[Eval/Unroll:3%]*) Axiom opt_unroll_test2 :
    opt_unroll FactBody <> FactBody.

  (*[Eval/Unroll:20%]*) Axiom opt_unroll_sound :
    forall (phi : environment) (c : cmd) (v v' : valuation),
      eval phi v c v' -> eval phi v (opt_unroll c) v'.
End S.

Global Arguments Nat.modulo !_ !_ /.
Global Arguments Nat.div !_ !_.
Global Arguments Nat.log2 !_ /.


Require Import FunctionalExtensionality.

Set Implicit Arguments.

Definition EqDec (T : Type) := forall (a b : T), {a = b} + {a <> b}.

(* generalized memory of any address / value type *)

Definition mem {A : Type} {AEQ : EqDec A} {V : Type} := A -> option V.
Definition empty_mem {A : Type} {AEQ : EqDec A} {V : Type} : @mem A AEQ V := fun a => None.

Section GenMem.
  Variable A : Type.
  Variable V : Type.
  Variable AEQ : EqDec A.

  Definition upd (m : @mem A AEQ V) (a : A) (v : V) : @mem A AEQ V :=
    fun a' => if AEQ a' a then Some v else m a'.

  Definition updSome (m : @mem A AEQ V) (a : A) (v : V) : @mem A AEQ V :=
    fun a' => if AEQ a' a then
      match m a with
      | None => None
      | Some _ => Some v
      end else m a'.

  Definition delete (m : @mem A AEQ V) (a : A) : @mem A AEQ V :=
    fun a' => if AEQ a' a then None else m a'.

  Definition insert (m : @mem A AEQ V) (a : A) (v : V) : @mem A AEQ V :=
    fun a' => if AEQ a' a then
      match m a with
      | None => Some v
      | Some _ => m a'
      end else m a'.

  Theorem upd_eq : forall m (a : A) (v : V) a',
    a' = a -> upd m a v a' = Some v.
  Proof.
    intros; subst; unfold upd.
    destruct (AEQ a a); tauto.
  Qed.

  Theorem updSome_eq : forall m (a : A) (v v' : V) a',
    m a = Some v' -> a' = a -> updSome m a v a' = Some v.
  Proof.
    intros; subst; unfold updSome.
    rewrite H.
    destruct (AEQ a a); tauto.
  Qed.

  Theorem insert_eq : forall m (a : A) (v v' : V) a',
    m a = None -> a' = a -> insert m a v a' = Some v.
  Proof.
    intros; subst; unfold insert.
    rewrite H.
    destruct (AEQ a a); congruence.
  Qed.

  Theorem upd_ne : forall m (a : A) (v : V) a',
    a' <> a -> upd m a v a' = m a'.
  Proof.
    intros; subst; unfold upd.
    destruct (AEQ a' a); tauto.
  Qed.

  Theorem updSome_ne : forall m (a : A) (v : V) a',
    a' <> a -> updSome m a v a' = m a'.
  Proof.
    intros; subst; unfold updSome.
    destruct (AEQ a' a); tauto.
  Qed.

  Theorem insert_ne : forall m (a : A) (v : V) a',
    a' <> a -> insert m a v a' = m a'.
  Proof.
    intros; subst; unfold insert.
    destruct (AEQ a' a); congruence.
  Qed.

  Theorem upd_repeat: forall m (a : A) (v v':V),
    upd (upd m a v') a v = upd m a v.
  Proof.
    intros; apply functional_extensionality; unfold upd; intros.
    destruct (AEQ a a); try congruence.
    destruct (AEQ x a); auto.
  Qed.

  Theorem updSome_repeat: forall m (a : A) (v v':V),
    updSome (updSome m a v') a v = updSome m a v.
  Proof.
    intros; apply functional_extensionality; unfold updSome; intros.
    destruct (AEQ a a); try congruence.
    destruct (AEQ x a); auto.
    destruct (m a); auto.
  Qed.

  Theorem insert_repeat: forall m (a : A) (v v':V),
    insert (insert m a v) a v' = insert m a v.
  Proof.
    intros; apply functional_extensionality; unfold insert at 1; intros.
    destruct (AEQ a a); try congruence.
    destruct (AEQ x a); auto.
    subst. unfold insert; simpl.
    destruct (AEQ a a); try congruence.
    destruct (m a); auto.
  Qed.

  Theorem upd_nop: forall m (a : A) (v : V),
    m a = Some v ->
    upd m a v = m.
  Proof.
    intros; apply functional_extensionality; intros.
    case_eq (AEQ a x); intros; subst.
    repeat erewrite upd_eq; eauto.
    repeat rewrite upd_ne; auto.
  Qed.

  Theorem updSome_nop: forall m (a : A) (v : V),
    m a = Some v ->
    updSome m a v = m.
  Proof.
    intros; apply functional_extensionality; intros.
    case_eq (AEQ a x); intros; subst.
    repeat erewrite updSome_eq; eauto.
    repeat rewrite updSome_ne; auto.
  Qed.

  Theorem updSome_none : forall m (a : A) (v : V),
    m a = None ->
    updSome m a v = m.
  Proof.
    unfold updSome; intros; apply functional_extensionality; intros.
    rewrite H.
    destruct (AEQ x a); subst; auto.
  Qed.

  Theorem upd_comm: forall m (a0 : A) (v0 : V) a1 v1, a0 <> a1
    -> upd (upd m a0 v0) a1 v1 = upd (upd m a1 v1) a0 v0.
  Proof.
    intros; apply functional_extensionality; unfold upd; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

  Theorem updSome_comm: forall m (a0 : A) (v0 : V) a1 v1, a0 <> a1
    -> updSome (updSome m a0 v0) a1 v1 = updSome (updSome m a1 v1) a0 v0.
  Proof.
    intros; apply functional_extensionality; unfold updSome; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

  Theorem updSome_insert_comm: forall m (a0 : A) (v0 : V) a1 v1, a0 <> a1
    -> updSome (insert m a0 v0) a1 v1 = insert (updSome m a1 v1) a0 v0.
  Proof.
    intros; apply functional_extensionality; unfold updSome, insert; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

  Theorem updSome_delete_comm: forall m (a0 : A) a1 (v1 : V), a0 <> a1
    -> updSome (delete m a0) a1 v1 = delete (updSome m a1 v1) a0.
  Proof.
    intros; apply functional_extensionality; unfold updSome, delete; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

  Theorem insert_comm: forall m (a0 : A) (v0 : V) a1 v1, a0 <> a1
    -> insert (insert m a0 v0) a1 v1 = insert (insert m a1 v1) a0 v0.
  Proof.
    intros; apply functional_extensionality; unfold insert; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

  Theorem insert_delete_comm: forall m (a0 : A) a1 (v1 : V), a0 <> a1
    -> insert (delete m a0) a1 v1 = delete (insert m a1 v1) a0.
  Proof.
    intros; apply functional_extensionality; unfold insert, delete; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

  Theorem delete_comm: forall m (a0 : A) a1, a0 <> a1
    -> delete (delete m a0) a1 = delete (delete m a1) a0.
  Proof.
    intros; apply functional_extensionality; unfold delete; intros.
    destruct (AEQ a1 a0); destruct (AEQ a0 a1); try congruence.
    case_eq (AEQ x a1); case_eq (AEQ x a0); intros; subst; try congruence.
  Qed.

End GenMem.

Hint Rewrite upd_eq using (solve [ auto ]) : upd.
Hint Rewrite upd_ne using (solve [ auto ]) : upd.


From Hammer Require Import Hammer.

Require Export Coq.Classes.Init.
Require Import Coq.Program.Basics.
Require Import Coq.Program.Tactics.

Generalizable Variables A B C D R S T U l eqA eqB eqC eqD.

Set Universe Polymorphism.

Definition crelation (A : Type) := A -> A -> Type.

Definition arrow (A B : Type) := A -> B.

Definition flip {A B C : Type} (f : A -> B -> C) := fun x y => f y x.

Definition iffT (A B : Type) := ((A -> B) * (B -> A))%type.

Section Defs.
Context {A : Type}.

Class Reflexive (R : crelation A) :=
reflexivity : forall x : A, R x x.

Definition complement (R : crelation A) : crelation A :=
fun x y => R x y -> False.

Typeclasses Opaque complement iffT.

Lemma complement_inverse R : complement (flip R) = flip (complement R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.complement_inverse".   reflexivity. Qed.

Class Irreflexive (R : crelation A) :=
irreflexivity : Reflexive (complement R).

Class Symmetric (R : crelation A) :=
symmetry : forall {x y}, R x y -> R y x.

Class Asymmetric (R : crelation A) :=
asymmetry : forall {x y}, R x y -> (complement R y x : Type).

Class Transitive (R : crelation A) :=
transitivity : forall {x y z}, R x y -> R y z -> R x z.

Class PreOrder (R : crelation A)  := {
PreOrder_Reflexive :> Reflexive R | 2 ;
PreOrder_Transitive :> Transitive R | 2 }.

Class StrictOrder (R : crelation A)  := {
StrictOrder_Irreflexive :> Irreflexive R ;
StrictOrder_Transitive :> Transitive R }.

Global Instance StrictOrder_Asymmetric `(StrictOrder R) : Asymmetric R.
Proof. hammer_hook "CRelationClasses" "CRelationClasses.StrictOrder_Asymmetric".   firstorder. Qed.

Class PER (R : crelation A)  := {
PER_Symmetric :> Symmetric R | 3 ;
PER_Transitive :> Transitive R | 3 }.

Class Equivalence (R : crelation A)  := {
Equivalence_Reflexive :> Reflexive R ;
Equivalence_Symmetric :> Symmetric R ;
Equivalence_Transitive :> Transitive R }.

Global Instance Equivalence_PER {R} `(Equivalence R) : PER R | 10 :=
{ PER_Symmetric := Equivalence_Symmetric ;
PER_Transitive := Equivalence_Transitive }.

Class Antisymmetric eqA `{equ : Equivalence eqA} (R : crelation A) :=
antisymmetry : forall {x y}, R x y -> R y x -> eqA x y.

Class subrelation (R R' : crelation A) :=
is_subrelation : forall {x y}, R x y -> R' x y.

Lemma subrelation_symmetric R `(Symmetric R) : subrelation (flip R) R.
Proof. hammer_hook "CRelationClasses" "CRelationClasses.subrelation_symmetric".   hnf. intros x y H'. red in H'. apply symmetry. assumption. Qed.

Section flip.

Lemma flip_Reflexive `{Reflexive R} : Reflexive (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.flip_Reflexive".   tauto. Qed.

Program Definition flip_Irreflexive `(Irreflexive R) : Irreflexive (flip R) :=
irreflexivity (R:=R).

Program Definition flip_Symmetric `(Symmetric R) : Symmetric (flip R) :=
fun x y H => symmetry (R:=R) H.

Program Definition flip_Asymmetric `(Asymmetric R) : Asymmetric (flip R) :=
fun x y H H' => asymmetry (R:=R) H H'.

Program Definition flip_Transitive `(Transitive R) : Transitive (flip R) :=
fun x y z H H' => transitivity (R:=R) H' H.

Program Definition flip_Antisymmetric `(Antisymmetric eqA R) :
Antisymmetric eqA (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.flip_Antisymmetric".   firstorder. Qed.

Lemma flip_PreOrder `(PreOrder R) : PreOrder (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.flip_PreOrder".   firstorder. Qed.

Lemma flip_StrictOrder `(StrictOrder R) : StrictOrder (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.flip_StrictOrder".   firstorder. Qed.

Lemma flip_PER `(PER R) : PER (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.flip_PER".   firstorder. Qed.

Lemma flip_Equivalence `(Equivalence R) : Equivalence (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.flip_Equivalence".   firstorder. Qed.

End flip.

Section complement.

Definition complement_Irreflexive `(Reflexive R)
: Irreflexive (complement R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.complement_Irreflexive".   firstorder. Qed.

Definition complement_Symmetric `(Symmetric R) : Symmetric (complement R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.complement_Symmetric".   firstorder. Qed.
End complement.

Class RewriteRelation (RA : crelation A).

Global Instance equivalence_rewrite_crelation `(Equivalence eqA) : RewriteRelation eqA.

Section Leibniz.
Global Instance eq_Reflexive : Reflexive (@eq A) := @eq_refl A.
Global Instance eq_Symmetric : Symmetric (@eq A) := @eq_sym A.
Global Instance eq_Transitive : Transitive (@eq A) := @eq_trans A.

Global Program Instance eq_equivalence : Equivalence (@eq A) | 10.
End Leibniz.

End Defs.

Instance: RewriteRelation impl.
Instance: RewriteRelation iff.

Hint Extern 1 (Reflexive (complement _)) => class_apply @irreflexivity : typeclass_instances.
Hint Extern 3 (Symmetric (complement _)) => class_apply complement_Symmetric : typeclass_instances.
Hint Extern 3 (Irreflexive (complement _)) => class_apply complement_Irreflexive : typeclass_instances.

Hint Extern 3 (Reflexive (flip _)) => apply flip_Reflexive : typeclass_instances.
Hint Extern 3 (Irreflexive (flip _)) => class_apply flip_Irreflexive : typeclass_instances.
Hint Extern 3 (Symmetric (flip _)) => class_apply flip_Symmetric : typeclass_instances.
Hint Extern 3 (Asymmetric (flip _)) => class_apply flip_Asymmetric : typeclass_instances.
Hint Extern 3 (Antisymmetric (flip _)) => class_apply flip_Antisymmetric : typeclass_instances.
Hint Extern 3 (Transitive (flip _)) => class_apply flip_Transitive : typeclass_instances.
Hint Extern 3 (StrictOrder (flip _)) => class_apply flip_StrictOrder : typeclass_instances.
Hint Extern 3 (PreOrder (flip _)) => class_apply flip_PreOrder : typeclass_instances.

Hint Extern 4 (subrelation (flip _) _) =>
class_apply @subrelation_symmetric : typeclass_instances.

Hint Resolve irreflexivity : ord.

Unset Implicit Arguments.

Ltac solve_crelation :=
match goal with
| [ |- ?R ?x ?x ] => reflexivity
| [ H : ?R ?x ?y |- ?R ?y ?x ] => symmetry ; exact H
end.

Hint Extern 4 => solve_crelation : crelations.

Ltac reduce_hyp H :=
match type of H with
| context [ _ <-> _ ] => fail 1
| _ => red in H ; try reduce_hyp H
end.

Ltac reduce_goal :=
match goal with
| [ |- _ <-> _ ] => fail 1
| _ => red ; intros ; try reduce_goal
end.

Tactic Notation "reduce" "in" hyp(Hid) := reduce_hyp Hid.

Ltac reduce := reduce_goal.

Tactic Notation "apply" "*" constr(t) :=
first [ refine t | refine (t _) | refine (t _ _) | refine (t _ _ _) | refine (t _ _ _ _) |
refine (t _ _ _ _ _) | refine (t _ _ _ _ _ _) | refine (t _ _ _ _ _ _ _) ].

Ltac simpl_crelation :=
unfold flip, impl, arrow ; try reduce ; program_simpl ;
try ( solve [ dintuition ]).

Local Obligation Tactic := simpl_crelation.

Program Instance impl_Reflexive : Reflexive impl.
Program Instance impl_Transitive : Transitive impl.

Instance iff_Reflexive : Reflexive iff := iff_refl.
Instance iff_Symmetric : Symmetric iff := iff_sym.
Instance iff_Transitive : Transitive iff := iff_trans.

Program Instance iff_equivalence : Equivalence iff.
Program Instance arrow_Reflexive : Reflexive arrow.
Program Instance arrow_Transitive : Transitive arrow.

Instance iffT_Reflexive : Reflexive iffT.
Proof. hammer_hook "CRelationClasses" "CRelationClasses.iffT_Reflexive".   firstorder. Defined.
Instance iffT_Symmetric : Symmetric iffT.
Proof. hammer_hook "CRelationClasses" "CRelationClasses.iffT_Symmetric".   firstorder. Defined.
Instance iffT_Transitive : Transitive iffT.
Proof. hammer_hook "CRelationClasses" "CRelationClasses.iffT_Transitive".   firstorder. Defined.

Local Open Scope list_scope.

Section Binary.
Context {A : Type}.

Definition relation_equivalence : crelation (crelation A) :=
fun R R' => forall x y, iffT (R x y) (R' x y).

Global Instance: RewriteRelation relation_equivalence.

Definition relation_conjunction (R : crelation A) (R' : crelation A) : crelation A :=
fun x y => prod (R x y) (R' x y).

Definition relation_disjunction (R : crelation A) (R' : crelation A) : crelation A :=
fun x y => sum (R x y) (R' x y).

Global Instance relation_equivalence_equivalence :
Equivalence relation_equivalence.
Proof. hammer_hook "CRelationClasses" "CRelationClasses.relation_equivalence_equivalence".   split; red; unfold relation_equivalence, iffT. firstorder.
firstorder.
intros. specialize (X x0 y0). specialize (X0 x0 y0). firstorder.
Qed.

Global Instance relation_implication_preorder : PreOrder (@subrelation A).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.relation_implication_preorder".   firstorder. Qed.

Class PartialOrder eqA `{equ : Equivalence A eqA} R `{preo : PreOrder A R} :=
partial_order_equivalence : relation_equivalence eqA (relation_conjunction R (flip R)).

Global Instance partial_order_antisym `(PartialOrder eqA R) : ! Antisymmetric A eqA R.
Proof with auto. hammer_hook "CRelationClasses" "CRelationClasses.partial_order_antisym".
reduce_goal.
apply H. firstorder.
Qed.

Lemma PartialOrder_inverse `(PartialOrder eqA R) : PartialOrder eqA (flip R).
Proof. hammer_hook "CRelationClasses" "CRelationClasses.PartialOrder_inverse".   unfold flip; constructor; unfold flip. intros. apply H. apply symmetry. apply X.
unfold relation_conjunction. intros [H1 H2]. apply H. constructor; assumption. Qed.
End Binary.

Hint Extern 3 (PartialOrder (flip _)) => class_apply PartialOrder_inverse : typeclass_instances.

Typeclasses Opaque relation_equivalence.


Require Import ExtLib.Structures.Applicative.

Definition lift1 {T U : Type} {F : Type -> Type}
           {Ap : Applicative.Applicative F}
(f : T -> U) (x : F T) : F U :=
  Applicative.ap (Applicative.pure f) x.

Definition lift2 {T U V : Type} {F : Type -> Type}
           {Ap : Applicative.Applicative F}
(f : T -> U -> V) (x : F T) (y : F U) : F V :=
  Applicative.ap (lift1 (F:=F) f x) y.

Definition lift3 {T U V W : Type} {F : Type -> Type}
           {Ap : Applicative.Applicative F}
(f : T -> U -> V -> W) (x : F T) (y : F U) (z : F V) : F W :=
  Applicative.ap (lift2 (F:=F) f x y) z.


Require Import HoTT.
Require Import HottCat.Pushforwards.
Require Import FunextAxiom.

Generalizable Variables T F A B C.

Reserved Notation "T $$ a" (at level 60, right associativity).
Reserved Notation "f $* x" (at level 60, right associativity).
Reserved Notation "f $ x" (at level 60, right associativity, only parsing).

Local Notation Endofunctor := (Type -> Type).
Notation "f $ x" := (f(x)) (at level 60, right associativity, only parsing).

Delimit Scope functor_scope with functor.
Open Scope functor_scope.

(*
Class TCategory := {
  Ob: Type;
  Hom: relation Ob;
  id: forall x: Ob, Hom x x;
  comp: forall x y z: Ob, Hom y z -> Hom x y -> Hom x z
  }.

Infix "-->" := Hom (at level 90, right associativity).
Coercion Ob: TCategory >-> Sortclass.

Global Instance Grpd: TCategory := {
  Ob := Type;
  Hom := fun x y: Type => x -> y;
  id := fun x: Type => idmap;
  comp := fun (x y z: Type) (f: y -> z) (g: x -> y) => f o g
  }.
*)

Class TFunctor (T: Endofunctor) := {
  fmap: forall {A B: Type}, (A -> B) -> (T A -> T B);
  pres_id: forall {A: Type}, @fmap A A idmap = idmap;
  pres_compose: forall {A B C: Type} (f: B -> C) (g: A -> B),
                fmap (f o g) = (fmap f) o (fmap g)
  }.

Arguments fmap T {_} A B _ _.
Arguments pres_id T {_} A .
Arguments pres_compose T {_ A B C} f g.
Notation "T $$ a" := (fmap T _ _ a) (at level 60, right associativity): functor_scope.
Bind Scope functor_scope with TFunctor.

Section Functors.
  Context `{TFunctor T}.

  Section Composition.
    (** A composition of functors is again a functor. *)

    Context `{TFunctor F}.
    Notation G := (F o T).

    Local Definition G_fmap: forall {A B: Type}, (A -> B) -> (G A -> G B)
      := fun A B f => F $$ T $$ f.

    Local Definition G_id: forall A: Type, (@G_fmap A A idmap) = idmap.
      intro; unfold G_fmap.
      rewrite (pres_id T).
      exact (pres_id F (T A)).
      Defined.

    Local Definition G_comp: forall (A B C: Type) (f: B -> C) (g: A -> B),
                G_fmap (f o g) = (G_fmap f) o (G_fmap g) .
      intros; unfold G_fmap.
      rewrite (pres_compose T).
      exact (pres_compose F (T $$ f) (T $$ g)).
      Defined.

    Global Instance tfunctor_compose `{TFunctor F}: TFunctor (F o T)
      := {| fmap := @G_fmap; pres_id := G_id; pres_compose := G_comp |}.

  End Composition.

End Functors.

Class TApplicative (T: Endofunctor) := {
  pure: forall {A: Type}, A -> T A;
  fzip: forall {A B: Type}, T (A -> B) -> T A -> T B;
  pure_id: forall {A: Type}, @fzip A A (pure idmap) = idmap;
  (** [f] $* [x] = [f x] *)
  homomorphism: forall {A B: Type} (f: A -> B) (x: A),
                  fzip (pure f) (pure x) = pure (f x);
  (** u $* [x] = [fun g => g x] $* u *)
  interchange: forall {A B: Type} (u: T (A -> B)) (x: A),
               fzip u (pure x) = fzip (pure (fun g => g x)) u;
  (** f $* (g $* x) = (([o] $* f) $* g) $* x *)
  composition: forall {A B C: Type} (f: T(B -> C)) (g: T(A -> B)) (x: T A),
    fzip f (fzip g x) = fzip (fzip (fzip (pure (fun u v => u o v)) f) g) x
  }.

Arguments pure T {_ A} _.
Arguments fzip T {_} A B _ _.

Notation "[ x ]" := (pure _ x): functor_scope.
Notation "f $* a" := (fzip _ _ _  f a) (at level 60, right associativity): functor_scope.
Bind Scope functor_scope with TApplicative.

Section Applicative.

  Context `{TApplicative T}.

  Section Applicative_is_functor.

    Let T_fmap := fun (A B: Type) (f: A -> B) => fzip T _ _ [f].
    Let T_pres_id := fun A: Type => @pure_id _ _ A.
    Let T_pres_compose: forall (A B C : Type) (f : B -> C) (g : A -> B),
              T_fmap A C (fun x : A => f (g x)) =
              (fun x : T A => T_fmap B C f (T_fmap A B g x)).
      intros; unfold T_fmap.
      by_extensionality x.
      transitivity ((([compose] $* [f]) $* [g]) $* x).
      - rewrite (homomorphism compose f).
        rewrite (homomorphism _ g).
        trivial.
      - exact (composition [f] [g] x)^.
      Defined.

    Global Instance applicative_is_functor: TFunctor T
      := Build_TFunctor T T_fmap T_pres_id T_pres_compose.

  End Applicative_is_functor.

  Lemma pure_natural `(f: A -> B) (x: A): (T $$ f) [x] = [f x].
    unfold "$$", applicative_is_functor.
    destruct H; eauto.
    Defined.

End Applicative.

Class TMonad (T: Endofunctor) := {
  bind: forall {A B: Type}, (T A) -> (A -> T B) -> T B;
  ret: forall {A: Type}, A -> T A;
  ret_unit_left: forall {A B: Type} (k: A -> T B) (a: A),
                 bind (ret a) k  =  k a;
  ret_unit_right: forall {A: Type} (m: T A), bind m ret = m;
  bind_assoc: forall {A B C: Type} (m: T A) (k: A -> T B) (h: B -> T C),
              bind m (fun x => bind (k x) h) = bind (bind m k) h
  }.

Notation "x >>= f" := (bind x f) (at level 60).
Notation "x <- y ; f" := (y >>= (fun x => f))
    (at level 100, f at level 200,
    format "'[v' x  '<-'  y ';' '/' f ']'").

Section Monads.

  Context `{TMonad T}.
  (*
  Section Monad_is_functor.

    Let T_fmap
    Global Instance monad_is_functor: TFunctor T :=
      fun `(f: A -> B) (t: T A) => x <- t; ret (f x).

  End Monad_is_functor.
  *)

  Section Monad_is_applicative.

    Let T_fzip := fun (A B: Type) (f: T (A -> B)) (t: T A)
                  => a <- t; g <- f; ret (g a).
    Let T_pure := @ret _ _.

    Notation "[ x ]" := (T_pure _ x): functor_scope.
    Notation "f $* a" := (T_fzip _ _  f a)
                         (at level 60, right associativity): functor_scope.

    Local Definition T_pure_id : `(T_fzip A A (T_pure (A -> A) idmap) = idmap).
      intro; unfold T_fzip, T_pure.
      by_extensionality t.
      assert (p := fun a: A =>
                   (ret_unit_left (fun g => ret (g a)) idmap)^).
      apply path_forall in p.
      apply (ap (fun x => (t >>= x) = t)) in p.
      path_induction.
      exact (ret_unit_right t).
      Defined.

    Local Definition T_homomorphism : forall (A B : Type) (f : A -> B)
                      (x : A), ([f] $* [x]) = [f x].
      intros. unfold T_fzip, T_pure.
      assert (p := ret_unit_left (fun a => g <- ret f; ret (g a)) x).
      assert (q := ret_unit_left (fun g => ret (g x)) f).
      path_induction; done.
      Defined.

    Local Definition T_interchange :
            forall (A B : Type) (u : T (A -> B)) (x : A),
            u $* [x] = [fun g => g x] $* u.
      intros; unfold T_fzip, T_pure.
      assert (p := (ret_unit_left (fun a => g <- u; ret (g a)) x)^).
      assert (q := fun a: A -> B =>
              (ret_unit_left (fun g => ret (g a)) (fun g => g x))^).
      apply path_forall in q.
      path_induction.
      done.
      Defined.

    Local Definition T_composition :
          forall (A B C: Type) (f: T(B -> C)) (g: T(A -> B)) (x: T A),
          f $* (g $* x) = (([compose] $* f) $* g) $* x.
      intros; unfold T_pure, T_fzip.
      Admitted.

    Global Instance monad_is_applicative: TApplicative T
    := {| pure := T_pure; fzip := T_fzip; pure_id := T_pure_id;
        homomorphism := T_homomorphism; interchange := T_interchange;
        composition := T_composition |}.

  End Monad_is_applicative.

  Definition join {A: Type} (t: T (T A)): T A
    := x <- t; x.

  Lemma join_is_action (A: Type):
        join o (T $$ join) = join o (join (A:=T A)).
    unfold join, fmap. simpl.
    apply path_forall. unfold "_ == _". intro.
    f_ap.
    assert (p := fun a => inverse $ ret_unit_left (fun g => ret (g a))
          (fun t : T (T A) => x0 <- t; x0)).
    apply path_forall in p.
    induction p.
    Admitted.

  Lemma ret_is_unit (A: Type) (x: T A): join ((T $$ ret) x) = x.
  Admitted.

End Monads.

Section Algebras.

  (** We need to define a T-algebra on Type (regarding T as a functor, not as a monad). The most important case is `T = List`, `T_act ([X_1, .. , X_n]) := X_1 * .. * X_n`, but there are at least a few other variants (corresponding to commutative product on Type or the structure of commutative ring wrt +,* ). *)
  Context `{TMonad T}.
  Variable T_act: T Type -> Type.

  (** We regard the type family is_op as a predicate stating that a function `f: T A -> A` is a "virtual operation", i.e. it is an iterated composition of some deining set of primitive operations or their deformations. E.g. if A is a strict semigroup and T = List, then `m: T A -> A` is a tuple (m_0, m_1, .. , m_k, ..) of k-ary products `m_k: A^k -> A` and a virtual operation is a pair (h: T A -> A, p: h = m). Note that in this case all iterated compositions of m are trivially equal to m and thus we can define a composition of virtual operations, extending the usual one. In general there will be nontrivial paths between iterated compositions of m_i, thus we need to choose composition on `sig is_op` beforehand. *)
(*  Definition VirtualOp {A: Type} (is_op: (T A -> A) -> Type) :=
    exists (f: T A -> A), is_op f.*)

  Section Composition.

    Context {A: Type}.
    Variable is_op : (T A -> A) -> Type.

    (** Morally a point in this type corresponds to T h, where h is some virtual operation. In the case T = List such points are lists [f_1, .. , f_k] of operations with arities [n_1, .. , n_k]. *)
    Definition T_is_op: (T (T A) -> T A) -> Type :=
      let t_fam := T_act o (T $$ is_op)
      in push_sigma t_fam (fzip T _ _).

    Definition VirtualCompose :=
      forall (f: sig T_is_op) (g: sig is_op),
      { h: sig is_op | h.1 o join = g.1 o f.1 }.

    Definition VirtualUnit :=
      forall (f: sig is_op), f.1 o ret = idmap.

  End Composition.

  Record TAlgebra (A: Type):=
  {
    is_op: (T A -> A) -> Type;
    comp: VirtualCompose is_op;
    unit: VirtualUnit is_op;
    unique_op: Contr (sig is_op)
    }.

  Arguments is_op {A} _ _.
  Arguments comp {A} _ _ _.
  Arguments unit {A} _ _.
  Arguments unique_op {A} _.

(*
  Record TAlg_map {A B: Type} (AT: TAlgebra A) (BT: TAlgebra B) :=
    {
      base_map: A -> B;
      op_map: forall (f: sig AT.(is_op)), { g: sig BT.(is_op) |
              g.1 o base_map = base_map o f.1 };
      compose_map: ???
    }.  *)

End Algebras.


Require Import Coq.Lists.List.
Require Import ExtLib.Core.RelDec.
Require Import ExtLib.Structures.Maps.
Require Import ExtLib.Structures.Monads.
Require Import ExtLib.Structures.Reducible.

Set Implicit Arguments.
Set Strict Implicit.

Section keyed.
  Variable K : Type.
  Variable K_le : K -> K -> Prop.
  Variable RD_K : K -> K -> comparison.

  Inductive twothree (T : Type) : Type :=
  | Leaf
  | Two : twothree T -> K -> T -> twothree T -> twothree T
  | Three : twothree T -> K -> T -> twothree T -> K -> T -> twothree T -> twothree T.

  Arguments Leaf {T}.
  Arguments Two {T} _ _ _ _.
  Arguments Three {T} _ _ _ _ _ _ _.

  Section modify.
    Variable V : Type.
    Variable k : K.
    Variable upd : V -> option V.
    Variable def : option V.

    Fixpoint remove_greatest (m : twothree V) {T} (k_oops : unit -> T) (k_ok : K -> V -> twothree V -> T) : T :=
      match m with
        | Leaf => k_oops tt
        | Two l k v r =>
          remove_greatest r (fun _ => k_ok k v l) (fun k' v' r' => k_ok k' v' (Two l k v r'))
        | Three l k v m k' v' r =>
          remove_greatest r (fun _ => k_ok k' v' (Two l k v m)) (fun k'' v'' r'' => k_ok k'' v'' (Three l k v m k' v' r''))
      end.

    Fixpoint twothree_modify (m : twothree V) {T} (k_ok : twothree V -> T) (k_splice_up : twothree V -> K -> V -> twothree V -> T) {struct m} : T :=
      match m with
        | Leaf =>
          match def with
            | Some v => k_splice_up Leaf k v Leaf
            | None => k_ok Leaf
          end
        | Two l k' v' r =>
          match RD_K k k' with
            | Eq =>
              match upd v' with
                | Some v' => k_ok (Two l k v' r)
                | None => remove_greatest l (fun _ => k_ok r) (fun k v l => k_ok (Two l k v r))
              end
            | Lt =>
              twothree_modify l (fun l => k_ok (Two l k' v' r))
                                (fun l'' k'' v'' r'' => k_ok (Three l'' k'' v'' r'' k' v' r))
            | Gt =>
              twothree_modify r (fun r => k_ok (Two l k' v' r))
                                (fun l'' k'' v'' r'' => k_ok (Three l k' v' l'' k'' v'' r''))
          end
        | Three l k1 v1 m k2 v2 r =>
          match RD_K k k1 with
            | Eq =>
              match upd v1 with
                | Some v' => k_ok (Three l k v' m k2 v2 r)
                | None =>
                  remove_greatest l (fun _ => k_ok (Two m k2 v2 r)) (fun k1 v1 l => k_ok (Three l k1 v2 m k2 v2 r))
              end
            | Lt =>
              twothree_modify l (fun l' => k_ok (Three l' k1 v1 m k2 v2 r))
                                (fun l' k' v' r' => k_splice_up (Two l' k' v' r') k1 v1 (Two m k2 v2 r))
            | Gt =>
              match RD_K k k2 with
                | Eq =>
                  match upd v2 with
                    | Some v2 => k_ok (Three l k1 v1 m k v2 r)
                    | None =>
                      remove_greatest m (fun _ => k_ok (Two l k1 v1 r))
                                        (fun k' v' m' => k_ok (Three l k1 v1 m' k' v' r))
                  end
                | Lt =>
                  twothree_modify m (fun m' => k_ok (Three l k1 v1 m' k2 v2 r))
                                    (fun l' k' v' r' => k_splice_up (Two l k1 v1 l') k' v' (Two r' k2 v2 r))
                | Gt =>
                  twothree_modify r (fun r' => k_ok (Three l k1 v1 m k2 v2 r'))
                                    (fun l' k' v' r' => k_splice_up (Two l k1 v1 m) k2 v2 (Two l' k' v' r'))
              end
          end
      end.

  End modify.

  Definition twothree_add {V} (k : K) (v : V) (m : twothree V) : twothree V :=
    twothree_modify k (fun _ => Some v) (Some v) m (fun m => m) Two.

  Definition twothree_remove {V} (k : K) (m : twothree V) : twothree V :=
    twothree_modify k (fun _ => None) None m (fun m => m) Two.

  Fixpoint twothree_find {V} (k : K) (m : twothree V) : option V :=
    match m with
      | Leaf => None
      | Two l k' v' r =>
        match RD_K k k' with
          | Eq => Some v'
          | Lt => twothree_find k l
          | Gt => twothree_find k r
        end
      | Three l k1 v1 m k2 v2 r =>
        match RD_K k k1 with
          | Eq => Some v1
          | Lt => twothree_find k l
          | Gt => match RD_K k k2 with
                    | Eq => Some v2
                    | Lt => twothree_find k m
                    | Gt => twothree_find k r
                  end
        end
    end.

  Section fold.
    Import MonadNotation.
    Local Open Scope monad_scope.

    Variables V T : Type.
    Variable f : K -> V -> T -> T.

    Fixpoint twothree_fold (acc : T) (map : twothree V) : T :=
      match map with
        | Leaf => acc
        | Two l k v r =>
          let acc := twothree_fold acc l in
          let acc := f k v acc in
          twothree_fold acc r
        | Three l k1 v1 m k2 v2 r =>
          let acc := twothree_fold acc l in
          let acc := f k1 v1 acc in
          let acc := twothree_fold acc m in
          let acc := f k2 v2 acc in
          twothree_fold acc r
      end.

  End fold.

  Definition twothree_union {V} (m1 m2 : twothree V) : twothree V :=
    twothree_fold twothree_add m2 m1.

  Global Instance Map_twothree V : Map K V (twothree V) :=
  { empty  := Leaf
  ; add    := twothree_add
  ; remove := twothree_remove
  ; lookup := twothree_find
  ; union  := twothree_union
  }.

  Global Instance Foldable_twothree V : Foldable (twothree V) (K * V) :=
    fun _ f b x => twothree_fold (fun k v => f (k,v)) b x.

End keyed.

(** Performance Test **)
(*
Module TEST.
  Definition m := twothree nat nat.
  Instance Map_m : Map nat (twothree nat).
    apply Map_twothree.
    apply Compare_dec.nat_compare.
  Defined.

  Definition z : m :=
    (fix fill n acc : m :=
      let acc := add n n acc in
      match n with
        | 0 => acc
        | S n => fill n acc
      end) 500 empty.

  Time Eval vm_compute in
    let z := z in
    (fix find_all n : unit :=
      let _ := lookup n z in
      match n with
        | 0 => tt
        | S n => find_all n
      end) 500.
End TEST.
*)


Require Import JMeq.
Require Export CommaCategoryFunctors.
Require Import Common Notations FEqualDep.

Set Implicit Arguments.

Generalizable All Variables.

Set Asymmetric Patterns.

Set Universe Polymorphism.

Local Open Scope category_scope.

Local Ltac slice_t :=
  repeat match goal with
           | [ |- @eq (SpecializedFunctor _ _) _ _ ] => apply Functor_eq; intros
           | [ |- @JMeq (sig _) _ (sig _) _ ] => apply sig_JMeq
           | [ |- (_ -> _) = (_ -> _) ] => apply f_type_equal
           | _ => progress (intros; simpl in *; trivial; subst_body)
           | _ => progress simpl_eq
           | _ => progress (repeat rewrite Associativity; repeat rewrite LeftIdentity; repeat rewrite RightIdentity)
           | _ => progress JMeq_eq
           | _ => apply f_equal
           | [ |- JMeq (?f ?x) (?f ?y) ] =>
             apply (@f_equal1_JMeq _ _ x y f)
           | [ |- JMeq (?f ?x) (?g ?y) ] =>
             apply (@f_equal_JMeq _ _ _ _ x y f g)
           | _ =>
             progress (
               destruct_type @CommaSpecializedCategory_Object;
               destruct_type @CommaSpecializedCategory_Morphism;
               destruct_sig
             )
         end.

Section FCompositionOf.
  Context `(A : @SpecializedCategory objA).
  Context `(B : @SpecializedCategory objB).
  Context `(C : @SpecializedCategory objC).

  Lemma CommaCategoryInducedFunctor_FCompositionOf s d d'
        (m1 : Morphism ((OppositeCategory (C ^ A)) * (C ^ B)) s d)
        (m2 : Morphism ((OppositeCategory (C ^ A)) * (C ^ B)) d d') :
    CommaCategoryInducedFunctor (Compose m2 m1)
    = ComposeFunctors (CommaCategoryInducedFunctor m2) (CommaCategoryInducedFunctor m1).
    Time slice_t. (* 44 s *)
  Qed.

  Lemma CommaCategoryInducedFunctor_FIdentityOf (x : (OppositeCategory (C ^ A)) * (C ^ B)) :
    CommaCategoryInducedFunctor (Identity x)
    = IdentityFunctor _.
    Time slice_t. (* 11 s *)
  Qed.
End FCompositionOf.


Set Implicit Arguments.

Require Import FCF.
Require Import PRF.
Require Import RndInList.

Local Open Scope list_scope.

Inductive in_oc_supp : forall (A B C : Set), C -> OracleComp A B C -> Prop :=
| in_oc_supp_Bind :
    forall (A B C C' : Set)(c : OracleComp A B C')(f : C' -> OracleComp A B C) x y,
      in_oc_supp x c ->
      in_oc_supp y (f x) ->
      in_oc_supp y (OC_Bind c f)
| in_oc_supp_Query :
    forall (A B : Set)(a : A)(b : B),
      in_oc_supp b (OC_Query B a)
| in_oc_supp_Ret :
  forall (A B C : Set)(c : Comp C) x,
    In x (getSupport c) ->
    in_oc_supp x (OC_Ret A B c)
| in_oc_supp_Run :
  forall (A A' B B' C S : Set)(eqds : EqDec S)(eqda : EqDec A)(eqdb : EqDec B)
    (c : OracleComp A B C)(oc : S -> A -> OracleComp A' B' (B * S)) s x s',
    in_oc_supp x c ->
    in_oc_supp (x, s') (OC_Run _ _ _ c oc s).

Theorem in_oc_supp_complete :
  forall (A B C : Set){eqdc : EqDec C}(c : OracleComp A B C) x
    (S : Set)(eqds : EqDec S)(o : S -> A -> Comp (B * S))(s : S) s',
    In (x, s') (getSupport (c _ _ o s)) ->
    in_oc_supp x c .

  induction c; intuition; repeat simp_in_support.
  econstructor.

  Local Opaque getSupport.
  simpl in *.
  repeat simp_in_support.
  destruct x.
  simpl.
  destruct p.
  simpl.

  econstructor.
  eapply (@IHc _ _ (S * S0)%type).
  eapply H1.

  simpl in *.
  repeat simp_in_support.
  econstructor.
  intuition.

  assert (EqDec C).
  eapply oc_EqDec.
  eapply c.
  intuition.
  assert (B * S).
  eapply comp_base_exists.
  eapply o0.
  trivial.
  intuition.
  intuition.
  intuition.
  assert (EqDec (B * S)).
  eapply comp_EqDec.
  eapply o0; trivial.
  eapply EqDec_pair_l.
  eapply H2.
  trivial.

  simpl in *.
  repeat simp_in_support.
  destruct x0.
  econstructor.
  eapply (@H0 _ S).
  eapply H2.
  eapply (@H _ _ _ S).
  eapply H3.

  Grab Existential Variables.
  eapply EqDec_pair_l.
  eauto.
  trivial.

Qed.

Local Open Scope nat_scope.

Inductive queries_at_most' : forall (A B C : Set), OracleComp A B C -> nat -> Prop :=
| qam_Bind :
  forall (A B C C' : Set)(eqdc' : EqDec C')(c : OracleComp A B C')(f : C' -> OracleComp A B C) q1 q2,
    queries_at_most' c q1 ->
    (forall c',
       in_oc_supp c' c ->
       queries_at_most' (f c') q2) ->
    queries_at_most' (OC_Bind c f) (q1 + q2)
| qam_Query :
  forall (A B : Set)(a : A),
  queries_at_most' (OC_Query B a) 1
| qam_Ret :
  forall (A B C : Set)(c : Comp C),
    queries_at_most' (OC_Ret A B c) 0
| qam_Run :
  forall (A A' B B' C S : Set)(eqds : EqDec S)(eqda : EqDec A)(eqdb : EqDec B)
    (c : OracleComp A B C)(oc : S -> A -> OracleComp A' B' (B * S)) s q1 q2,
    queries_at_most' c q1 ->
    (forall s a, queries_at_most' (oc s a) q2) ->
    queries_at_most' (OC_Run _ _ _ c oc s) (q1 * q2)
| qam_le :
  forall (A B C : Set)(c : OracleComp A B C) q1 q2,
    queries_at_most' c q1 ->
    q1 <= q2 ->
    queries_at_most' c q2.

Theorem qam_count_gen' :
  forall (A B C : Set)(c : OracleComp A B C)(q : nat),
    queries_at_most' c q ->
    forall (S : Set)(count : S -> nat)(eqds : EqDec S)(o : S -> A -> Comp (B * S))(s : S)(n : nat),
      (forall a b x y,
        In (a, b) (getSupport (o x y)) ->
          count b <= n + (count x)) ->
      forall a b,
      In (a, b) (getSupport (c _ _ o s)) ->
      count b <= q * n + (count s).

  Opaque getSupport.
  induction 1; intuition; simpl in *.

  (* Bind case *)
  repeat simp_in_support.
  destruct x.
  rewrite mult_plus_distr_r.
  pose proof H4.
  eapply (IHqueries_at_most' _ count) in H4.
  eapply (H1 _ _ _ count) in H5.
  rewrite H5.
  rewrite plus_comm.
  rewrite <- (plus_assoc).
  rewrite (plus_comm (q2 * n)).
  rewrite plus_assoc.
  eapply plus_le_compat.
  eauto.
  eapply le_refl.
  intuition.
  intuition.

  (* query case *)
  rewrite plus_0_r.
  eapply H; eauto.

  (* ret case *)
  Transparent getSupport.
  repeat simp_in_support.
  intuition.

  (* run case *)
  repeat simp_in_support.

  destruct x.
  simpl in *.

  specialize (IHqueries_at_most' (S * S0)%type (fun p =>  count (snd p)) _
  (fun (x : S * S0) (y : A) =>
                p <-$ (oc (fst x) y) S0 eqds0 o (snd x);
                ret (fst (fst p), (snd (fst p), snd p)))
  (s, s0) (q2 * n)
  ).

  eapply le_trans.
  eapply IHqueries_at_most'.
  intuition.
  simpl.
  repeat simp_in_support.
  destruct x.
  simpl in *.
  eapply H1.
  eauto.
  eauto.
  eauto.
  simpl.
  rewrite mult_assoc.
  intuition.

  (* trans case *)
  eapply le_trans.
  eapply IHqueries_at_most'.
  eauto.
  eauto.
  eapply plus_le_compat; intuition.
  eapply mult_le_compat; intuition.

  Grab Existential Variables.
  eapply in_oc_supp_complete; eauto.

Qed.

(*
Definition countOracle{A B State : Set}{eqdb : EqDec B}{eqds : EqDec State}(defB : B)(o : State -> A -> Comp (B * State)) p a :=
  [s, n] <-2 p;
  match n with
    | 0 => ret (defB, p)
    | S n' =>
      [b', s'] <-$2 o s a;
        ret (b', (s', n'))
  end.

Theorem countOracle_equiv :
  forall (A B C : Set)(defB : B)(eqdb : EqDec B)(eqdc : EqDec C)(c : OracleComp A B C) n,
    queries_at_most' c n ->
    forall (S : Set)(eqd : EqDec S)(o : S -> A -> Comp (B * S))(s : S),
      comp_spec
        (fun a b => fst a = fst b /\ snd a = fst (snd b))
        (c _ _ o s)
        (c _ _ (countOracle defB o) (s, n)).

  induction 1; intuition; simpl in *.

  admit.
  admit.
  admit.

  comp_skip.
  admit.
  admit.
  assert (EqDec C).
  eapply EqDec_pair_l.
  eauto.
  trivial.

  assert (B * S).
  eapply oc_base_exists.
  eapply oc

  specialize (@IHqueries_at_most' eqdb0 H2 defB (S * S0)%type).
  eapply H1.
  eapply oc_comp_spec_eq.

Qed.

*)

Local Open Scope rat_scope.

Theorem evalDist_bind_event_le :
  forall (A : Set)(c : Comp A)(f : A -> Comp bool)(evta : A -> bool) (k1 k2 : Rat),
    Pr[a <-$ c; ret (evta a)] <= k1 ->
    (forall a, In a (getSupport c) -> evta a = false -> Pr[f a] <= k2) ->
    Pr[a <-$ c; f a] <= k1 + k2.

  intuition.
  simpl in *.
  rewrite (sumList_partition evta).

  eapply ratAdd_leRat_compat.
  assert ( sumList (getSupport c)
     (fun a : A => evalDist c a * Pr  [f a ] * (if evta a then 1 else 0)) <=
            sumList (getSupport c)
        (fun b : A =>
         evalDist c b * (if EqDec_dec bool_EqDec (evta b) true then 1 else 0)) ).

  eapply sumList_le; intuition.
  destruct (EqDec_dec bool_EqDec (evta a) true ).
  rewrite e.

  eapply ratMult_leRat_compat; intuition.
  eapply leRat_trans.
  eapply ratMult_leRat_compat.
  eapply leRat_refl.
  eapply evalDist_le_1.
  rewrite ratMult_1_r.
  intuition.

  destruct (evta a); intuition.
  repeat rewrite ratMult_0_r.
  eapply rat0_le_all.
  rewrite H1.
  eapply H.

  assert ( sumList (getSupport c)
     (fun a : A => evalDist c a * Pr  [f a ] * (if evta a then 0 else 1)) <=
            sumList (getSupport c)
     (fun a : A => k2 * (evalDist c a * (if evta a then 0 else 1)))).

  eapply sumList_le.
  intuition.
  case_eq (evta a); intuition.
  eapply eqRat_impl_leRat.
  repeat rewrite ratMult_0_r.
  intuition.
  rewrite H0.
  eapply eqRat_impl_leRat.
  repeat rewrite ratMult_1_r.
  eapply ratMult_comm.
  trivial.
  trivial.

  rewrite H1.
  rewrite sumList_factor_constant_l.

  eapply leRat_trans.
  Focus 2.
  eapply eqRat_impl_leRat.
  eapply ratMult_1_r.
  eapply ratMult_leRat_compat; intuition.

  assert (sumList (getSupport c)
     (fun a : A => evalDist c a * (if evta a then 0 else 1)) <=
          sumList (getSupport c)
     (fun a : A => evalDist c a)).

  eapply sumList_le.
  intuition.
  eapply leRat_trans.
  Focus 2.
  eapply eqRat_impl_leRat.
  eapply ratMult_1_r.
  eapply ratMult_leRat_compat; intuition.
  destruct (evta a); intuition.
  eapply rat0_le_all.
  rewrite H2.
  eapply evalDist_sum_le_1.
Qed.

Theorem oc_eventProb :
  forall (A B C : Set)(c : OracleComp A B C) n,
    queries_at_most' c n ->
    forall
    (S : Set)(eqds : EqDec S)(o : S -> A -> Comp (B * S))
    (count : S -> nat)(evt : S -> bool)(s : S)(k : nat -> Rat) i,
      (forall (n1 n2 : nat), (n1 <= n2)%nat -> (k n1 <= k n2)) ->
      evt s = false ->
      (forall s a, evt s = false -> Pr[p <-$ o s a; ret (evt (snd p))] <= (k (i + (count s))%nat)) ->
      (forall s s' a b, In (b, s') (getSupport (o s a)) ->
                        count s' <= i + (count s))%nat ->
      Pr[p <-$ c _ _ o s; ret (evt (snd p))] <= (n / 1) * (k (i * n + (count s))%nat).

  Local Opaque evalDist.

  induction 1; intuition; simpl in *.
  inline_first.

  assert (
       Pr
   [a <-$ c S eqds o s;
    p <-$ ([z, s']<-2 a; (f z) S eqds o s'); ret evt (snd p) ] <=
   q1 / 1 * k (i * q1 + count s)%nat +
   q2 / 1 * k (i * q2 + (i * q1 + count s))%nat
                     ).

  eapply evalDist_bind_event_le.
  eapply (@IHqueries_at_most' _ _ _ count evt); intuition.
  intros.
  comp_simp.
  eapply leRat_trans.
  eapply (@H1 _ _ _ _ _ count evt).
  eapply H2.
  trivial.
  intuition.
  intuition.
  eapply ratMult_leRat_compat; intuition.
  eapply H2.
  eapply plus_le_compat; intuition.

  Require Import RndInList.
  eapply le_trans.
  eapply (qam_count_gen' _ _ _ _ _ i).
  intuition.
  eapply H5.
  eauto.
  eauto.
  rewrite mult_comm.
  intuition.

  rewrite H6.
  clear H6.

  eapply leRat_trans.
  Focus 2.
  eapply eqRat_impl_leRat.
  symmetry.
  rewrite ratAdd_num.
  eapply ratMult_distrib_r.
  eapply ratAdd_leRat_compat;
  eapply ratMult_leRat_compat; intuition.
  eapply H2.
  rewrite mult_plus_distr_l.
  rewrite plus_comm.
  repeat rewrite <- plus_assoc.
  eapply plus_le_compat; intuition.

  rewrite H1; intuition.
  rewrite <- ratMult_1_l.
  eapply ratMult_leRat_compat; intuition.
  eapply leRat_terms; intuition.

  inline_first.
  eapply distro_irr_le.
  intuition.
  comp_simp.
  simpl.
  rewrite H0.
  rewrite evalDist_ret_0.
  eapply rat0_le_all.
  congruence.

  inline_first.

  assert (EqDec C).
  eapply oc_EqDec; intros.
  eapply c.
  assert (B * S).
  eapply oc_base_exists.
  eapply oc.
  trivial.
  trivial.
  intuition.
  assert (B' * S0).
  eapply comp_base_exists.
  eapply o.
  trivial.
  trivial.
  intuition.
  intuition.
  trivial.

  assert
    (
      (evalDist
        (Bind
           (c (prod S S0) (pair_EqDec eqds eqds0)
              (fun (x : prod S S0) (y : A) =>
               Bind ((oc (fst x) y) S0 eqds0 o (snd x))
                 (fun p : prod (prod B S) S0 =>
                  Ret (EqDec_dec (pair_EqDec eqdb (pair_EqDec eqds eqds0)))
                    (pair (fst (fst p)) (pair (snd (fst p)) (snd p)))))
              (pair s s0))
           (fun a : prod C (prod S S0) =>
            Bind
              (Ret
                 (EqDec_dec
                    (pair_EqDec
                       (pair_EqDec
                          (oc_EqDec c
                             (fun x : A =>
                              fst
                                (oc_base_exists (oc s x)
                                   (fun y : A' =>
                                    fst (comp_base_exists (o s0 y)))))
                             (fun x : A =>
                              EqDec_pair_l
                                (oc_EqDec (oc s x)
                                   (fun y : A' =>
                                    fst (comp_base_exists (o s0 y)))
                                   (fun y : A' =>
                                    EqDec_pair_l (comp_EqDec (o s0 y)) s0)) s))
                          eqds) eqds0))
                 (pair (pair (fst a) (fst (snd a))) (snd (snd a))))
              (fun p : prod (prod C S) S0 =>
               Ret (EqDec_dec bool_EqDec) (evt (snd p))))) true)
   ==
   Pr
   [a <-$
    c (S * S0)%type (pair_EqDec eqds eqds0)
      (fun (x : S * S0) (y : A) =>
       p <-$ (oc (fst x) y) S0 eqds0 o (snd x);
       ret (fst (fst p), (snd (fst p), snd p))) (s, s0);
     ret (evt (snd (snd a))) ]
   ).

  comp_skip.
  simpl.
  intuition.
  rewrite H7.
  clear H7.

  eapply leRat_trans.
  eapply (@IHqueries_at_most' _ _
                             ((fun (x : S * S0) (y : A) =>
       p <-$ (oc (fst x) y) S0 eqds0 o (snd x);
       ret (fst (fst p), (snd (fst p), snd p))))
                             (fun p => count (snd p)) (fun p => evt (snd p))
         _ (fun x => q2 / 1 * k (x)%nat) (i * q2)%nat).
  intuition.
  eapply ratMult_leRat_compat; intuition.
  trivial.
  intuition.

  comp_inline_l.
  assert (
      Pr
   [a1 <-$ (oc (fst (a, b)) a0) S0 eqds0 o (snd (a, b));
    p <-$ ret (fst (fst a1), (snd (fst a1), snd a1)); ret evt (snd (snd p)) ]
   ==
   Pr
   [a1 <-$ (oc (fst (a, b)) a0) S0 eqds0 o (snd (a, b));
    ret evt (snd a1) ]
   ).

  comp_skip.
  simpl.
  intuition.
  rewrite H8.
  clear H8.
  simpl.
  eapply leRat_trans.

  eapply (@H1 ); intuition.
  reflexivity.
  intuition.
  repeat simp_in_support.
  destruct x.
  simpl in *.
  rewrite mult_comm.
  eapply qam_count_gen'.
  eapply H0.
  intuition.
  eapply H5.
  eapply H7.
  eapply H8.
  simpl.
  rewrite <- ratMult_assoc.
  eapply ratMult_leRat_compat; intuition.
  eapply eqRat_impl_leRat.
  rewrite <- ratMult_num_den.
  eapply eqRat_terms; intuition.
  eapply H2.
  eapply plus_le_compat; intuition.
  rewrite <- mult_assoc.
  eapply mult_le_compat; intuition.
  rewrite mult_comm.
  intuition.

  rewrite IHqueries_at_most'.
  eapply ratMult_leRat_compat; intuition.
  eapply leRat_terms; intuition.
  intuition.
  trivial.
  intuition.
  intuition.

  Grab Existential Variables.
  trivial.
  eapply in_oc_supp_complete; eauto.

Qed.

Theorem oc_eventProb_0_1 :
  forall (S : Set)(count : S -> nat)(evt : S -> bool)(k : nat -> Rat)
         (A B C : Set)(c : OracleComp A B C) n,
    queries_at_most' c n ->
    forall
    (eqds : EqDec S)(o : S -> A -> Comp (B * S))
    (s : S),
      (forall (n1 n2 : nat), (n1 <= n2)%nat -> (k n1 <= k n2)) ->
      evt s = false ->
      (forall s a, evt s = false -> Pr[p <-$ o s a; ret (evt (snd p))] <= (k (1 + (count s))%nat)) ->
      (forall s s' a b, In (b, s') (getSupport (o s a)) ->
                        count s' <= 1 + (count s))%nat ->
      count s = 0%nat ->
      Pr[p <-$ c _ _ o s; ret (evt (snd p))] <= (n / 1) * (k n).

  intuition.
  eapply leRat_trans.
  eapply oc_eventProb; intuition.
  eauto.
  rewrite H4.
  rewrite mult_1_l.
  rewrite plus_0_r.
  intuition.
Qed.

(*
Definition maxCount_o(D R S : Set){eqdr : EqDec R}{eqds : EqDec S}(count : S -> nat)(o : S -> D -> Comp (R * S))
           (defR : R)(m : nat)(s : S)(d : D) :=
  if (ge_dec (count s) m) then (ret (defR, s)) else
    o s d.

Theorem maxCount_o_equiv :
    forall (D R C: Set)(A : OracleComp D R C) (n1: nat),
      queries_at_most A n1 ->
      forall (S : Set)(count : S -> nat)(o : S -> D -> Comp (R * S))(defR : R)(eqdr : EqDec R)(eqds : EqDec S)(RndR : Comp R) s x n2,
        (n1 + (count s) <= n2)%nat ->
        evalDist (A _ _ o s) x ==
        evalDist (A _ _ (maxCount_o count o defR n2) s) x.

  Local Opaque evalDist.
  induction 1; intuition; simpl in *.

  comp_skip.
  comp_simp.
  eapply H1.

  specialize (qam_count_gen H (fun (p : list (A * B)) => length p) _ (randomFunc RndR eqdd) ls); intuition.
  eapply le_trans.
  eapply plus_le_compat.
  eapply le_refl.
  eapply H4.
  intuition; simpl in *.
  unfold randomFunc in *.
  destruct (arrayLookup eqdd x y); repeat simp_in_support.
  omega.
  simpl.
  assert ( (S (length x) <= 1+ length x)%nat).
  omega.
  eapply H5.
  eapply H3.
  omega.

  unfold randomFunc, randomFunc_max.
  destruct (arrayLookup eqdd ls a).
  eapply evalDist_ret_eq; intuition.
  comp_skip.
  destruct (ge_dec (length ls) n2).
  omega.
  eapply evalDist_ret_eq; intuition.

  eapply eqRat_refl.

  comp_skip.
  remember (map

  eapply (@IHqueries_at_most _ _ RndR ls.

  trivial.

  erewrite H4.

  Print RndInList.

  rewrite <- H2.
  eapply plus_le_compat; intuition.
  omega.

Qed.
*)

Section RandPermSwitching.

  Variable A D R : Set.
  Variable RndR : Comp R.
  Hypothesis RndR_wf : well_formed_comp RndR.

  Variable A2 : OracleComp D R bool.

  Hypothesis eqdd : EqDec D.
  Hypothesis eqdr : EqDec R.

  Variable fBad : list (D * R) -> D -> R -> bool.
  Variable badProb : nat -> Rat.
  Hypothesis badProb_mono : forall n1 n2, (n1 <= n2)%nat -> badProb n1 <= badProb n2.

  Hypothesis A2_wf : well_formed_oc A2.
  Variable A2_queries : nat.
  Hypothesis A2_queries_correct : queries_at_most' A2 A2_queries.
  Hypothesis badProb_correct :
    forall d f,
      Pr [r <-$ RndR;
           ret (fBad f d r)] <= badProb (S (length f)).

   Hypothesis goodExists :
    forall d f,
      (length f <= A2_queries) %nat ->
      exists r,
        In r (getSupport RndR) /\
        fBad f d r = false.

  (*
  Hypothesis fBad_single :
    forall d r,
      fBad ((d, r) :: nil) = false.
*)

  (*
  Hypothesis goodExists :
    forall d f,
      arrayLookup _ f d = None ->
      exists r,
        In r (getSupport RndR) /\
        fBad f d r = false.
*)

  Definition rndNotBad(d : D)(f : list (D * R)):=
    Repeat RndR (fun r => negb (fBad f d r)).

  Definition GenRP(s : list (D * R))(d : D) : Comp (R * list (D * R)) :=
    match arrayLookup _ s d with
      | Some r => ret (r, s)
      | None => r <-$ (rndNotBad d s); ret (r, (d, r) :: s)
    end.

  Definition RPS_G0  :=
    [b, _] <-$2 A2 _ _ (@randomFunc D R RndR _ ) nil;
    ret b.

  Definition RPS_G1 :=
    [b, _] <-$2 A2 _ _ GenRP nil;
    ret b.

   Definition randomFunc_bad (s : list (D * R) * bool) d :=
    [s, bad] <-2 s;
    match (arrayLookup _ s d) with
        | Some r => ret (r, (s, bad))
        | None =>
          r <-$ RndR;
          newF <- ((d, r) :: s);
          bad' <- fBad s d r;
          ret (r, (newF, bad || bad'))
    end.

   Definition GenRP_bad (s : list (D * R) * bool) d :=
    [s, bad] <-2 s;
    match (arrayLookup _ s d) with
        | Some r => ret (r, (s, bad))
        | None =>
          r <-$ RndR;
          newF <- ((d, r) :: s);
          bad' <- fBad s d r;
          if (bad') then
            (r <-$ rndNotBad d s;
            ret (r, ((d, r) :: s, bad || bad')))
          else
            ret (r, (newF, bad || bad'))
    end.

  Definition RPS_G_1 :=
    [b, s] <-$2 A2 _ _ (randomFunc_bad) (nil, false);
    ret (b, snd s).

  Definition RPS_G_2 :=
    [b, s] <-$2 A2 _ _ (GenRP_bad) (nil, false);
    ret (b, snd s).

  Theorem randomFunc_bad_wf :
    forall a b,
      well_formed_comp (randomFunc_bad a b).

    intuition.
    unfold randomFunc_bad.
    destruct (arrayLookup eqdd a0 b0); wftac.

  Qed.

  Theorem GenRP_bad_wf :
    forall a b,
      (length (fst a) <= A2_queries)%nat ->
      well_formed_comp (GenRP_bad a b).

    intuition.
    simpl in *.
    case_eq (arrayLookup eqdd a0 b0); intuition; wftac.
    unfold rndNotBad; wftac.
    edestruct goodExists ; eauto.
    intuition.
    econstructor.
    unfold eq_dec; intuition.
    eapply (EqDec_dec _).
    trivial.
    eapply filter_In; intuition.
    eauto.
    rewrite H4.
    trivial.

  Qed.

  Theorem rf_rp_oc_eq_until_bad :
    comp_spec
      (fun y1 y2 =>
         snd (snd y1) = snd (snd y2) /\
          (snd (snd y1) = false -> fst (snd y1) = fst (snd y2) /\ fst y1 = fst y2))
      (A2 (list (D * R) * bool)%type _
          randomFunc_bad (nil, false))
      (A2 (list (D * R) * bool)%type _
          GenRP_bad (nil, false)).

    eapply comp_spec_consequence.
    eapply (@oc_comp_spec_eq_until_bad _ _ _ _ _ _ _ _ _ _ _ _ _
                                       (fun a => snd a)
                                       (fun a => snd a)
                                       (fun a b => (fst a = fst b))).
    intuition.
    eapply randomFunc_bad_wf.
    intuition.
    eapply GenRP_bad_wf.
    admit.

    intuition.
    simpl in *.
    subst.

    case_eq (arrayLookup eqdd a0 a1); intuition.
    eapply comp_spec_ret; intuition.

    comp_skip.
    eapply comp_base_exists; eauto.
    eapply comp_base_exists; eauto.

    case_eq ( fBad a0 a1 b); intuition.
    comp_irr_r.
    unfold rndNotBad.
    edestruct goodExists.
    eauto.
    intuition.
    econstructor.
    unfold eq_dec; intuition.
    eapply (EqDec_dec _).
    trivial.
    eapply filter_In; intuition.
    eauto.
    rewrite H5.
    trivial.
    eapply comp_spec_ret.
    simpl.
    split.
    intuition.
    intros.
    rewrite orb_true_r in H4.
    discriminate.
    eapply comp_spec_ret; intuition.
    intuition.
    simpl in *.
    case_eq (arrayLookup eqdd a1 d); intuition.
    rewrite H1 in H0.
    repeat simp_in_support.
    trivial.
    rewrite H1 in H0.
    repeat simp_in_support.
    trivial.

    intuition.
    simpl in *.
    case_eq (arrayLookup eqdd a1 d ); intuition;
    rewrite H1 in H0;
    repeat simp_in_support;
    trivial.

    trivial.
    trivial.

    intuition.

    Grab Existential Variables.
    trivial.

  Qed.

  Theorem RPS_G_1_2_bad_same :
     Pr  [x <-$ RPS_G_1; ret snd x ] == Pr  [x <-$ RPS_G_2; ret snd x ].

    unfold RPS_G_1, RPS_G_2.
    inline_first.

    eapply comp_spec_impl_eq.
    comp_skip.
    eapply rf_rp_oc_eq_until_bad.

    intuition.
    simpl in *.
    intuition.
    destruct b.
    simpl in *.
    destruct p; simpl in *.
    subst.
    comp_simp.
    simpl.
    eapply comp_spec_ret; intuition.

  Qed.

  Theorem RPG_G_1_2_eq_until_bad :
    forall a : bool,
      evalDist RPS_G_1 (a, false) == evalDist RPS_G_2 (a, false).

    intuition.
    unfold RPS_G_1, RPS_G_2.
    eapply comp_spec_impl_eq.
    comp_skip.
    eapply rf_rp_oc_eq_until_bad.

    intuition.
    simpl in *.
    intuition.
    destruct b.
    simpl in *.
    destruct p; simpl in *.
    eapply comp_spec_ret; intuition.
    pairInv.
    subst.
    intuition.
    subst.
    intuition.
    subst.
    pairInv.
    intuition.
    subst.
    trivial.

  Qed.

  Theorem RPS_G_1_bad_small :
    Pr  [x <-$ RPS_G_1; ret snd x ] <= A2_queries / 1 * badProb A2_queries.

    unfold RPS_G_1.
    inline_first.

    assert (Pr
   [a <-$
    A2 (list (D * R) * bool)%type _
      (randomFunc_bad) (nil, false);
    x <-$ ([b, s]<-2 a; ret (b, snd s)); ret snd x ]
   ==
   Pr
   [a <-$
    A2 (list (D * R) * bool)%type _
      (randomFunc_bad) (nil, false);
    ret snd (snd a) ]
   ).
    comp_skip.
    comp_simp.
    reflexivity.
    rewrite H.
    clear H.

    eapply leRat_trans.
    eapply (@oc_eventProb_0_1 _ (fun p => length (fst p)) _ badProb).
    eauto.
    eapply badProb_mono.
    trivial.
    intuition.
    simpl in *.
    subst.
    case_eq (arrayLookup eqdd a a0); intuition.
    comp_simp.
    simpl.
    rewrite evalDist_ret_0.
    eapply rat0_le_all.
    congruence.

    simpl in *.
    eapply leRat_trans.
    Focus 2.
    eapply badProb_correct.
    inline_first.
    comp_skip.
    simpl.
    eapply leRat_refl.
    intuition.
    simpl in *.
    destruct (arrayLookup eqdd a a1);
    repeat simp_in_support.
    omega.
    simpl.
    intuition.

    trivial.
    intuition.

  Qed.

  Theorem RPS_G_1_2_close :
  | Pr[x <-$ RPS_G_1; ret fst x] - Pr[x <-$ RPS_G_2; ret fst x] | <= (A2_queries / 1)  * badProb A2_queries.

    intuition.
    eapply leRat_trans.
    eapply fundamental_lemma_h.

    (* badness is the same. *)
    eapply RPS_G_1_2_bad_same.

    eapply RPG_G_1_2_eq_until_bad.

    eapply RPS_G_1_bad_small.

  Qed.

  Theorem RPS_G0_equiv :
    Pr[RPS_G0] == Pr[x <-$ RPS_G_1; ret fst x].

    unfold RPS_G0, RPS_G_1.
    inline_first.
    eapply (comp_spec_eq_impl_eq).
    comp_skip.
    eapply (@oc_comp_spec_eq _ _ _ _ _ _ _ _ _ _ _ _ _ _ (fun a b => a = fst b)); trivial.
    intuition.
    subst.
    unfold randomFunc, randomFunc_bad.
    comp_simp.
    simpl.
    case_eq ( arrayLookup eqdd l a ); intuition.
    eapply comp_spec_ret; intuition.
    comp_skip.
    eapply comp_base_exists; eauto.
    eapply comp_base_exists; eauto.
    eapply comp_spec_ret; intuition.

    comp_simp.
    simpl in *; intuition.
    subst.
    eapply comp_spec_eq_refl.

  Qed.

  Theorem RPS_G1_equiv :
    Pr[x <-$ RPS_G_2; ret fst x] == Pr[RPS_G1].

    unfold RPS_G_2, RPS_G1.
    inline_first.
    eapply comp_spec_eq_impl_eq.
    comp_skip.
    eapply (@oc_comp_spec_eq _ _ _ _ _ _ _ _ _ _ _ _ _ _ (fun a b => fst a = b)); trivial.
    intuition.
    unfold GenRP, GenRP_bad.
    subst.
    simpl.
    case_eq (arrayLookup eqdd a0 a); intuition.
    eapply comp_spec_ret; intuition.

    unfold rndNotBad.

    assert (
        comp_spec
          eq
          (r <-$ (r <-$ RndR; if (negb (fBad a0 a r)) then (ret r) else (Repeat RndR (fun r : R => negb (fBad a0 a r))));

           ret (r, (a, r) :: a0))
          (r <-$ Repeat RndR (fun r : R => negb (fBad a0 a r));
           ret (r, (a, r) :: a0))
      ).

    eapply eq_impl_comp_spec_eq.
    intuition.
    symmetry.
    comp_skip.
    eapply repeat_unroll_eq.
    trivial.
    edestruct goodExists; eauto.
    intuition.
    econstructor.
    eapply filter_In; intuition.
    eauto.
    rewrite H2.
    trivial.
    eapply comp_spec_eq_trans_r.
    Focus 2.
    eapply H0.
    clear H0.
    inline_first.
    comp_skip.
    eapply comp_base_exists; eauto.
    eapply comp_base_exists; eauto.
    case_eq (fBad a0 a b0); intuition.
    simpl.
    comp_skip.
    eapply comp_spec_ret; intuition.
    simpl.
    comp_simp.
    eapply comp_spec_ret; intuition.
    comp_simp.
    simpl in *.
    intuition; subst.
    eapply comp_spec_eq_refl.

  Qed.

  Theorem RPS_G0_1_close :
  | Pr[RPS_G0] - Pr[RPS_G1] | <= (A2_queries / 1)  * badProb A2_queries.

    rewrite RPS_G0_equiv.
    rewrite <- RPS_G1_equiv.
    eapply RPS_G_1_2_close.
  Qed.

End RandPermSwitching.