From Perennial.program_proof.mvcc Require Import
     wrbuf_prelude wrbuf_repr
     index_proof
     tuple_repr tuple_own tuple_free tuple_write_lock.

Section heap.
Context `{!heapGS Σ, !mvcc_ghostG Σ}.

(*****************************************************************)
(* func swap(ents []WrEnt, i, j uint64)                          *)
(*****************************************************************)
Local Lemma wp_swap (entsS : Slice.t) (ents : list wrent) (i j : u64) :
  {{{ slice.is_slice_small entsS (structTy WrEnt) 1 (wrent_to_val <$> ents) ∗
      ⌜(int.nat i < length ents ∧ int.nat j < length ents)%nat⌝
  }}}
    swap (to_val entsS) #i #j
  {{{ (ents' : list wrent), RET #();
      slice.is_slice_small entsS (structTy WrEnt) 1 (wrent_to_val <$> ents') ∗
      ⌜ents' ≡ₚ ents⌝
  }}}.
Proof.
  iIntros (Φ) "[HentsS %Hbounds] HΦ".
  destruct Hbounds as [Hi Hj].
  apply list_lookup_lt in Hi as Hx.
  destruct Hx as [x Hx].
  apply list_lookup_lt in Hj as Hy.
  destruct Hy as [y Hy].
  iDestruct (is_slice_small_sz with "HentsS") as "%HentsLen".
  rewrite fmap_length in HentsLen.
  wp_call.

  (***********************************************************)
  (* tmp := ents[i]                                          *)
  (* ents[i] = ents[j]                                       *)
  (* ents[j] = tmp                                           *)
  (***********************************************************)
  wp_apply (wp_SliceGet with "[$HentsS]").
  { iPureIntro. by rewrite list_lookup_fmap Hx. }
  iIntros "[HentsS %Htyx]".
  wp_apply (wp_SliceGet with "[$HentsS]").
  { iPureIntro. by rewrite list_lookup_fmap Hy. }
  iIntros "[HentsS %Htyy]".
  wp_apply (wp_SliceSet with "[$HentsS]").
  { iPureIntro. split; last done. by rewrite list_lookup_fmap fmap_is_Some. }
  iIntros "HentsS".
  wp_apply (wp_SliceSet with "[$HentsS]").
  { iPureIntro. split; last done.
    rewrite -list_fmap_insert list_lookup_fmap fmap_is_Some.
    apply lookup_lt_is_Some_2.
    rewrite insert_length.
    word.
  }
  iIntros "HentsS".
  wp_pures.
  iApply "HΦ".
  iFrame.
  do 2 rewrite -list_fmap_insert.
  iFrame.
  iPureIntro.
  apply list_insert_insert_swap; done.
Qed.

(*********************************************************************)
(* func (wrbuf *WrBuf) sortEntsByKey()                               *)
(*********************************************************************)
Theorem wp_wrbuf__sortEntsByKey wrbuf mods :
  {{{ own_wrbuf_xtpls wrbuf mods }}}
    WrBuf__sortEntsByKey #wrbuf
  {{{ RET #(); own_wrbuf_xtpls wrbuf mods }}}.
Proof.
  iIntros (Φ) "Hwrbuf HΦ".
  wp_call.

  (***********************************************************)
  (* ents := wrbuf.ents                                      *)
  (* var i uint64 = 1                                        *)
  (***********************************************************)
  iNamed "Hwrbuf".
  wp_loadField.
  wp_pures.
  wp_apply (wp_ref_to); first by auto.
  iIntros (iRef) "HiRef".
  wp_pures.
  iDestruct (is_slice_small_acc with "HentsS") as "[HentsS HentsC]".
  iDestruct (is_slice_small_sz with "HentsS") as "%HentsLen".
  rewrite fmap_length in HentsLen.

  (***********************************************************)
  (* for i < uint64(len(ents)) {                             *)
  (*     var j uint64 = i                                    *)
  (*     for j > 0 {                                         *)
  (*         if ents[j - 1].key <= ents[j].key {             *)
  (*             break                                       *)
  (*         }                                               *)
  (*         swap(ents, j - 1, j)                            *)
  (*         j--                                             *)
  (*     }                                                   *)
  (*     i++                                                 *)
  (* }                                                       *)
  (***********************************************************)
  set P := (λ (b : bool), ∃ (ents' : list wrent) (i : u64),
               "HentsS" ∷ is_slice_small entsS (struct.t WrEnt) 1 (wrent_to_val <$> ents') ∗
               "HiRef"  ∷ iRef ↦[uint64T] #i ∗
               "%Hperm" ∷ ⌜ents' ≡ₚ ents⌝
           )%I.
  wp_apply (wp_forBreak_cond P with "[] [HentsS HiRef]").
  { (* Outer loop body. *)
    clear Φ.
    iIntros (Φ) "!> HP HΦ".
    iNamed "HP".
    wp_load.
    wp_apply (wp_slice_len).
    wp_if_destruct; last first.
    { (* Outer loop condition. *)
      iApply "HΦ".
      subst P. simpl.
      eauto 10 with iFrame.
    }
    wp_load.
    wp_apply (wp_ref_to); first by auto.
    iIntros (jRef) "HjRef".
    wp_pures.

    (* Inner loop. *)
    set Q := (λ (b : bool), ∃ (ents'' : list wrent) (j : u64),
                "HentsS" ∷ is_slice_small entsS (struct.t WrEnt) 1 (wrent_to_val <$> ents'') ∗
                "HjRef"  ∷ jRef ↦[uint64T] #j ∗
                "%Hperm" ∷ ⌜ents'' ≡ₚ ents'⌝ ∗
                "%Hle" ∷ ⌜int.Z j ≤ int.Z i⌝
            )%I.
    wp_apply (wp_forBreak_cond Q with "[] [HentsS HjRef]").
    { (* Inner loop body. *)
      clear Φ.
      iIntros (Φ) "!> HQ HΦ".
      iNamed "HQ".
      wp_load.
      wp_if_destruct; last first.
      { (* Inner loop condition. *)
        iApply "HΦ".
        subst Q. simpl.
        eauto 10 with iFrame.
      }
      wp_pures.

      assert (Hboundj : (int.nat j < length ents'')%nat).
      { apply Permutation_length in Hperm, Hperm0. word. }

      (* Read key of entry at [j - 1]. *)
      wp_load.
      destruct (list_lookup_lt _ (wrent_to_val <$> ents'') (int.nat (word.sub j (U64 1)))) as [entx Hlookupx].
      { rewrite fmap_length. word. }
      wp_apply (wp_SliceGet with "[$HentsS]"); first done.
      iIntros "[HentsS %Hty]".
      wp_pures.
      apply val_to_wrent_with_val_ty in Hty as (kx & vx & wx & tx & Hentx).
      subst entx.
      wp_pures.

      (* Read key of entry at [j]. *)
      wp_load.
      destruct (list_lookup_lt _ (wrent_to_val <$> ents'') (int.nat j)) as [enty Hlookupy].
      { rewrite fmap_length. word. }
      wp_apply (wp_SliceGet with "[$HentsS]"); first done.
      iIntros "[HentsS %Hty]".
      wp_pures.
      apply val_to_wrent_with_val_ty in Hty as (ky & vy & wy & ty & Henty).
      subst enty.
      wp_pures.

      wp_if_destruct.
      { (* Early return of inner loop. *)
        iApply "HΦ". subst Q. simpl.
        eauto with iFrame.
      }
      (* Swap entries at [j - 1] and at [j]. *)
      do 2 wp_load.
      wp_apply (wp_swap with "[$HentsS]").
      { iPureIntro. word. }

      iIntros (ents1) "[HentsS %Hperm1]".
      wp_load.
      wp_store.
      iApply "HΦ".
      subst Q. simpl.
      do 2 iExists _.
      iFrame.
      iPureIntro.
      split; last by word.
      by rewrite -Hperm0.
    }
    { (* Inner loop entry. *)
      subst Q. simpl.
      eauto with iFrame.
    }

    (* Increment [i]. *)
    iIntros "HQ".
    subst Q. simpl.
    iNamed "HQ".
    wp_load.
    wp_store.
    iApply "HΦ".
    subst P. simpl.
    do 2 iExists _. iFrame.
    iPureIntro.
    by rewrite <- Hperm.
  }
  { (* Outer loop entry. *)
    subst P. simpl.
    eauto with iFrame.
  }

  (* RETURN. *)
  iIntros "HP".
  subst P. simpl.
  iNamed "HP".
  iDestruct ("HentsC" with "HentsS") as "HentsS".
  wp_pures.
  iApply "HΦ".
  do 2 iExists _.
  iFrame.
  iPureIntro.
  split.
  { (* Prove [NoDup]. *)
    by rewrite Hperm.
  }
  { (* Prove [list_to_map]. *)
    rewrite Hmods.
    apply list_to_map_proper; last by rewrite Hperm.
    by apply NoDup_wrent_to_key_dbval.
  }
Qed.

End heap.


Require Import GHC.Base.
Require Import GHC.Num.
Require Import GHC.Real.
Require Import Coq.micromega.OrderedRing.
Require Import Proofs.GHC.Base.
Require Import OrdTactic.

(*Instances of GHC.Real do not necessarily need to be an ordered ring, but if they are, they
  should satify the following laws, which allows us to use OrderedRing tactics and theorems  *)

Class RealRing (b: Type) {Heq: Eq_ b} {Hord: Ord b} {Heqlaws: EqLaws b} {Hnum: Num b}
  {Hreal: Real b}  := {
  real_ring: @SOR b #0 #1 op_zp__ op_zt__ op_zm__ negate (fun x y => x == y = true)
    (fun x y => x <= y = true) (fun x y => x < y = true);
  rplus_eq: forall x y : b, (x == y) = true ->
    forall x0 y0 : b, (x0 == y0) = true -> (_+_ x x0 == _+_ y y0) = true;
  rtimes_eq: forall x y : b, (x == y) = true ->
    forall x0 y0 : b, (x0 == y0) = true -> (_*_ x x0 == _*_ y y0) = true;
  negate_eq: forall x y : b, (x == y) = true -> (negate x == negate y) = true;
  rle_eq: forall x y : b, (x == y) = true ->
    forall x0 y0 : b, (x0 == y0) = true -> (x <= x0) = true <-> (y <= y0) = true;
  rlt_eq: forall x y : b, (x == y) = true ->
    forall x0 y0 : b, (x0 == y0) = true -> (x < x0) = true <-> (y < y0) = true
}.


Module Imp.

Definition Int := nat.

Definition Id := nat.

Require Import Arith.

Definition beq_id := beq_nat.

Inductive AExp : Set :=
  | BaseInt : Int -> AExp
  | BaseId : Id -> AExp
  | Plus : AExp -> AExp -> AExp
  | Div : AExp -> AExp -> AExp
  | AExpHole : AExp.

Inductive BExp : Set :=
  | BaseBool : bool -> BExp
  | Lt : AExp -> AExp -> BExp
  | Not : BExp -> BExp
  | And : BExp -> BExp -> BExp
  | BExpHole : BExp.

Inductive Stmt : Set :=
  | Skip : Stmt
  | Assign : Id -> AExp -> Stmt
  | Seq : Stmt -> Stmt -> Stmt
  | IfThenElse : BExp -> Stmt -> Stmt -> Stmt
  | WhileDo : BExp -> Stmt -> Stmt
  | StmtHole : Stmt.

Require Export List.

Definition Ids := list Id.

Inductive Pgm : Set :=
  | Program : Ids -> Stmt -> Pgm.

Inductive K : Set :=
  | KAExp : AExp -> K
  | KBExp : BExp -> K
  | KStmt : Stmt -> K
  | KIds : Ids -> K
  | KPgm : Pgm -> K
  | KSeq : K -> K -> K
  | KNothing : K.

Inductive MapBaseIdBaseInt : Set :=
  | EmptyMap : MapBaseIdBaseInt
  | ConsMap : Id -> Int -> MapBaseIdBaseInt -> MapBaseIdBaseInt.

Definition Cfg : Set := prod K MapBaseIdBaseInt.

Require Import List.

Local Notation "[ ]" := nil : list_scope.
Local Notation "[ a ; .. ; b ]" := (a :: .. (b :: []) ..) : list_scope.

Definition iIds : Ids := [ 1 ; 2 ; 3 ].

Definition iPgm : Pgm := (Program iIds (Assign 1 (BaseInt 2))).

Definition iCfg : Cfg := (KSeq (KPgm iPgm) KNothing, EmptyMap).

Fixpoint isResultAExp (aexp : AExp) : bool :=
  match aexp with
    | BaseInt _ => true
    | _ => false
  end.

Fixpoint isResultBExp (bexp : BExp) : bool :=
  match bexp with
    | BaseBool _ => true
    | _ => false
  end.

Fixpoint heatAExp (aexp : AExp) : K :=
  match aexp with
    | BaseInt i => KAExp (BaseInt i)
    | BaseId i => KAExp (BaseId i)
    | Plus aexp1 aexp2 =>
      if isResultAExp aexp1 then
        if isResultAExp aexp2 then
          KAExp (Plus aexp1 aexp2)
        else
          KSeq (KAExp aexp2) (KAExp (Plus aexp1 AExpHole))
      else
        KSeq (KAExp aexp1) (KAExp (Plus AExpHole aexp2))
    | Div aexp1 aexp2 =>
      if isResultAExp aexp1 then
        if isResultAExp aexp2 then
          KAExp (Div aexp1 aexp2)
        else
          KSeq (KAExp aexp2) (KAExp (Div aexp1 AExpHole))
      else
        KSeq (KAExp aexp1) (KAExp (Div AExpHole aexp2))
    | AExpHole =>
      KAExp AExpHole
  end.

Fixpoint heatBExp (bexp : BExp) : K :=
  match bexp with
    | BaseBool b => KBExp (BaseBool b)
    | Lt aexp1 aexp2 =>
      if isResultAExp aexp1 then
        if isResultAExp aexp2 then
          KBExp (Lt aexp1 aexp2)
        else
          KSeq (KAExp aexp2) (KBExp (Lt aexp1 AExpHole))
      else
        KSeq (KAExp aexp1) (KBExp (Lt AExpHole aexp2))
    | Not bexp =>
      if isResultBExp bexp then
        KBExp (Not bexp)
      else
        KSeq (KBExp bexp) (KBExp (Not BExpHole))
    | And bexp1 bexp2 =>
      if isResultBExp bexp1 then
        KBExp (And bexp1 bexp2)
      else
        KSeq (KBExp bexp1) (KBExp (And BExpHole bexp2))
    | BExpHole =>
      KBExp BExpHole
  end.

Fixpoint heatStmt (stmt : Stmt) : K :=
  match stmt with
    | Skip => KStmt Skip
    | Assign id aexp =>
      if isResultAExp aexp then
        KStmt (Assign id aexp)
      else
        KSeq (KAExp aexp) (KStmt (Assign id AExpHole))
    | Seq stmt1 stmt2 =>
      KStmt (Seq stmt1 stmt2)
    | IfThenElse bexp stmt1 stmt2 =>
      if isResultBExp bexp then
        KStmt (IfThenElse bexp stmt1 stmt2)
      else
        KSeq (KBExp bexp) (KStmt (IfThenElse BExpHole stmt1 stmt2))
    | WhileDo bexp stmt =>
      KStmt (WhileDo bexp stmt)
    | StmtHole =>
      KStmt StmtHole
  end.

Definition heatIds (ids : Ids) : K := KIds ids.

Fixpoint isResultStmt (stmt : Stmt) : bool :=
  match stmt with
    _ => false
  end.

Fixpoint heatPgm (pgm : Pgm) : K :=
  match pgm with
    | Program ids stmt =>
      KPgm (Program ids stmt)
  end.

Fixpoint heat  (k : K) : K :=
  match k with
    | KAExp aexp => heatAExp aexp
    | KBExp bexp => heatBExp bexp
    | KStmt stmt => heatStmt stmt
    | KIds Ids => heatIds Ids
    | KPgm pgm => heatPgm pgm
    | KSeq k1 k2 => KSeq (heat k1) k2
    | KNothing => KNothing
  end.

Fixpoint heatCfg (c : Cfg) : Cfg :=
  match c with
    | (k, m) => (heat k, m)
  end.

Fixpoint lookup (m : MapBaseIdBaseInt) (x : Id) : option Int :=
  match m with
    | EmptyMap => None
    | ConsMap y i rest =>
      if beq_id x y then
        Some i
      else
        lookup rest x
  end.

Definition update (m : MapBaseIdBaseInt) (x : Id) (i : Int) : MapBaseIdBaseInt :=
  ConsMap x i m.

Fixpoint keys (m : MapBaseIdBaseInt) : list Id :=
  match m with
    | EmptyMap => []
    | ConsMap x _ rest => x :: (keys rest)
  end.

Definition rw_lookup (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KAExp (BaseId x)) rest, map) =>
      match lookup map x with
        | Some i => Some (KSeq (KAExp (BaseInt i)) rest, map)
        | None => None
      end
    | _ => None
  end.

Definition rw_plus (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KAExp (Plus aexp1 aexp2)) rest, map) =>
      match (aexp1, aexp2) with
        | (BaseInt i1, BaseInt i2) =>
          Some (KSeq (KAExp (BaseInt (i1 + i2))) rest, map)
        | _ =>
          None
      end
    | _ => None
  end.

Definition rw_div (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KAExp (Div aexp1 aexp2)) rest, map) =>
      match (aexp1, aexp2) with
        | (BaseInt i1, BaseInt i2) =>
          if beq_nat i2 0 then
            None
          else
            Some (KSeq (KAExp (BaseInt (i1 + i2))) rest, map)
        | _ =>
          None
      end
    | _ => None
  end.

Definition rw_lt (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KBExp (Lt aexp1 aexp2)) rest, map) =>
      match (aexp1, aexp2) with
        | (BaseInt i1, BaseInt i2) =>
          Some (KSeq (KBExp (BaseBool (leb i1 i2))) rest, map)
        | _ =>
          None
      end
    | _ => None
  end.

Definition rw_not (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KBExp (Not bexp)) rest, map) =>
      match bexp with
        | BaseBool b =>
          Some (KSeq (KBExp (BaseBool (negb b))) rest, map)
        | _ =>
          None
      end
    | _ => None
  end.

Definition rw_and1 (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KBExp (And bexp1 bexp2)) rest, map) =>
      match bexp1 with
        | BaseBool true =>
            Some (KSeq (KBExp bexp2) rest, map)
        | _ =>
          None
      end
    | _ => None
  end.

Definition rw_and2 (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KBExp (And bexp1 bexp2)) rest, map) =>
      match bexp1 with
        | BaseBool false =>
            Some (KSeq (KBExp (BaseBool false)) rest, map)
        | _ =>
          None
      end
    | _ => None
 end.

Definition rw_skip (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KStmt Skip) rest, map) =>
      Some (rest, map)
    | _ => None
  end.

Definition rw_assign (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KStmt (Assign id aexp)) rest, map) =>
      match aexp with
        | BaseInt i =>
          match lookup map id with
            | Some _ =>
              Some (rest, update map id i)
            | None =>
              None
          end
        | _ =>
          None
      end
    | _ =>
      None
  end.

Definition rw_sequential (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KStmt (Seq stmt1 stmt2)) rest, map) =>
      Some (KSeq (KSeq (KStmt stmt1) (KStmt stmt2)) rest, map)
    | _ =>
      None
  end.

Definition rw_if_true (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KStmt (IfThenElse (BaseBool true) stmt1 stmt2)) rest, map) =>
      Some (KSeq (KStmt stmt1) rest, map)
    | _ =>
      None
  end.

Definition rw_if_false (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KStmt (IfThenElse (BaseBool false) stmt1 stmt2)) rest, map) =>
      Some (KSeq (KStmt stmt2) rest, map)
    | _ =>
      None
  end.

Definition rw_while (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KStmt (WhileDo bexp stmt)) rest, map) =>
      Some (KSeq (KStmt (IfThenElse bexp (Seq stmt (WhileDo bexp stmt)) Skip)) rest, map)
    | _ =>
      None
  end.

Fixpoint contains l x : bool :=
  match l with
    | [] => false
    | y :: rest =>
      if beq_nat x y then
        true
      else
        contains rest x
  end.

Definition rw_program_still_vars (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KPgm (Program (x :: xs) stmt)) rest, map) =>
      if contains (keys map) x then
        None
      else
        Some (KSeq (KPgm (Program xs stmt)) rest, update map x 0)
    | _ => None
  end.

Definition rw_program_no_vars (c : Cfg) : option Cfg :=
  match c with
    | (KSeq (KPgm (Program [] stmt)) rest, map) =>
      Some (KSeq (KStmt stmt) rest, map)
    | _ => None
  end.

Eval simpl in heatCfg iCfg.

Definition rewrite_rules := [
  rw_program_still_vars;
  rw_program_no_vars;
  rw_lookup;
  rw_plus;
  rw_div;
  rw_lt;
  rw_not;
  rw_and1;
  rw_and2;
  rw_skip;
  rw_assign;
  rw_sequential;
  rw_if_true;
  rw_if_false;
  rw_while
  ].

Definition execute_rule (c : Cfg) rule : option Cfg :=
  rule c.

Fixpoint execute_one_rule (c : Cfg) rules : option Cfg :=
  match rules with
    | [] => None
    | rule :: rest =>
      match execute_rule c rule  with
        | Some c' => Some c'
        | None => execute_one_rule c rest
      end
  end.

Definition execute_one_step (c : Cfg) : option Cfg :=
  execute_one_rule (heatCfg c) rewrite_rules.

Fixpoint execute_many_steps (n : nat) (c : Cfg) : option Cfg :=
  match n with
    | 0 => Some c
    | S n' =>
      match execute_one_step c with
        | Some c' => execute_many_steps n' c'
        | None => None
      end
  end.

Definition hCfg := (heatCfg iCfg).
Eval compute in hCfg.

Definition hCfg1 := execute_one_step hCfg.
Eval compute in hCfg1.

Definition coerce (f : Cfg -> option Cfg) (c : option Cfg) : option Cfg :=
  match c with
    | None => None
    | Some c' => f c'
  end.

Definition hCfg2 := coerce execute_one_step hCfg1.
Eval compute in hCfg2.

Definition hCfg3 := coerce execute_one_step hCfg2.
Eval compute in hCfg3.

Definition hCfg4 := coerce execute_one_step hCfg3.
Eval compute in hCfg4.

Definition hCfg5 := coerce execute_one_step hCfg4.
Eval compute in hCfg5.

End Imp.

Extraction Imp.


From Coq Require Import Eqdep_dec. (* Needed to prove decidable equality on RegName *)
From Coq.micromega Require Import ZifyClasses.
From stdpp Require Import gmap fin_maps list.
From Coq Require Import ssreflect.
From cap_machine Require Import stdpp_extra.
From machine_utils Require Export finz.

(* No longer a coercion in Coq >= 8.14*)
Local Coercion Z.of_nat : nat >-> Z.

(* We assume a fixed set of registers, and a finite set of memory addresses.

   The exact size of the address space does not matter, it could be made a
   parameter of the machine.
*)

Definition RegNum: nat := 31.
Definition MemNum: Z := 2000000.
Global Opaque MemNum.

(* ---------------------------------- Registers ----------------------------------------*)

Inductive RegName: Type :=
| PC
| R (n: nat) (fin: n <=? RegNum = true).

Global Instance reg_eq_dec : EqDecision RegName.
Proof. intros r1 r2.  destruct r1,r2; [by left | by right | by right |].
       destruct (Nat.eq_dec n n0).
       + subst n0. left.
         assert (forall (b: bool) (n m: nat) (P1 P2: n <=? m = b), P1 = P2).
         { intros. apply eq_proofs_unicity.
           intros; destruct x; destruct y; auto. }
         rewrite (H _ _ _ fin fin0). reflexivity.
       + right. congruence.
Defined.

Lemma reg_eq_sym (r1 r2 : RegName) : r1 ≠ r2 → r2 ≠ r1. Proof. auto. Qed.

Program Definition n_to_regname (n : nat) : option RegName :=
  match Nat.le_dec n RegNum with left _ => Some (R n _) | right _ => None end.
Next Obligation.
  intros. eapply Nat.leb_le; eauto.
Defined.

Global Instance reg_countable : Countable RegName.
Proof.
  refine {| encode r := encode match r with
                               | PC => inl ()
                               | R n fin => inr n
                               end ;
            decode n := match (decode n) with
                        | Some (inl ()) => Some PC
                        | Some (inr n) => n_to_regname n
                        | None => None
                        end ;
            decode_encode := _ |}.
  intro r. destruct r; auto.
  rewrite decode_encode.
  unfold n_to_regname.
  destruct (Nat.le_dec n RegNum).
  - do 2 f_equal. apply eq_proofs_unicity. decide equality.
  - exfalso. by apply (Nat.leb_le n RegNum) in fin.
Defined.

(* Instances for [zify]: make [lia] work on registers *)
(* TODO: separate the proof parts into lemmas *)

Definition Z_of_regname (r: RegName): Z.
  destruct r. exact 0.
  exact (S n).
Defined.

Instance RegName_InjTyp : InjTyp RegName Z.
  refine (mkinj _ _ Z_of_regname (fun n => n <= RegNum + 1)%Z _).
  intros [|]. cbn. lia. cbn. apply Nat.leb_le in fin. lia.
Defined.
Add Zify InjTyp RegName_InjTyp.

Instance Op_RegName_eq : BinRel (@eq RegName).
  refine ({| TR := @eq Z; TRInj := _ |}).
  cbn. intros r1 r2. split.
  - intros ->; eauto.
  - destruct r1; destruct r2; eauto; cbn; try apply Nat.leb_le in fin; try lia.
    intros ->%Nat2Z.inj%eq_add_S.
    f_equal. apply eq_proofs_unicity. intros [|] [|]; eauto.
Defined.
Add Zify BinRel Op_RegName_eq.

(* Names for registers *)
Definition r_t0 : RegName := R 0 eq_refl.
Definition r_t1 : RegName := R 1 eq_refl.
Definition r_t2 : RegName := R 2 eq_refl.
Definition r_t3 : RegName := R 3 eq_refl.
Definition r_t4 : RegName := R 4 eq_refl.
Definition r_t5 : RegName := R 5 eq_refl.
Definition r_t6 : RegName := R 6 eq_refl.
Definition r_t7 : RegName := R 7 eq_refl.
Definition r_t8 : RegName := R 8 eq_refl.
Definition r_t9 : RegName := R 9 eq_refl.
Definition r_t10 : RegName := R 10 eq_refl.
Definition r_t11 : RegName := R 11 eq_refl.
Definition r_t12 : RegName := R 12 eq_refl.
Definition r_t13 : RegName := R 13 eq_refl.
Definition r_t14 : RegName := R 14 eq_refl.
Definition r_t15 : RegName := R 15 eq_refl.
Definition r_t16 : RegName := R 16 eq_refl.
Definition r_t17 : RegName := R 17 eq_refl.
Definition r_t18 : RegName := R 18 eq_refl.
Definition r_t19 : RegName := R 19 eq_refl.
Definition r_t20 : RegName := R 20 eq_refl.
Definition r_t21 : RegName := R 21 eq_refl.
Definition r_t22 : RegName := R 22 eq_refl.
Definition r_t23 : RegName := R 23 eq_refl.
Definition r_t24 : RegName := R 24 eq_refl.
Definition r_t25 : RegName := R 25 eq_refl.
Definition r_t26 : RegName := R 26 eq_refl.
Definition r_t27 : RegName := R 27 eq_refl.
Definition r_t28 : RegName := R 28 eq_refl.
Definition r_t29 : RegName := R 29 eq_refl.
Definition r_t30 : RegName := R 30 eq_refl.
Definition r_t31 : RegName := R 31 eq_refl.

(* A list of all general purpuse registers (if regnum=31) *)
Definition all_registers : list RegName :=
  [r_t0;r_t1;r_t2;r_t3;r_t4;r_t5;r_t6;r_t7;r_t8;r_t9;r_t10;r_t11;r_t12;r_t13;
     r_t14;r_t15;r_t16;r_t17;r_t18;r_t19;r_t20;r_t21;r_t22;r_t23;r_t24;r_t25;r_t26;
       r_t27;r_t28;r_t29;r_t30;r_t31;PC].

(* Set of all registers *)
Definition all_registers_s : gset RegName := list_to_set all_registers.

Lemma all_registers_NoDup :
  NoDup all_registers.
Proof.
  unfold all_registers.
  repeat (
    apply NoDup_cons_2;
    first (repeat (rewrite not_elem_of_cons; split; [done|]); apply not_elem_of_nil)
  ).
  by apply NoDup_nil.
Qed.

(* Spec for all_registers *)

Lemma all_registers_correct r1 :
  r1 ∈ all_registers.
Proof.
  rewrite /all_registers.
  destruct r1.
  - do 32 (apply elem_of_cons; right).
      by apply elem_of_list_singleton.
  - induction n.
    + apply elem_of_cons; left.
      apply f_equal. apply eq_proofs_unicity. decide equality.
    + apply elem_of_list_lookup_2 with (S n).
      repeat (destruct n;
                first (simpl;do 2 f_equal;apply eq_proofs_unicity;decide equality)).
      simpl in *. inversion fin.
Qed.

Lemma all_registers_s_correct r:
  r ∈ all_registers_s.
Proof.
  rewrite /all_registers_s elem_of_list_to_set.
  apply all_registers_correct.
Qed.

Lemma all_registers_correct_sub r : NoDup r → r ⊆+ all_registers.
Proof.
  intros Hdup.
  apply NoDup_submseteq;auto. intros r' Hin.
  apply all_registers_correct.
Qed.

Instance setunfold_all_regs:
  forall x, SetUnfoldElemOf x all_registers_s True.
Proof.
  intros. constructor. split; auto.
  intro. eapply all_registers_s_correct.
Qed.

Lemma all_registers_union_l s :
  s ∪ all_registers_s = all_registers_s.
Proof.
  apply (anti_symm subseteq). 2: set_solver.
  rewrite elem_of_subseteq. intros ? _.
  apply all_registers_s_correct.
Qed.

Lemma all_registers_union_r s :
  all_registers_s ∪ s = all_registers_s.
Proof. rewrite union_comm_L. apply all_registers_union_l. Qed.

Lemma all_registers_subseteq s :
  s ⊆ all_registers_s.
Proof.
  rewrite elem_of_subseteq. intros ? _. apply all_registers_s_correct.
Qed.

Lemma regmap_full_dom {A} (r: gmap RegName A):
  (∀ x, is_Some (r !! x)) →
  dom r = all_registers_s.
Proof.
  intros Hfull. apply (anti_symm subseteq); rewrite elem_of_subseteq.
  - intros rr _. apply all_registers_s_correct.
  - intros rr _. rewrite elem_of_dom. apply Hfull.
Qed.

(* -------------------------------- Memory addresses -----------------------------------*)

Notation Addr := (finz MemNum).
Declare Scope Addr_scope.
Delimit Scope Addr_scope with a.

Notation "a1 <= a2 < a3" := (@finz.le_lt MemNum a1 a2 a3) : Addr_scope.
Notation "a1 <= a2" := (@finz.le MemNum a1 a2) : Addr_scope.
Notation "a1 <=? a2" := (@finz.leb MemNum a1 a2) : Addr_scope.
Notation "a1 < a2" := (@finz.lt MemNum a1 a2) : Addr_scope.
Notation "a1 <? a2" := (@finz.ltb MemNum a1 a2) : Addr_scope.
Notation "a1 + z" := (@finz.incr MemNum a1 z) : Addr_scope.
Notation "a ^+ off" := (@finz.incr_default MemNum a off) (at level 50) : Addr_scope.

Notation z_to_addr := (@finz.of_z MemNum).
Notation z_of := (@finz.to_z MemNum).

Notation za := (@finz.FinZ MemNum 0%Z eq_refl eq_refl).
Notation top := (finz.largest za : Addr).
Notation "0" := (za) : Addr_scope.

Notation eqb_addr := (λ (a1 a2: Addr), Z.eqb a1 a2).
Notation "a1 =? a2" := (eqb_addr a1 a2) : Addr_scope.

Notation addr_incr_eq := (finz_incr_eq).

Global Open Scope general_if_scope.
(* ---------------------------------- OTypes ----------------------------------------*)

(* Number of otypes in our system *)
Definition ONum: nat := 1000.
Global Opaque ONum.
Notation OType := (finz ONum).
Declare Scope OType_scope.
Delimit Scope OType_scope with ot.

Notation "a1 <= a2 < a3" := (@finz.le_lt ONum a1 a2 a3) : OType_scope.
Notation "a1 <= a2" := (@finz.le ONum a1 a2) : OType_scope.
Notation "a1 <=? a2" := (@finz.leb ONum a1 a2) : OType_scope.
Notation "a1 < a2" := (@finz.lt ONum a1 a2) : OType_scope.
Notation "a1 <? a2" := (@finz.ltb ONum a1 a2) : OType_scope.
Notation "a1 + z" := (@finz.incr ONum a1 z) : OType_scope.
Notation "a ^+ off" := (@finz.incr_default ONum a off) (at level 50) : OType_scope.

Notation z_to_otype := (@finz.of_z ONum).
Notation z_of_ot := (@finz.to_z ONum).

Notation za_ot := (@finz.FinZ ONum 0%Z eq_refl eq_refl).
Notation top_ot := (finz.largest za : OType).
Notation "0" := (za_ot) : OType_scope.

Notation eqb_otype := (λ (a1 a2: OType), Z.eqb a1 a2).
Notation "a1 =? a2" := (eqb_otype a1 a2) : OType_scope.

Notation otype_incr_eq := (finz_incr_eq).


Require Import Setoid.
Require Import Bool.
Require Import List.
Require Import ZArith.
Require Import Nat.
(*Load our files*)
Require Import Tools.
Require Import ListExtras.
Require Import Zvec.
Require Import SimpleLieAlgebras.
Require Import Data.
Require Import TypeAll.
Open Scope nat_scope.

Theorem thm_main_B :
  forall g,
    lie_algebra_type g = lie_B_type
    -> forall lambda,
         is_radical g lambda = true
         -> (Is_nonmixed g lambda)
            \/ (Is_mixed g lambda).
Admitted.
Theorem thm_main_D :
  forall g,
    lie_algebra_type g = lie_D_type
    -> forall lambda,
         is_radical g lambda = true
         -> (Is_nonmixed g lambda)
            \/ (Is_mixed g lambda).
Admitted.


Require Import Crypto.Arithmetic.PrimeFieldTheorems.
Require Import Crypto.Specific.montgomery32_2e256m2e224p2e192p2e96m1_8limbs.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.


From stars Require Import definitions.

Inductive re (X : Type) :=
  | RE_None
  | RE_Empty
  | RE_Literal (x : X)
  | RE_Or (a b : re X)
  | RE_Seq (a b : re X)
  | RE_Star (a : re X).

Arguments RE_None {_}.
Arguments RE_Empty {_}.
Arguments RE_Literal {_}.
Arguments RE_Or {_}.
Arguments RE_Seq {_}.
Arguments RE_Star {_}.

Global Instance : FMap re := λ X Y f, fix go (a : re X) :=
  match a with
  | RE_None => RE_None
  | RE_Empty => RE_Empty
  | RE_Literal x => RE_Literal (f x)
  | RE_Or b c => RE_Or (go b) (go c)
  | RE_Seq b c => RE_Seq (go b) (go c)
  | RE_Star b => RE_Star (go b)
  end.

Notation "a ∣ b" := (RE_Or a b)
  (left associativity, at level 53, format "a ∣ b").
Notation "a ⋅ b" := (RE_Seq a b)
  (left associativity, at level 52, format "a ⋅ b").
Notation "a ∗" := (RE_Star a)
  (left associativity, at level 51, format "a ∗").

Section Regular_Expressions.

Variable X : Type.
Notation re := (re X).

Fixpoint re_in (w : list X) (a : re) :=
  match a with
  | RE_None => False
  | RE_Empty => w = []
  | RE_Literal x => w = [x]
  | RE_Or b c => re_in w b \/ re_in w c
  | RE_Seq b c => ∃ u v, w = u ++ v /\ re_in u b /\ re_in v c
  | RE_Star b => w = [] \/ ∃ vs, w = concat vs /\ Forall (λ v, re_in v b) vs
  end.

Global Instance : Zero re := RE_None.
Global Instance : One re := RE_Empty.
Global Instance : Empty re := RE_None.
Global Instance : ElemOf (list X) re := re_in.
Global Instance : Equiv re := λ a b, ∀ w, w ∈ a <-> w ∈ b.

Global Instance : Add re :=
  λ a b, match a, b with
  | RE_None, _ => b
  | _, RE_None => a
  | RE_Empty, RE_Empty => RE_Empty
  | RE_Empty, RE_Star _ => b
  | RE_Star _, RE_Empty => a
  | _, _ => RE_Or a b
  end.

Global Instance : Mul re :=
  λ a b, match a with
  | RE_None => RE_None
  | RE_Empty => b
  | _ => match b with
    | RE_None => RE_None
    | RE_Empty => a
    | _ => RE_Seq a b
    end
  end.

Global Instance : Star re :=
  λ a, match a with
  | RE_None => RE_Empty
  | RE_Empty => RE_Empty
  | RE_Star _ => a
  | _ => RE_Star a
  end.

Global Instance : @Equivalence re (≡).
Proof. firstorder. Qed.

Section Simplifications.

Lemma re_seq_1_l a :
  1⋅a ≡ a.
Proof.
intros w; cbn; split.
intros (u&v&->&->&Hv); done.
intros Hw; exists [], w; done.
Qed.

Lemma re_seq_1_r a :
  a⋅1 ≡ a.
Proof.
intros w; cbn; split.
intros (u&v&->&Hu&->); simplify_list_eq; done.
intros Hw; exists w, []; simplify_list_eq; done.
Qed.

Lemma re_star_0 :
  0∗ ≡ 1.
Proof.
intros w; cbn; firstorder; subst.
destruct x; decompose_Forall; done.
Qed.

Lemma re_star_1 :
  1∗ ≡ 1.
Proof.
intros w; cbn; firstorder; subst.
apply concat_nil_Forall; done.
Qed.

Lemma re_star_star a :
  a∗∗ ≡ a∗.
Proof.
intros w; cbn; split; intros [H|(vs&->&H)]; auto.
- induction vs as [|v vs]. left; done. decompose_Forall.
  apply IHvs in H0 as [->|(vs'&->&Hvs')]; clear IHvs.
  + simplify_list_eq; firstorder.
  + right; destruct H as [->|(us&->&Hus)]; cbn.
    exists vs'; done. exists (us ++ vs'); split.
    symmetry; apply concat_app. decompose_Forall; done.
- right; exists vs; split; [done|]. eapply Forall_impl.
  apply H. cbn; intros v Hv; right; exists [v]; split.
  cbn; simplify_list_eq; done. decompose_Forall; done.
Qed.

Theorem equiv_re_add a b :
  a + b ≡ a∣b.
Proof.
destruct a, b; cbn; try firstorder.
Qed.

Theorem equiv_re_mul a b :
  a * b ≡ a⋅b.
Proof.
destruct a, b; cbn; try reflexivity; symmetry.
all: try apply re_seq_1_l; try apply re_seq_1_r.
all: firstorder.
Qed.

Theorem equiv_re_star a :
  a{*} ≡ a∗.
Proof.
destruct a; cbn; try reflexivity; symmetry.
apply re_star_0. apply re_star_1. apply re_star_star.
Qed.

End Simplifications.

Section Properties.

Lemma assoc_re_seq a b c :
  a⋅(b⋅c) ≡ a⋅b⋅c.
Proof.
intros w; cbn; split; intros (u&v&->&H).
+ destruct H as (Ha&v1&v2&->&Hv1&Hv2).
  exists (u ++ v1), v2; repeat split; try done.
  apply app_assoc. exists u, v1; done.
+ destruct H as ((u1&u2&->&Hu1&Hu2)&Hv).
  exists u1, (u2 ++ v); repeat split; try done.
  symmetry; apply app_assoc. exists u2, v; done.
Qed.

Lemma left_expand_re_star a :
  a∗ ≡ 1∣a⋅a∗.
Proof.
intros w; cbn; split; intros H.
+ destruct H as [->|(vs&->&H)]; [left; done|].
  destruct vs as [|v vs]; [left; done|right]; exists v, (concat vs).
  decompose_Forall; repeat split; try done. right; exists vs; done.
+ destruct H as [->|(u&v&->&Hu&H)]; [left; done|right].
  destruct H as [->|(vs&->&Hvs)]; [exists [u]|exists (u :: vs)];
  cbn; repeat split; decompose_Forall; done.
Qed.

Lemma right_expand_re_star a :
  a∗ ≡ 1∣a∗⋅a.
Proof.
intros w; cbn; split; intros H.
+ destruct H as [->|(vs&->&H)]; [left; done|].
  destruct vs as [|v vs _] using rev_ind; [left; done|right].
  decompose_Forall; exists (concat vs), v; repeat split.
  rewrite concat_app; simplify_list_eq; done.
  right; exists vs; done. done.
+ destruct H as [->|(u&v&->&H&Hv)]; [left; done|right].
  destruct H as [->|(us&->&Hus)]; [exists [v]|exists (us ++ [v])];
  cbn; repeat split; simplify_list_eq; decompose_Forall; try done.
  rewrite concat_app; simplify_list_eq; done.
Qed.

Lemma left_intro_re_star a b :
  a⋅b∣b ≡ b -> a∗⋅b∣b ≡ b.
Proof.
intros H w; cbn; split; [|auto].
intros [(u&v&->&[->|(us&->&Hus)]&Hv)|Hw]; simplify_list_eq; try done.
induction us as [|u us]; simplify_list_eq. done. decompose_Forall.
apply H; cbn; left; eexists; eexists; auto.
Qed.

Lemma right_intro_re_star a b :
  b⋅a∣b ≡ b -> b⋅a∗∣b ≡ b.
Proof.
intros H w; cbn; split; [|auto].
intros [(u&v&->&Hu&[->|(vs&->&Hvs)])|Hw]; simplify_list_eq; try done.
induction vs as [|v vs] using rev_ind; simplify_list_eq. done.
rewrite concat_app; simplify_list_eq; rewrite app_assoc.
decompose_Forall; apply H; cbn; left; eexists; eexists; auto.
Qed.

End Properties.

Section Rewriting.

Global Instance : Proper ((≡) ==> (≡) ==> (≡)) RE_Or.
Proof. firstorder. Qed.

Global Instance : Proper ((≡) ==> (≡) ==> (≡)) RE_Seq.
Proof. firstorder. Qed.

Global Instance : Proper ((≡) ==> (≡)) RE_Star.
Proof.
intros a b Hab w; cbn; split; intros [->|(vs&->&H)]; try (left; done); right;
exists vs; split; try done; eapply Forall_impl; try apply H; apply Hab.
Qed.

Global Instance : Proper ((≡) ==> (≡) ==> (≡)) (@add re _).
Proof. intros a b Hab c d Hcd; rewrite ?equiv_re_add, Hab, Hcd; done. Qed.

Global Instance : Proper ((≡) ==> (≡) ==> (≡)) (@mul re _).
Proof. intros a b Hab c d Hcd; rewrite ?equiv_re_mul, Hab, Hcd; done. Qed.

Global Instance : Proper ((≡) ==> (≡)) (@star re _).
Proof. intros a b Hab; rewrite ?equiv_re_star, Hab; done. Qed.

End Rewriting.

Local Ltac expand := rewrite ?equiv_re_add, ?equiv_re_mul, ?equiv_re_star.

Global Instance : Kleene_Algebra re.
Proof.
split. split. split. c. split. 1,3: split. 1,4: split. 1,3: c.
1,3-4,7-11,14: intro; intros; expand; firstorder.
- intros a b c; expand; apply assoc_re_seq.
- intros a; expand; apply re_seq_1_l.
- intros a; expand; apply re_seq_1_r.
- intros a; expand; apply left_expand_re_star.
- intros a; expand; apply right_expand_re_star.
- intros a b; expand; apply left_intro_re_star.
- intros a b; expand; apply right_intro_re_star.
Qed.

Section Evaluation.

Context `{Kleene_Algebra X}.

Fixpoint re_eval (a : re) : X :=
  match a with
  | RE_None => 0
  | RE_Empty => 1
  | RE_Literal x => x
  | RE_Or b c => re_eval b + re_eval c
  | RE_Seq b c => re_eval b * re_eval c
  | RE_Star b => (re_eval b){*}
  end.

End Evaluation.

End Regular_Expressions.


Require Import Coq.Classes.Morphisms.
Require Import Coq.Classes.RelationClasses.

Require Import Domains.Continuous Domains.Domain Domains.Order.

(** Embedding-projection pairs. These are the morphisms in the
    category Cpo^ep. We solve recursive domain equations can by taking
    suprema of ω-chains in this category ("EPChains"). *)
Section epPair.
  Record EPPair (A B : Type) `{CPO A} `{CPO B} : Type :=
    { embed : A -c> B
    ; proj  : B -c> A
    ; embed_proj_id : proj ∘ embed =o= id_cfun A
    ; proj_embed_le : embed ∘ proj <o= id_cfun B
    }.
End epPair.

Arguments embed {_ _ _ _ _ _ _ _}.
Arguments proj {_ _ _ _ _ _ _ _}.

Notation "A '-*>' B" := (EPPair A B) (right associativity, at level 99).

(* Ordering on EPPairs. *)
Instance OTEPPair {A B : Type} `{cpoA : CPO A} `{cpoB : CPO B}
  : OType (A -*> B) :=
  {| oleq := fun ep1 ep2 => embed ep1 <o= embed ep2 /\ proj ep1 <o= proj ep2 |}.
Proof.
  constructor.
  - intros ep. split; intros x y H0; destruct ep as [[] [] ? ?]; auto.
  - intros ep1 ep2 ep3 [H0 H1] [H2 H3]; split.
    + etransitivity; eauto; apply H2; reflexivity.
    + etransitivity; eauto; apply H3; reflexivity.
Defined.

(* The identity EPPair. *)
Program Definition id_ep (A : Type) `{CPO A} : A -*> A :=
  {| embed := id_cfun A; proj := id_cfun A |}.
Next Obligation. rewrite id_cfun_l; reflexivity. Qed.

(* Composition of EPPairs as a meta operation (not a continuous
   function). *)
Program Definition compose_ep
        {A B C : Type} `{cpoA : CPO A} `{cpoB : CPO B}
        `{cpoC : CPO C} (ep1 : A -*> B) (ep2 : B -*> C) : A -*> C :=
  {| embed := embed ep2 ∘ embed ep1
   ; proj  := proj ep1 ∘ proj ep2 |}.
Next Obligation.
  rewrite compose_cfun_assoc.
  rewrite <- (compose_cfun_assoc (embed ep2)).
  destruct ep2. simpl. rewrite embed_proj_id0.
  generalize id_cfun_r. intros H0.
  rewrite id_cfun_r; destruct ep1; auto.
Qed.
Next Obligation.
  cut ((embed ep2 ∘ embed ep1) ∘ (proj ep1 ∘ proj ep2) <o= id_cfun C).
  { firstorder. }
  rewrite compose_cfun_assoc.
  rewrite <- (compose_cfun_assoc (proj ep1)).
  destruct ep1. rewrite proj_embed_le0.
  rewrite id_cfun_r. destruct ep2; auto.
Qed.

Notation "f '∘∘' g" := (compose_ep g f) (at level 65) : cfun_scope.
Open Scope cfun_scope.

(* Program Definition postcompose_ep {A B C} `{CPO A} `{CPO B} `{CPO C} *)
(*         (g : B -*> C) : (A -c> B) -*> A -c> C := *)
(*   {| embed := postcompose_cfun (embed g) *)
(*    ; proj  := postcompose_cfun (proj g) |}. *)
(* Next Obligation. *)
(*   split; intros f h Hleq1 x y Hleq2; destruct g as [e p [Hep1 Hep2] Hpe]; *)
(*     simpl in *; try apply Hep1; try apply Hep2; apply Hleq1; apply Hleq2. *)
(* Qed. *)
(* Next Obligation. destruct g as [e p ? Hpe]; apply Hpe; auto. Qed. *)

Lemma id_ep_l {A B} `{CPO A} `{CPO B} (f : A -*> B) :
  id_ep B ∘∘ f =o= f.
Proof. split; split; destruct f as [[] []]; firstorder. Qed.

Lemma id_ep_r {A B} `{CPO A} `{CPO B} (f : A -*> B) :
  f ∘∘ id_ep A =o= f.
Proof. split; split; destruct f as [[] []]; firstorder. Qed.

Program Definition pair_ep {A B C D} `{CPO A} `{CPO B} `{CPO C} `{CPO D}
        (f : A -*> C) (g : B -*> D) : (A*B) -*> (C*D) :=
  {| embed := pair_cfun (embed f) (embed g)
   ; proj  := pair_cfun (proj f) (proj g) |}.
Next Obligation.
  split; split; destruct a1, a2, f as [[] []], g as [[] []];
    simpl; firstorder.
Qed.
Next Obligation.
  destruct f as [[] []], g as [[] []]; simpl; firstorder.
Qed.

(* A bit annoying to get it to use the correct types here. *)
Instance Proper_pair_ep {A B C D} `{cpoA : CPO A} `{cpoB : CPO B}
         `{cpoC : CPO C} `{cpoD : CPO D} :
  Proper (@oeq _ (@OTEPPair _ _ _ _ cpoA _ _ cpoC) ==>
               @oeq  _ (@OTEPPair _ _ _ _ cpoB _ _ cpoD) ==>
               oeq) pair_ep.
Proof.
  intros h k Heq. intros h' k' Heq'.
  split; split; intros [x1 x2] [y1 y2] [Hleq1 Hleq2]; split;
    destruct h, k, h', k', Heq as [Heq1 Heq2], Heq' as [Heq3 Heq4];
    simpl in *; try apply Heq1; try apply Heq2;
      try apply Heq3; try apply Heq4; auto.
Qed.

Lemma pair_ep_id {A B} `{cpoA : CPO A} `{cpoB : CPO B} :
  pair_ep (@id_ep _ _ _ cpoA) (@id_ep _ _ _ cpoB) =o= id_ep _.
Proof. split; split; intros [] []; firstorder. Qed.

Lemma pair_ep_compose
      {A B C D E F} `{CPO A} `{CPO B} `{CPO C} `{CPO D} `{CPO E} `{CPO F}
      (f : A -*> C) (g : B -*> D) (h : C -*> E) (k : D -*> F) :
  pair_ep (h ∘∘ f) (k ∘∘ g) =o= pair_ep h k ∘∘ pair_ep f g.
Proof.
  split; split; intros [] [] ?; split;
    destruct f as [[] []], g as [[] []], h as [[] []], k as [[] []];
    simpl; firstorder.
Qed.

Program Definition hom_ep {A B C} `{CPO A} `{CPO B} `{CPO C}
        (f : B -*> C) : (A -> B) -*> (A -> C) :=
  {| embed := postcompose_cfun' (embed f)
   ; proj  := postcompose_cfun' (proj f) |}.
Next Obligation. apply postcompose_cfun'_iso; destruct f; auto. Qed.
Next Obligation. destruct f; apply proj_embed_le0; auto. Qed.

Program Definition chom_ep {A B C D} `{CPO A} `{CPO B} `{CPO C} `{CPO D}
        (f : A -*> B) (g : C -*> D) : (A -c> C) -*> (B -c> D) :=
  {| embed := {| cfun_app := fun h => embed g ∘ h ∘ proj f |}
   ; proj  := {| cfun_app := fun k => proj g ∘ k ∘ embed f |} |}.
Next Obligation.
  intros h k Hleq1 x y Hleq2.
  destruct f, g. simpl.
  destruct embed1, h, proj0, k. simpl in *.
  apply cfun_Proper. apply Hleq1. rewrite Hleq2. reflexivity.
Qed.
Next Obligation.
  destruct f, g; split; simpl; intros x y Hleq; destruct embed1;
    rewrite cfun_Continuous; destruct H2; apply supremumProper;
      intros n; apply cfun_Proper; destruct A0; simpl; destruct (x0 n);
        apply cfun_Proper0; destruct proj0; rewrite Hleq; reflexivity.
Qed.
Next Obligation.
  intros h k Hleq1 x y Hleq2.
  destruct f, g. simpl.
  destruct embed0, h, proj1, k. simpl in *.
  apply cfun_Proper1. apply Hleq1. rewrite Hleq2. reflexivity.
Qed.
Next Obligation.
  destruct f, g; split; simpl; intros x y Hleq; destruct proj1;
    rewrite cfun_Continuous; destruct H1; apply supremumProper;
      intros n; apply cfun_Proper; destruct A0; simpl; destruct (x0 n);
        apply cfun_Proper0; destruct embed0; rewrite Hleq; reflexivity.
Qed.
Next Obligation.
  split; simpl; intros h k Hleq1 x y Hleq2; destruct f, g;
    apply embed_proj_id1; apply Hleq1; apply embed_proj_id0; apply Hleq2.
Qed.
Next Obligation.
  destruct f, g; apply proj_embed_le1; apply H3;
    apply proj_embed_le0; apply H4.
Qed.

Lemma chom_ep_id {A B} `{CPO A} `{CPO B} :
  chom_ep (id_ep A) (id_ep B) =o= id_ep _.
Proof. split; split; simpl; auto. Qed.

Lemma chom_ep_comp {A B C D E F} `{CPO A} `{CPO B} `{CPO C}
      `{CPO D} `{CPO E} `{CPO F}
      (f : A -*> B) (g : B -*> C)
      (h : D -*> E) (k : E -*> F) :
  chom_ep (g ∘∘ f) (k ∘∘ h) =o= chom_ep g k ∘∘ chom_ep f h.
Proof.
  split; split; simpl; intros f' g' Hleq1 x y Hleq2;
    destruct f as [[fe feProp ?] [fp fpProp ?] fep fpe];
    destruct g as [[ge geProp ?] [gp gpProp ?] gep gpe];
    destruct h as [[he heProp ?] [hp hpProp ?] hep hpe];
    destruct k as [[ke keProp ?] [kp kpProp ?] kep kpe];
    try apply keProp; try apply heProp; try apply hpProp;
      try apply kpProp; try apply keProp; try apply heProp;
        apply Hleq1; rewrite Hleq2; reflexivity.
Qed.

Lemma hom_ep_id {A B} `{CPO A} `{CPO B} :
  hom_ep (id_ep A) =o= id_ep _.
Proof. split; split; simpl; auto. Qed.

Lemma hom_ep_comp {T A B C} `{CPO T} `{CPO A} `{CPO B} `{CPO C}
      (f : A -*> B) (g : B -*> C) :
  @hom_ep T A C _ _ _ _ _ _ _ _ _ (g ∘∘ f) =o=
  hom_ep g ∘∘ hom_ep f.
Proof.
  split; split; simpl; intros h k Hleq x;
    destruct g as [[ge geProp ?] [gp gpProp ?] gep gpe];
    destruct f as [[fe feProp ?] [fp fpProp ?] fep fpe];
    try apply geProp; try apply feProp; try apply fpProp;
      try apply gpProp; apply Hleq.
Qed.

Program Definition tuple_bot_ep {A B} `{CPO A} `{CPO B}
        (f : unit -*> A) (g : unit -*> B) : unit -*> (A*B) :=
  {| embed := tuple_cfun (embed f) (embed g)
   ; proj  := const_cfun tt |}.
Next Obligation. rewrite const_cfun_compose; firstorder. Qed.
Next Obligation.
  split.
  - destruct f as [e p Hep Hpe]; simpl in *; specialize (Hpe a0 a H1);
      destruct (p @c@ a0); apply Hpe.
  - destruct g as [e p Hep Hpe]; simpl in *; specialize (Hpe b0 b H2);
      destruct (p @c@ b0); apply Hpe.
Qed.

Program Definition chom_bot_ep {A B} `{CPO A} `{CPO B}
        (ep : unit -*> B) : unit -*> (A -c> B) :=
  {| embed := const_cfun (const_cfun (embed ep @c@ tt))
   ; proj  := const_cfun tt |}.
Next Obligation. split; simpl; auto. Qed.
Next Obligation.
  destruct ep; simpl in *.
  assert (Hleq: a2 @c@ a3 <o= a2 @c@ a3) by reflexivity.
  specialize (proj_embed_le0 (a2 @c@ a3) (a2 @c@ a3) Hleq).
  destruct (proj0 @c@ (a2 @c@ a3)); apply proj_embed_le0.
Qed.

Program Definition hom_bot_ep {A B} `{CPO A} `{CPO B}
        (ep : unit -*> B) : unit -*> (A -> B) :=
  {| embed := const_cfun (fun _ => embed ep @c@ tt)
   ; proj  := const_cfun tt |}.
Next Obligation. rewrite const_cfun_compose; split; simpl; auto. Qed.
Next Obligation.
  destruct ep; simpl in *.
  assert (Hleq: a2 x <o= a2 x) by reflexivity.
  specialize (proj_embed_le0 _ _ Hleq); destruct (proj0 @c@ a2 x).
  apply proj_embed_le0.
Qed.

Lemma embed_inj {A B} `{CPO A} `{CPO B} (ep : A -*> B) x y :
  embed ep @c@ x =o= embed ep @c@ y ->
  x =o= y.
Proof.
  simpl.
  intros H1.
  destruct ep. simpl in *.
  assert (proj0 @c@ (embed0 @c@ x) =o= proj0 @c@ (embed0 @c@ y)).
  { destruct proj0. simpl. rewrite H1. reflexivity. }
  rewrite 2!oeq_id_pointwise in H2; auto.
Qed.

Section epChain.
  Class OTypeSequence (f : nat -> Type) : Type :=
    oTypeSeq :> forall n, OType (f n).
  Class SupremumSequence (f : nat -> Type) {g : OTypeSequence f} : Type :=
    supremumSeq :> forall n, Supremum (f n).
  Class CPOSequence (f : nat -> Type) {g : OTypeSequence f}
        {h : SupremumSequence f} : Type :=
    cpoSeq :> forall n, CPO (f n).
  Class EPChain (f : nat -> Type) {g : OTypeSequence f}
        {h : SupremumSequence f} {i : CPOSequence f} : Type :=
    chain :> forall (n : nat), f n -*> f (S n).
End epChain.

Section constEPChain.
  Context (T : Type) {oT : OType T} {supT : Supremum T} {cpoT : CPO T}.
  Definition ConstTypeSequence : nat -> Type := fun _ => T.
  Global Instance ConstOTypeSequence : OTypeSequence ConstTypeSequence :=
    fun _ => oT.
  Global Instance ConstSupremumSequence : SupremumSequence ConstTypeSequence :=
    fun _ => supT.
  Global Instance ConstCPOSequence : CPOSequence ConstTypeSequence :=
    fun _ => cpoT.
  Global Instance ConstEPChain : EPChain ConstTypeSequence :=
    fun n => id_ep T.
End constEPChain.

(** Product EP chain. *)
Definition ProductTypeSequence (f1 f2 : nat -> Type) : nat -> Type :=
  fun n => prod (f1 n) (f2 n).

Instance ProductOTypeSequence (f1 f2 : nat -> Type)
         {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
  : OTypeSequence (ProductTypeSequence f1 f2) :=
  fun n => @OTpair _ _ (g1 n) (g2 n).

Instance ProductSupremumSequence (f1 f2 : nat -> Type)
         {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
         (h1 : SupremumSequence f1) (h2 : SupremumSequence f2)
  : SupremumSequence (ProductTypeSequence f1 f2) :=
  fun n => ProductSupremum (f1 n) (f2 n).

Instance ProductCPOSequence (f1 f2 : nat -> Type)
         {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
         {h1 : SupremumSequence f1} {h2 : SupremumSequence f2}
         (i1 : CPOSequence f1) (i2 : CPOSequence f2)
  : CPOSequence (ProductTypeSequence f1 f2) :=
  fun n => ProductCPO (f1 n) (f2 n).

Instance ProductEPChain `(e1 : EPChain) `(e2 : EPChain) :
  EPChain (ProductTypeSequence f f0) :=
  fun n => pair_ep (chain n) (chain n).

(** Constant LHS continuous function space EP chain. *)
Definition ConstHomTypeSequence T `{CPO T} (f : nat -> Type)
           {g : OTypeSequence f} {h : SupremumSequence f}
           {i : CPOSequence f} : nat -> Type :=
  fun n => T -> f n.

Instance ConstHomOTypeSequence T `{CPO T} (f : nat -> Type)
         {g : OTypeSequence f} {h : SupremumSequence f}
         {i : CPOSequence f} : OTypeSequence (ConstHomTypeSequence T f) :=
  fun n => _.

Instance ConstHomSupremumSequence T `{CPO T} (f : nat -> Type)
         {g : OTypeSequence f} {h : SupremumSequence f}
         {i : CPOSequence f} : SupremumSequence (ConstHomTypeSequence T f) :=
  fun n => _.

Instance ConstHomCPOSequence T `{CPO T} (f : nat -> Type)
         {g : OTypeSequence f} {h : SupremumSequence f}
         {i : CPOSequence f} : CPOSequence (ConstHomTypeSequence T f) :=
  fun n => _.

Instance ConstHomEPChain T `{CPO T} `(G : EPChain)
  : EPChain (ConstHomTypeSequence T f) :=
  fun n => hom_ep (chain n).

(** Continuous function space EP chain. *)
Definition HomTypeSequence (f1 f2 : nat -> Type)
           {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
           {h1 : SupremumSequence f1} {h2 : SupremumSequence f2}
           {i1 : CPOSequence f1} {i2 : CPOSequence f2} : nat -> Type :=
  fun n => f1 n -c> f2 n.

Instance HomOTypeSequence (f1 f2 : nat -> Type)
         {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
         {h1 : SupremumSequence f1} {h2 : SupremumSequence f2}
         (i1 : CPOSequence f1) (i2 : CPOSequence f2)
  : OTypeSequence (HomTypeSequence f1 f2) :=
  fun n => _.

Instance HomSupremumSequence (f1 f2 : nat -> Type)
         {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
         {h1 : SupremumSequence f1} {h2 : SupremumSequence f2}
         (i1 : CPOSequence f1) (i2 : CPOSequence f2)
  : SupremumSequence (HomTypeSequence f1 f2) :=
  fun n => _.

Instance HomCPOSequence (f1 f2 : nat -> Type)
         {g1 : OTypeSequence f1} {g2 : OTypeSequence f2}
         {h1 : SupremumSequence f1} {h2 : SupremumSequence f2}
         (i1 : CPOSequence f1) (i2 : CPOSequence f2)
  : CPOSequence (HomTypeSequence f1 f2) :=
  fun n => _.

Instance HomEPChain `(e1 : EPChain) `(e2 : EPChain)
  : EPChain (HomTypeSequence f f0) :=
  fun n => chom_ep (chain n) (chain n).