From Categories Require Import Essentials.Notations.
From Categories Require Import Essentials.Types.
From Categories Require Import Essentials.Facts_Tactics.
From Categories Require Import Category.Main.
From Categories Require Import Functor.Main.
From Categories Require Import Cat.Cat.
From Categories Require Import Ext_Cons.Prod_Cat.Prod_Cat Ext_Cons.Prod_Cat.Operations.
From Categories Require Import Basic_Cons.Product.
From Categories Require Import Basic_Cons.Exponential.
From Categories Require Import NatTrans.NatTrans NatTrans.Func_Cat.
From Categories Require Import Cat.Product.

(** The exponential in cat is jsut the functor category. *)

Local Open Scope functor_scope.

(** Evaluation functor. *)
Program Definition Exp_Cat_Eval (C C' : Category) :
  ((Func_Cat C C') × C) –≻ C' :=
{|
  FO := fun x => ((fst x) _o (snd x))%object;
  FA := fun A B f => (((fst B) _a (snd f)) ∘ (@Trans _ _ _ _ (fst f) _))%morphism
|}.

Next Obligation. (* F_compose *)
Proof.
  repeat rewrite F_compose.
  repeat rewrite assoc.
  match goal with
      [|- (?A ∘ ?B = ?A ∘ ?C)%morphism] => cutrewrite (B = C); trivial
  end.
  repeat rewrite assoc_sym.
  match goal with
      [|- (?A ∘ ?B = ?C ∘ ?B)%morphism] => cutrewrite (A = C); trivial
  end.
  rewrite Trans_com; trivial.
Qed.

(* Exp_Cat_Eval defined *)

(** The arrow map of curry functor. *)
Program Definition Exp_Cat_morph_ex_A
        {C C' C'' : Category} (F : (C'' × C) –≻  C')
        (a b : C'') (h : (a –≻ b)%morphism)
  :
    ((Fix_Bi_Func_1 a F) –≻ (Fix_Bi_Func_1 b F))%nattrans :=
{|
  Trans := fun c => (F _a (h, id _ c))%morphism
|}.

(* Exp_Cat_morph_ex_A defined *)

Local Hint Extern 1 => apply NatTrans_eq_simplify; cbn.

(** The curry functor. *)
Program Definition Exp_Cat_morph_ex
        {C C' C'' : Category}
        (F : (C'' × C) –≻ C')
  :
    C'' –≻ (Func_Cat C C') :=
{|
  FO := fun a => Fix_Bi_Func_1 a F;
  FA := Exp_Cat_morph_ex_A F
|}.

(** Proof that currying after uncurrying gives back the same functor. *)
Lemma Exp_cat_morph_ex_eval_id
      {C C' C'' : Category}
      (u : C'' –≻ (Func_Cat C C'))
  :
    (u =
     Exp_Cat_morph_ex
       (
         (Exp_Cat_Eval C C')
           ∘ ((×ᶠⁿᶜ _ Cat_Has_Products) @_a (_, _) (_, _) (u, id Cat C))
       )
    )%morphism.
Proof.
  Func_eq_simpl.
  {
    extensionality a; extensionality b; extensionality h; cbn.
    apply NatTrans_eq_simplify.
    extensionality m.
    apply JMeq_eq.
    apply (@JMeq_trans _ _ _ _ (Trans (u _a h)%morphism m)).
    + match goal with [H : _ = _ |-_] => destruct H end; trivial.
    + cbn; auto.
  }
  {
    FunExt; cbn.
    Func_eq_simpl; FunExt; cbn.
    auto.
  }
Qed.

(** The exponential for category of categories (functor categories). *)
Program Definition Cat_Exponential (C C' : Cat) : (C ⇑ C')%object :=
{|
  exponential := Func_Cat C C';
  eval := Exp_Cat_Eval C C';
  Exp_morph_ex := fun C'' F => @Exp_Cat_morph_ex C C' C'' F
|}.

Next Obligation. (* Exp_morph_com *)
Proof.
  Func_eq_simpl.
  FunExt; cbn.
  rewrite <- F_compose; cbn.
  auto.
Qed.

Local Obligation Tactic := idtac.

Next Obligation. (* Exp_morph_unique *)
Proof.
  intros C C' z f u u' H1 H2; simpl in *.
  match type of H1 with
    _ = ?A => match type of H2 with
               _ = ?B => assert (A = B) as H3; auto; clear H1 H2
             end
  end.
  assert (H4 := @f_equal _ _ Exp_Cat_morph_ex _ _ H3).
  etransitivity; [apply Exp_cat_morph_ex_eval_id|].
  etransitivity; [|symmetry; apply Exp_cat_morph_ex_eval_id].
  trivial.
Qed.

(* Cat_Exponentials defined *)

Program Instance Cat_Has_Exponentials : Has_Exponentials Cat := Cat_Exponential.


Require Import Common.
Require DList.

Inductive type : Type :=
| Tlift : type0 -> type
| Ttuple : list type -> type
| Tunion : list type -> type.

Fixpoint eval_type (t : type) :=
  match t with
    | Tlift t => eval_type0 t
    | Ttuple l =>
        ETuple.of_list eval_type l
    | Tunion l =>
        List.fold_right (fun x acc => eval_type x + acc)%type False l
  end.

Inductive sync : Type :=
  | Treg : forall (t : type), sync
  | Tregfile : forall (n : nat) (t : type), sync.

Definition state := list sync.

Definition eval_sync (s : sync) :=
  match s with
    | Treg t => eval_type t
    | Tregfile n t => Regfile.T n (eval_type t)
  end.

Definition eval_state := Tuple.of_list eval_sync.

Notation Int n := (Tlift (Tint n)).
Notation Bool := (Tlift (Tbool)).
Notation Unit := (Tlift (Tunit)).

Inductive primitive (Phi : state) : list type -> type -> Type:=
  | register_read : forall t, var Phi (Treg t) ->  primitive Phi nil t
  | register_write : forall t, var Phi (Treg t) -> primitive Phi (t:: nil) Unit
  (* register file primitives *)
  | regfile_read : forall n t (v : var Phi (Tregfile n t)) p, primitive Phi ([Tlift (Tint p)])%list  t
  | regfile_write : forall n t (v : var Phi (Tregfile n t)) p, primitive Phi ([Tlift (Tint p); t])%list  Unit.

Section s.

  Variable Phi : state.

  Section t.
    Variable Var : type -> Type.
    Inductive expr :  type -> Type :=
    | Evar : forall t (v : Var t), expr t
    | Ebuiltin : forall args res (f : builtin args res),
                   DList.T  (fun j => expr (Tlift j)) (args) ->
                   expr  (Tlift res)
    | Econstant : forall  ty (c : constant0 ty), expr (Tlift ty)
    (* Tuple operations *)
    | Efst : forall l t, expr (Ttuple (t::l)) -> expr t
    | Esnd : forall l t, expr (Ttuple (t::l)) -> expr (Ttuple l)
    | Enth : forall l t (m : var l t), expr (Ttuple l) -> expr t
    | Etuple : forall l (exprs : DList.T (expr) l), expr (Ttuple l)

    (* Tagged-unions constructors *)
    | Econstr : forall l t (c : var l t), expr t -> (expr (Tunion l)).

    Inductive action : type -> Type :=
    | Return : forall t (exp : expr t), action t
    | Bind :
      forall t u
        (a : action  t)
        (f : Var t -> action u),
        action u
    | Assert : forall (e : expr Bool), action Unit
    | Primitive :
      forall args res (p : primitive Phi args res)
        (exprs : DList.T (expr) args),
        action res
    | Try : forall (a : action Unit), action Unit

    | Case : forall td l t (c : var l t) (e : expr (Tunion l)),
               (Var t -> action td)
               -> action td.
  End t.
  Definition Action t := forall Var, action Var t.
  Definition Expr t := forall Var, expr Var t.

  Notation eval_type_list := (DList.to_tuple eval_type).

  Definition eval_expr (t : type) (e : expr eval_type t) : eval_type t.
  refine (
      let fix eval_expr t e {struct e}:=
          match e with
            | Evar t v => v
            | Ebuiltin args res f exprs =>
                let exprs :=
                               DList.to_tuple
                               (fun (x : type0) (X : expr eval_type (Tlift x)) =>
                                  eval_expr (Tlift x) X)
                               exprs

                      in
                        builtin_denotation args res f exprs
            | Econstant ty c => c
            | Etuple l exprs =>
                DList.to_etuple eval_expr exprs
            | Enth l t v e =>
                ETuple.get l t v (eval_expr (Ttuple l) e)
            | Efst l t e =>
                ETuple.fst _ _ (eval_expr _ e)
            | Esnd l t e =>
                ETuple.snd _ _ (eval_expr _ e)
            | Econstr l t c e =>
                _
          end
      in eval_expr t e).

  refine (let fix fold {t : type}  (l : list type)  (c : var l t): eval_type t -> eval_type (Tunion l) :=
              match c as c' in (var E' t') return (eval_type t' -> eval_type (Tunion E')) with
                  | var_0 E t => fun e => inl (e) : eval_type (Tunion (t::E))
                  | var_S E t t' v' => fun e => inr (fold _ v' e) : eval_type (Tunion (t::E))
              end
          in fold  l c (eval_expr _ e1)
         ).

  Defined.

End s.

Delimit Scope expr_scope with expr.
Delimit Scope action_scope with action.

(** * Actions *)

Arguments Bind {Phi Var t u} _%action _%action.
Notation "'DO' X <- A ; B" := (Bind A (fun X => B)) (at level 200, X ident, A at level 100, B at level 200) : action_scope.

Arguments Assert {Phi Var} e%expr.

Definition When {Phi Var t} e a : action Phi Var t := @Bind Phi Var _ _ (Assert e) (fun _ => a).
Arguments When {Phi Var t} _%expr _%action.
Notation " 'WHEN' e ; B" := (When e B) (at level 200, e at level 100, B at level 200).

Arguments Return {Phi Var t} _%expr.
Notation " 'RETURN' e" := (Return e) (at level 200, e at level 100).

Arguments Primitive {Phi Var} args res _ _%expr.

(** * Primitives *)
Arguments register_read {Phi t} _.
Notation "'read' [: v ]" := (Primitive nil _ (register_read v) DList.nil) (no associativity).

Arguments register_write {Phi t} _.
Notation "'write' [: v <- w ]" := (Primitive (cons _ nil) Unit (register_write v) (DList.cons (w)%expr (DList.nil))) (no associativity).

Arguments regfile_read {Phi n t} _ p.
Notation "'read' M [: v ]" := (Primitive ([Tlift (W _)])%list _ (regfile_read M _ ) (DList.cons (v)%expr (DList.nil))) (no associativity).

Arguments regfile_write {Phi n t} _ p.
Notation "'write' M [: x <- v ]" := (Primitive ([Tlift (W _); _])%list _ (regfile_write M _ ) (DList.cons (x)%expr (DList.cons (v)%expr (DList.nil)))) (no associativity).

(** * Expressions  *)
Arguments Enth {Var l t} m _%expr.
Arguments Efst {Var l t} _%expr.
Arguments Esnd {Var l t} _%expr.
Arguments Etuple {Var l} exprs%dlist.
Notation "[< a ; b >]" := (Etuple ([a%expr ; b%expr])%dlist)(at level 0) : expr_scope.

Arguments Evar  {Var t} _.
Notation "! x" := (Evar x) (at level  10) : expr_scope .

Arguments Ebuiltin {Var} {args res} _ _%expr.
Notation "{< f ; x ; y >}" := (Ebuiltin f (DList.cons  (x)%expr (DList.cons (y)%expr DList.nil))).
Notation "{< f ; x >}" := (Ebuiltin f (DList.cons (x)%expr DList.nil)).

Notation "~ x" :=  ({< BI_negb ; x >}) : expr_scope.
Notation "a || b" := ({< BI_orb ; a ; b >}) : expr_scope.
Notation "a && b" := ({< BI_andb ; a ; b >}) : expr_scope.
Notation "a - b" := ({< BI_minus _ ; a ; b >}) : expr_scope.
Notation "a + b" := ({< BI_plus _ ; a ; b >}) : expr_scope.
Notation "a = b" := ({< BI_eq _ ; a ; b >}) : expr_scope.
Notation "a < b" := ({< BI_lt _ ; a ; b >}) : expr_scope.
Notation "x <= y" := ((x < y) || (x = y))%expr : expr_scope.
Notation "x <> y" := (~(x = y))%expr : expr_scope.

Arguments Econstant {Var ty} _.
Notation "#i x" := (Econstant (Cword x)) (at level 0).
Notation "#b x" := (Econstant (Cbool x)) (at level 0).
Notation "#" := Econstant.

Definition Ctt : constant0 Tunit := tt.

Definition constrs t := match t with Tunion l => l | _ => nil end.
Arguments Case {Phi Var _} l t _ _  _.

(* Notation follow x B := (fun R =>                            *)
(*                           ( *)
(*                             DO x <- Return (Efst R);  *)
(*                             B (Esnd R) *)
(*                           )%action                            *)
(*                        ) (only parsing).  *)

(* Notation "|- x => B" := (follow x (fun _ => B)) (at level 0) : action_scope.  *)
(* Notation "|- x , y => B" := (follow x (follow y (fun _ => B))) (at level 0) : action_scope.  *)
(* Notation "|- x , y , z => B" := (follow x (follow y (follow z (fun _ => B)))) (at level 0) : action_scope.  *)

Notation "'MATCH' E 'IN' T 'WITH' constr 'OF' x ==> B" :=
(
  (DO XE <- (Return E);
   Case  (constrs T) _ (constr) (!XE)%expr
         (fun x => B )
  )%action
) (at level 200, E at level 80, T at level 80) : action_scope.

Notation "'MATCH' E 'IN' T 'WITH' constr 'OF' x , y ==> B" :=
(
  (DO XE <- (Return E);
   Case  (constrs T) _ (constr) (!XE)%expr
         (fun XY =>
            DO x <- Return (Efst (!XY)%expr);
            DO y <- Return (Efst (Esnd (!XY)%expr));
            B
         )
  )%action
) (at level 200, E at level 80, T at level 80) : action_scope.

Notation "'MATCH' E 'IN' T 'WITH' constr 'OF' x , y , z ==> B" :=
(
  (DO XE <- (Return E);
   Case  (constrs T) _ (constr) (!XE)%expr
         (fun XYZ =>
            DO x <- Return (Efst (!XYZ)%expr);
            DO y <- Return (Efst (Esnd (!XYZ)%expr));
            DO z <- Return (Efst( Esnd (Esnd (!XYZ))%expr));
            B
         )
  )%action
) (at level 200, E at level 80, T at level 80) : action_scope.

Record TRS : Type := mk_TRS
  {
    Phi : state;
    rules : list (Action Phi Unit)
  }.

Module Diff.
  Section t.
    Variable Phi : state.
    Definition T := Tuple.of_list (option ∘ eval_sync) Phi.
    Definition add (Delta : T) t (v : var Phi t) w : option T :=
      match Tuple.get Phi t v Delta  with
        | None =>  Some  (Tuple.set _ _ Phi t v (Some w) Delta)
        | Some _ => None
      end.

    Definition init : T := Tuple.init _ _ (fun _ => None) Phi.
    Definition apply : T  -> Tuple.of_list eval_sync Phi -> Tuple.of_list eval_sync Phi :=
                Tuple.map2 (fun s delta old => match delta with
                                             | None => old
                                             | Some delta => delta
                                           end
                            ) Phi.
  End t.

End Diff.

Module Sem.

  Module Dyn.
    (** dynamic semantics:
    - continuation passing style: the current diff is threaded through the execution
    - illegal behaviors are denoted by None *)
    Section t.
      Variable Phi : state.
      (* semantics :  Action Phi t -> T t  *)
      Definition T t := eval_state Phi -> Diff.T Phi -> option (t * Diff.T Phi).
      Definition Return {t} (e : t) : T t  := fun st d => Some (e, d).
      Definition Bind {s t} : T s -> (s -> T t) -> T t :=
        fun (x : T s) (f : s -> T t) (st : eval_state Phi) (d : Diff.T Phi) =>
          match x st d with
              Some (e, d') =>  f e st d'
            | None => None
          end.
      Definition Fail {t} : T t := fun _ _ => None.

      Definition Try : T unit -> T unit :=
        fun a st Delta =>
          match a st Delta with
            | None => Some (tt, Delta)
            | x =>  x
          end.

      Definition Run (CMD : T unit) (st : eval_state Phi) :=
        match CMD st (Diff.init Phi) with
          | None => None
          | Some Delta => Some (Diff.apply Phi (snd Delta) st)
        end.

      Definition primitive_denote  args res (p : primitive Phi args res) (exprs : ETuple.of_list eval_type args) : T (eval_type res).
      refine (match
          p in (primitive _ l t) return (ETuple.of_list eval_type l -> T (eval_type t))
        with
          | register_read t v =>
              fun _ (st : eval_state Phi)
                  (Delta : Diff.T Phi) => Some (Tuple.get Phi (Treg t) v st, Delta)
          | register_write t v =>
              fun (exprs : ETuple.of_list eval_type [t]) (_ : eval_state Phi) (Delta : Diff.T Phi) =>
                let w := exprs in
                  do Delta2 <- Diff.add Phi Delta (Treg t) v w;
                  Some (tt, Delta2)
          | regfile_read n t v p =>
              fun (exprs : ETuple.of_list eval_type [Tlift (W p)])
                  (st : eval_state Phi) (Delta : Diff.T Phi) =>
                let rf := Tuple.get Phi (Tregfile n t) v st in
                let adr := exprs in
                  let v := Regfile.get rf (Word.unsigned adr) in Some (v, Delta)
          | regfile_write n t v p =>
              fun (exprs : ETuple.of_list eval_type [Tlift (W p); t])
                  (st : eval_state Phi) (Delta : Diff.T Phi) =>
                let rf := Tuple.get Phi (Tregfile n t) v st in
                let rf := match exprs with
                              (adr, (w)) => Regfile.set rf (Word.unsigned adr) w
                          end
                in
                  do Delta2 <- Diff.add Phi Delta (Tregfile n t) v rf; Some (tt, Delta2)

        end exprs).
      Defined.
    End t.
  End Dyn.

  (** dynamic semantics:
    - all guard failures are represented by None.
    - continuation passing style: the current diff is threaded through the execution
  *)
  Section t.
    Variable Phi : state.
    Definition eval_action (t : type) (a : action Phi eval_type t) :
      (Dyn.T Phi (eval_type t)).

    refine (
        let fix eval_action (t : type) (a : action Phi eval_type t) :
            Dyn.T Phi ( eval_type t) :=
            match a with
              | Return t exp => Dyn.Return Phi (eval_expr _ exp)
              | Bind t u a f =>
                  let act := eval_action _ a in
                    let f' := (fun e => eval_action u (f e)) in
                      Dyn.Bind Phi act f'
              | Assert e =>
                  let g := eval_expr _ e in
                    match g with
                      | true => Dyn.Return Phi tt
                      | false => Dyn.Fail  Phi
                    end
              | Primitive args res p exprs =>
                  Dyn.primitive_denote Phi args res p (DList.to_etuple eval_expr exprs)
              | Try a =>
                  let a := eval_action _ a in
                    Dyn.Try Phi a
              | Case td l t c e f =>
                  _
          end
      in  eval_action t a).

    refine (let fix fold {t : type}  (l : list type)  (c : var l t) :
            eval_type (Tunion l) -> (eval_type t -> Dyn.T Phi (eval_type td)) ->
            Dyn.T Phi (eval_type td) :=
            match c as c' in (var E' t') return
               eval_type (Tunion E') -> (eval_type t' -> Dyn.T Phi (eval_type td)) ->
               Dyn.T Phi (eval_type td) with
               | var_0 E t => fun disj f => match disj with
                                          | inl x => f x
                                          | inr _ => Dyn.Fail Phi
                                          end
               | var_S E t t' v' => fun disj f =>
                                     match disj with
                                         | inl _ => Dyn.Fail Phi
                                         | inr x => fold _ v' x f
                                     end
            end
          in fold  l c (eval_expr _ e) (fun x => eval_action _ (f x))
         ).
    Defined.

  End t.
  Arguments eval_action {Phi} {t} _%action _ _.
End Sem.

Section run.

  Variable T : TRS.
  Notation rule := (Action (Phi T) Unit).

  Definition run_rule (R : rule) :=
    fun st => Sem.Dyn.Run _ (Sem.eval_action (R eval_type)) st.

  Fixpoint first_rule (l : list rule) x :=
    match l with
      | nil => Some x
      | cons R q =>
          match run_rule R x with
            | None => first_rule q x
            | Some x => Some x
          end
    end.

  Fixpoint iter_option {A} n (f : A -> option A) x :=
    match n with
      | 0 => Some x
      | S n => match f x with | None => Some x | Some x => iter_option n f x end
    end.

  Fixpoint run_unfair n x :=
    match n with
      | 0 => Some x
    | S n =>
        match first_rule (rules T) x with
          | None => Some x
          | Some x => run_unfair n x
        end
  end.

End run.

Module WF.
  Section t.
    Variable va vb : type -> Type.
    Variable Phi : state.
    Reserved Notation "l & Gamma |- x == y" (at level 80).

    Inductive
      equiv_expr : forall t, expr va t  -> expr vb t -> Prop := .

    Notation ctxt l := (DList.T  (fun x => (va x * vb x)%type) l).

    Inductive
      equiv_action : forall (l : list type), ctxt l ->   forall t, action Phi va t  -> action Phi vb t -> Prop :=
    | EA_Return : forall l Gamma t ea eb,
                    equiv_expr t ea eb ->
                    l & Gamma |- @Return Phi va t ea == @Return Phi vb t eb
    | EA_Bind  : forall l Gamma t u a1 a2 f1 f2,
                   l & Gamma |- a1 == a2 ->
                      (forall v1 v2, (t::l)%list &  ((v1,v2)::Gamma)%dlist |- f1 v1 == f2 v2) ->
                      l & Gamma |- @Bind Phi va t u a1 f1 == @Bind Phi vb t u a2 f2
    | EA_Assert : forall l Gamma ea eb,
                    equiv_expr _ ea eb ->
                    l & Gamma |- @Assert Phi va ea == @Assert Phi vb eb
    | EA_Primitive : forall l Gamma args res p lea leb,
                       l & Gamma |- @Primitive Phi va args res p lea ==  @Primitive Phi vb args res p leb
    | EA_Try : forall l Gamma a1 a2,
                 l & Gamma |- a1 == a2 ->
                    l & Gamma |- @Try Phi va a1 == @Try Phi vb a2
    | EA_Case : forall l Gamma td constrs t c ea eb fa fb,
                  equiv_expr _ ea eb ->
                  (forall v1 v2, (t::l)%list &  ((v1,v2)::Gamma)%dlist |- fa v1 == fb v2) ->
                  l & Gamma |- (@Case Phi va td constrs t c ea fa) == (@Case Phi vb td constrs t c eb fb)
                        where "l & Gamma |- x == y" := (equiv_action l Gamma _ x y).
End t.
End WF.
Notation " l & Gamma |- x == y" := (WF.equiv_action _ _ _ l Gamma _ x y) (no associativity, at level 80).

Open Scope action_scope.


(* This code is copyrighted by its authors; it is distributed under  *)
(* the terms of the LGPL license (see LICENSE and description files) *)

(* Contribution to the Coq Library   V6.3 (July 1999)                    *)
(****************************************************************************)
(*                                                                          *)
(*                                                                          *)
(*                   Solange Coupet-Grimal & Line Jakubiec                  *)
(*                                                                          *)
(*                                                                          *)
(*              Laboratoire d'Informatique de Marseille                     *)
(*               CMI-Technopole de Chateau-Gombert                          *)
(*                   39, Rue F. Joliot Curie                                *)
(*                   13453 MARSEILLE Cedex 13                               *)
(*           e-mail:{Solange.Coupet,Line.Jakubiec}@lim.univ-mrs.fr          *)
(*                                                                          *)
(*                                                                          *)
(*                                Coq V5.10                                 *)
(*                              May 30th 1996                               *)
(*                                                                          *)
(****************************************************************************)
(*                              Comp_Synth .v                               *)
(****************************************************************************)

Require Import Factorization_Synth.
Require Import Comparator_Relation.

Parameter BASE : BT.
Definition b := base BASE.
Definition Num := num BASE.
Definition Val_bound := val_bound BASE.

Lemma Comparator :
 forall (n : nat) (o : order) (X Y : Num n),
 {o' : order | R (exp b n) o (Val_bound n X) (Val_bound n Y) o'}.
intros n o X Y.
unfold R in |- *; unfold b in |- *; unfold Val_bound in |- *.
apply factorization_for_synthesis.
exact is_factorizable.
exact (is_proper BASE).
Defined.

(*************************************************************************************)


Set Implicit Arguments.

Require Import monads.
(*Require ne_list.*)
Require ne_tree.
Require Import util.

Definition C := ne_tree.T.

Definition ret {A: Set}: A -> C A := @ne_tree.Leaf A.

Fixpoint bind (A B: Set) (m: C A) (k: A -> C B): C B :=
  match m with
  | ne_tree.Leaf a => k a
  | ne_tree.Node ts => ne_tree.Node (ne_list.map (fun x => bind x k) ts)
  end.

Let runit (a b: Set) (x: a) (f: a -> C b): bind (ret x) f = f x.
Proof. auto. Qed.

Fixpoint lunit A (f: C A) {struct f}: bind f ret = f :=
  match f return bind f ret = f with
  | ne_tree.Leaf x => refl_equal
  | ne_tree.Node l =>
    eq_ind_r (fun l0 => ne_tree.Node l0 = ne_tree.Node l) refl_equal
      ((fix F (l: ne_list.L (C A)) :=
        match l return ne_list.map (fun u => bind u ret) l = l with
        | ne_list.one x => eq_ind_r (fun c => ne_list.one c = ne_list.one x) refl_equal (lunit x)
        | ne_list.cons x y => eq_ind_r (fun c => ne_list.cons c (ne_list.map (fun x => bind x ret) y) = ne_list.cons x y) (eq_ind_r (fun l => ne_list.cons x l = ne_list.cons x y) refl_equal (F y)) (lunit x)
        end) l)
  end.

Let assoc (a b c: Set) (n: C a) (f: a -> C b) (g: b -> C c):
  bind (bind n f) g = bind n (fun x: a => bind (f x) g).
Proof with auto.
  intros.
  generalize n. clear n.
  apply (ne_tree.alt_rect2 (fun n => bind (bind n f) g = bind n (fun x: a => bind (f x) g)) (fun l => ne_list.map (fun x => bind (bind x f) g) l = ne_list.map (fun x => bind x (fun x0 => bind (f x0) g)) l)); intros; simpl...
      rewrite ne_list.map_map.
      unfold compose.
      rewrite H...
    rewrite H...
  rewrite H.
  rewrite H0...
Qed.

Definition M: Monad := Build_Monad C bind (@ret) runit lunit assoc.

Lemma ext: extMonad M.
Proof with auto.
  unfold extMonad.
  induction x using ne_tree.alt_ind; simpl...
  replace (ne_list.map (fun x: C A => bind x f) l) with (ne_list.map (fun x: C A => bind x g) l)...
  generalize H0. clear H0.
  induction l; simpl; intros; rewrite H0...
  rewrite IHl...
Qed.

Lemma bind_Leaf (A B: Set) (x: A) (f: A -> M B): bind (ne_tree.Leaf x) f = f x.
Proof. auto. Qed.

Lemma bind_Node (A B: Set) (x: ne_list.L (ne_tree.T A)) (f: A -> M B):
  bind (ne_tree.Node x) f = ne_tree.Node (ne_list.map (fun x0: C A => bind x0 f) x).
Proof. auto. Qed.

Lemma bind_Node_one (X Y: Set) (t: M X) (g: X -> M Y):
  bind (ne_tree.Node (ne_list.one t)) g = ne_tree.Node (ne_list.one (t >>= g)).
Proof. auto. Qed.

Lemma bind_Node_cons (X Y: Set) (t: M X) (l: ne_list.L (M X)) (g: X -> M Y):
  bind (ne_tree.Node (ne_list.cons t l)) g = ne_tree.Node (ne_list.cons (bind t g) (ne_list.map (fun x => bind x g) l)).
Proof. auto. Qed.

Lemma bind_map (X Y: Set) (f: X -> Y) (x: M X): bind x (ret ∘ f) = ne_tree.map f x.
Proof with try reflexivity.
  induction x...
    simpl.
    rewrite IHx...
  simpl.
  rewrite IHx.
  f_equal.
  simpl in IHx0.
  inversion_clear IHx0...
Qed.

Definition deterministic (X: Set) (x: M X) (v: X): Prop := x = ne_tree.Leaf v.

Lemma deterministic_ret (A: Set) (a: A): deterministic (ret a) a.
Proof. unfold deterministic. auto. Qed.

Lemma ex_deterministic_ret (X: Set) (x: X): exists u, deterministic (ret x) u.
Proof. intros. exists x. apply deterministic_ret. Qed.

Lemma deterministic_bind:
  forall (A: Set) (a: M A) (z: A), deterministic a z ->
  forall (B: Set) (b: A -> M B) (v: B), deterministic (b z) v ->
    exists w, deterministic (a >>= b) w.
Proof with auto. unfold deterministic. intros. subst. simpl. exists v... Qed.

Lemma deterministic_bind_weak:
  forall (A: Set) (a: M A) (z: A), deterministic a z ->
  forall (B: Set) (b: A -> M B), (forall q, exists v, deterministic (b q) v) ->
  exists w, deterministic (a >>= b) w.
Proof with auto. intros. destruct (H0 z). apply (deterministic_bind H _ H1). Qed.

Lemma ex_deterministic_bind:
  forall (A: Set) (a: M A) (z: A), deterministic a z ->
  forall (B: Set) (b: A -> M B), (exists v, deterministic (b z) v) ->
    exists w, deterministic (a >>= b) w.
Proof. intros. destruct H0. apply (deterministic_bind H _ H0). Qed.

Lemma ex_deterministic_bind_weak:
  forall (A: Set) (a: M A) (z: A), deterministic a z ->
  forall (B: Set) (b: A -> M B), (forall q, exists v, deterministic (b q) v) ->
    exists w, deterministic (a >>= b) w.
Proof. intros. apply (deterministic_bind_weak H). assumption. Qed.

Definition pick T: ne_list.L T -> M T := @ne_tree.Node T ∘ ne_list.map (@ne_tree.Leaf T).

Lemma In_bind_inv (X Y: Set) (f: X -> M Y) (x: M X) r:
  ne_tree.In r (bind x f) -> exists z, ne_tree.In z x /\ ne_tree.In r (f z).
Proof with eauto.
  induction x in r |- *...
    simpl.
    intros.
    inversion_clear H.
    inversion_clear H0.
    destruct (IHx r H).
    destruct H0...
  intros.
  inversion_clear H.
  inversion_clear H0.
    destruct (IHx r H).
    destruct H0...
  destruct (IHx0 r (ne_tree.InNode H)).
  destruct H0.
  inversion_clear H0...
Qed.

Lemma bind_eq (X X' Y XM: Set)
  (f: X -> M Y) (f': X' -> M Y)
  (xm: X -> XM) (xm': X' -> XM):
  (forall p q, xm p = xm' q -> f p = f' q) ->
  forall (x: M X) (x': M X'),
  ne_tree.map xm x = ne_tree.map xm' x' ->
    bind x f = bind x' f'.
Proof with auto.
  induction x; simpl; destruct x'; simpl; intros; try discriminate.
      apply H.
      inversion_clear H0...
    destruct l; simpl; intros.
      simpl.
      intros.
      rewrite (IHx t)...
      inversion_clear H0...
    discriminate.
  destruct l0.
    discriminate.
  simpl in H0.
  simpl.
  inversion H0.
  clear H0.
  f_equal.
  replace (bind x f) with (bind t f').
    Focus 2.
    symmetry.
    apply IHx...
  replace (ne_list.map (fun x0: C X => bind x0 f) l) with (ne_list.map (fun x0: C X' => bind x0 f') l0)...
  assert (ne_tree.map xm (ne_tree.Node l) = ne_tree.map xm' (ne_tree.Node l0)).
    simpl.
    rewrite H3...
  cset (IHx0 (ne_tree.Node l0) H0).
  inversion_clear H1...
Qed.

Lemma map_bind (X Y Z: Set) (f: Y -> Z) (g: X -> M Y) (x: M X):
  ne_tree.map f (bind x g) =
  bind x (fun xx => ne_tree.map f (g xx)).
Proof with auto.
  induction x...
    simpl.
    rewrite IHx...
  simpl.
  rewrite IHx.
  f_equal.
  simpl in IHx0.
  inversion IHx0...
Qed.

Coercion ne_tree_isa_monad (A: Set) (a: ne_tree.T A): M A := a.


Require Import Coq.Lists.List.
Require Import ExtLib.Data.Option.
Require Import ExtLib.Tactics.Consider.
Require Import ExtLib.Tactics.Injection.

Set Implicit Arguments.
Set Strict Implicit.

Section iteration.
  Context (T U : Type) (f : T -> option U).

  Fixpoint first_success  (ls : list T) : option U :=
    match ls with
      | nil => None
      | l :: ls =>
        match f l with
          | None => first_success ls
          | x => x
        end
    end.

  Lemma first_success_sound
  : forall ls val,
      first_success ls = Some val ->
      exists l,
        In l ls /\ f l = Some val.
  Proof.
    induction ls; simpl; intros.
    - congruence.
    - consider (f a); intros.
      + exists a. inv_all; subst. auto.
      + apply IHls in H0. destruct H0. intuition. eauto.
  Qed.

  Variable (f' : T -> U -> option U).

  Fixpoint all_success (ls : list T) (acc : U)
  : option U :=
    match ls with
      | nil => Some acc
      | l :: ls =>
        match f' l acc with
          | None => None
          | Some x => all_success ls x
        end
    end.

  Fixpoint filter_map (ls : list T) : list U :=
    match ls with
      | nil => nil
      | l :: ls => match f l with
                     | None => filter_map ls
                     | Some x => x :: filter_map ls
                   end
    end.

  Lemma filter_map_sound
  : forall ls,
      Forall (fun x => exists l, In l ls /\ f l = Some x)
             (filter_map ls).
  Proof.
    clear. induction ls; simpl; intros.
    { constructor. }
    { assert (Forall (fun x : U => exists l : T, (a = l \/ In l ls) /\ f l = Some x)
                     (filter_map ls)).
      { eapply Forall_impl; eauto.
        simpl in *. intuition.
        destruct H as [ ? [ ? ? ] ].
        eexists; split. right. eassumption. eauto. }
      { consider (f a); intros; auto.
        constructor; eauto. } }
  Qed.

  Lemma in_filter_map_iff : forall ls x,
                              List.In x (filter_map ls) <->
                              exists y, f y = Some x /\ List.In y ls.
  Proof.
    clear.
    induction ls; simpl.
    { intuition. destruct H; intuition. }
    { intuition.
      { consider (f a); intros.
        { destruct H0. subst. eauto.
          eapply IHls in H0. destruct H0. intuition; eauto. }
        { eapply IHls in H0. destruct H0; intuition; eauto. } }
      { destruct H. destruct H.
        destruct H0; subst.
        { rewrite H. left. auto. }
        { destruct (f a); try right; apply IHls; eauto. } } }
  Qed.

End iteration.


Require Import Coq.Strings.String
        Coq.Sets.Ensembles
        Coq.Bool.Bool.
Require Import Fiat.Common
        Fiat.Common.BoolFacts
        Fiat.Common.LogicFacts
        Fiat.Computation.Core
        Fiat.Computation.Monad
        Fiat.Computation.SetoidMorphisms
        Fiat.Computation.Refinements.Tactics.

(** General Lemmas about the behavior of [computes_to], [refine], and
    [refineEquiv]. *)

Local Arguments impl _ _ / .

Section general_refine_lemmas.

  Lemma refine_under_bind {A B}
  : forall c (x y : A -> Comp B),
    (forall a, c ↝ a -> refine (x a) (y a))
    -> refine (a <- c; x a) (a <- c; y a).
  Proof. t_refine.  Qed.

  (* Combines refine_under_bind and refine_bind:
generates another subgoal for the first expression,
and gives you the computational hypothesis for the second *)
  Lemma refine_under_bind_both {A B}
    : forall c c' (x y : A -> Comp B),
      refine c c' ->
      (forall a, c ↝ a -> refine (x a) (y a))
      ->
      refine (a <- c; x a) (a <- c'; y a).
  Proof. t_refine. Qed.

  Lemma refine_pick A (P : A -> Prop) c (H : forall x, c ↝ x -> P x)
    : @refine A ({x : A | P x })%comp
              c.
  Proof. t_refine. Qed.

  Lemma refine_pick_val A a (P : A -> Prop)
    : P a -> @refine A ({x | P x })%comp
                     (ret a).
  Proof.
    t_refine.
  Qed.

  Lemma refine_bind_pick :
    forall (A B : Type) (P : Ensemble A),
    forall x y : A -> Comp B,
      (forall a, P a -> refine (x a) (y a)) ->
      refine (a <- {a | P a};
              x a)
             (a <- {a | P a};
              y a).
  Proof. t_refine. Qed.

  Lemma refine_pick_pick A (P1 P2 : A -> Prop)
        (H : forall x, P2 x -> P1 x)
    : @refine _
              { x : A | P1 x }%comp
              { x : A | P2 x }%comp.
  Proof. t_refine. Qed.

  Lemma refineEquiv_pick_pick A (P1 P2 : A -> Prop)
        (H : forall x, P1 x <-> P2 x)
    : @refineEquiv _
                   { x : A | P1 x }%comp
                   { x : A | P2 x }%comp.
  Proof. t_refine. Qed.

  Lemma refineEquiv_pick_pair A B (PA : A -> Prop) (PB : B -> Prop)
    : @refineEquiv _
                   { x : A * B | PA (fst x) /\ PB (snd x) }%comp
                   (a <- { a : A | PA a };
                    b <- { b : B | PB b };
                    ret (a, b))%comp.
  Proof. t_refine. Qed.

  Lemma refineEquiv_pick_pair_fst_dep A B (PA : A * B -> Prop) (PB : B -> Prop)
    : @refineEquiv _
                   {x | PA x /\ PB (snd x)}%comp
                   (b <- { b | PB b };
                    a <- { a | PA (a, b) };
                    ret (a, b))%comp.
  Proof.
    t_refine.
  Qed.

  Lemma refineEquiv_pick_pair_snd_dep A B (PA : A -> Prop) (PB : A * B -> Prop)
    : @refineEquiv _
                   { x | PA (fst x) /\ PB x }%comp
                   (a <- { a | PA a };
                    b <- { b | PB (a, b) };
                    ret (a, b))%comp.
  Proof. t_refine. Qed.

  Lemma refineEquiv_pick_pair_pair A B C
        (PA : A -> Prop) (PB : B -> Prop) (PC : C -> Prop)
    : @refineEquiv _
                   { x | (PA (fst (fst x))
                          /\ PB (snd (fst x)))
                         /\ PC (snd x)}%comp
                   (a <- { a | PA a };
                    b <- { b | PB b };
                    c <- { c | PC c };
                    ret (a, b, c))%comp.
  Proof. t_refine. Qed.

  Definition refineEquiv_split_ex A B
             (P : A -> Prop) (P' : A -> B -> Prop)
    : @refineEquiv _
                   { b | exists a, P a /\ P' a b }%comp
                   (a <- { a | P a /\ exists b, P' a b };
                    { b | P' a b })%comp.
  Proof.
    split; intros v Comp_v; computes_to_inv;
    intuition; destruct_ex;
    repeat (computes_to_econstructor; eauto);
    intuition eauto.
  Qed.

  Definition refineEquiv_pick_contr_ret A (P : A -> Prop)
             (x : A) (H : unique P x)
    : @refineEquiv _
                   { y | P y }
                   (ret x).
  Proof. t_refine. Qed.

  Definition refineEquiv_pick_eq
             A (x : A)
    : @refineEquiv _
                   { y | y = x }%comp
                   (ret x).
  Proof. t_refine. Qed.

  Definition refineEquiv_pick_eq' A (x : A)
    : @refineEquiv _ { y | x = y }%comp
                   (ret x).
  Proof. t_refine. Qed.

  Definition refine_pick_eq
             A (x : A)
    : @refine _
              { y | y = x }%comp
              (ret x).
  Proof. apply refineEquiv_pick_eq. Qed.

  Definition refine_pick_eq' A (x : A)
    : @refine _ { y | x = y }%comp
              (ret x).
  Proof. apply refineEquiv_pick_eq'. Qed.

  Definition refineEquiv_split_func_ex
    : forall A B (P : A -> Prop) (f : A -> B),
      @refineEquiv _ { b | exists a, P a /\ b = f a}%comp
                   (a <- { a | P a};
                    ret (f a))%comp.
  Proof.
    split; intros v Comp_v; computes_to_inv;
    intuition; destruct_ex;
    repeat (computes_to_econstructor; eauto);
    intuition eauto.
    subst; eauto.
  Qed.

  Definition refineEquiv_split_func_ex2
             A A' B (P : A -> Prop) (P' : A' -> Prop)
             (f : A -> A' -> B)
    : refineEquiv { b | exists a, P a /\ exists a', P' a' /\ b = f a a'}
                  (a <- { a | P a};
                   a' <- { a' | P' a'};
                   ret (f a a')).
  Proof.
    split; intros v Comp_v; computes_to_inv;
    intuition; destruct_ex;
    computes_to_econstructor; eauto;
    intuition; destruct_ex; intuition; try subst;
    computes_to_econstructor; eauto.
  Qed.

  Definition refineEquiv_split_func_ex2'
             A A' B (P : A -> Prop) (P' : A' -> Prop)
             (f : A -> A' -> B)
    : refineEquiv { b | exists a, P a /\ exists a', P' a' /\ f a a' = b}
                  (a <- { a | P a};
                   a' <- { a' | P' a'};
                   ret (f a a')).
  Proof.
    split; intros v Comp_v; computes_to_inv;
    intuition; destruct_ex;
    computes_to_econstructor; eauto;
    intuition; destruct_ex; intuition; try subst;
    computes_to_econstructor; eauto.
  Qed.

  Definition refineEquiv_pick_computes_to A (c : Comp A)
    : refineEquiv { v | c ↝ v } c.
  Proof. t_refine. Qed.

  Definition refine_pick_computes_to A (c : Comp A)
    : refine { v | c ↝ v } c.
  Proof. t_refine. Qed.

  Lemma split_refineEquiv_fst_proj1_sig A B P Q
    : refineEquiv { x : { x : A * B | P x } | Q (fst (proj1_sig x)) }
                  (x <- { x : A | Q x };
                   y <- { y : B | P (x, y) };
                   pf <- { pf : P (x, y) | True };
                   ret (exist P _ pf)).
  Proof. t_refine. Qed.

  Definition split_refine_fst_proj1_sig A B P Q
    := proj1 (@split_refineEquiv_fst_proj1_sig A B P Q).

  Lemma split_refineEquiv_snd_proj1_sig A B P Q
    : refineEquiv { x : { x : A * B | P x } | Q (snd (proj1_sig x)) }
                  (x <- { x : B | Q x };
                   y <- { y : A | P (y, x) };
                   pf <- { pf : P (y, x) | True };
                   ret (exist P _ pf)).
  Proof. t_refine. Qed.

  Definition refineEquiv_pick_computes_to_and A (c : Comp A)
             (P : Ensemble A)
    : refineEquiv { v | c ↝ v /\ P v}
                  (v <- c;
                   { a | a = v /\ P a}).
  Proof. t_refine. Qed.

  Definition split_refine_snd_proj1_sig A B P Q
    := proj1 (@split_refineEquiv_snd_proj1_sig A B P Q).

  Lemma split_refineEquiv_proj1_sig A P Q
    : refineEquiv { x : { x : A | P x } | Q (proj1_sig x) }
                  (x <- { x | P x /\ Q x };
                   p <- { _ : P x | True };
                   ret (exist P x p)).
  Proof. t_refine. Qed.

  Lemma refineEquiv_pick_forall_eq
        A B (a : A) (P : A -> B -> Prop)
    : @refineEquiv _
                   (Pick (fun b => forall a', a = a' -> P a' b))
                   (Pick (P a)).
  Proof. t_refine. Qed.

  Lemma refine_if_bool_eta :
    forall (u : bool),
      refine (if u then (ret true) else (ret false))
             (ret u).
  Proof. destruct u; reflexivity. Qed.

  Open Scope comp_scope.

  Lemma refinement_step {A} (c c' : Comp A) :
    refine c c'
    -> Refinement of c'
    -> Refinement of c.
  Proof.
    intros refine_c_c' Refinement_c'.
    destruct Refinement_c' as [c'' refine_c'_c''].
    intros; econstructor.
    etransitivity; eassumption.
  Defined.

  Lemma refine_pick_forall_Prop
        A B C (a : A) (Q : C -> A -> Prop) (P : A -> C -> B -> Prop)
        b
    :
      (forall c : C, Q c a -> refine {x : B | P a c x} b) ->
      refine {b' : B | forall c : C, Q c a -> P a c b'} b.
  Proof.
    unfold refine; intros; computes_to_econstructor; intros.
    generalize (H _ H1 _ H0); intros.
    computes_to_inv; assumption.
  Qed.

  Lemma refine_pick_eq_ex_bind {A B : Type}
        (P : A -> B -> Prop)
        (a : A)
    : refine (a' <- {a' | a' = a /\ exists b, P a' b};
              {b | P a' b})
             {b | P a b}.
  Proof.
    t_refine.
  Qed.

  (* This helper lemma makes terms more ameneable to
     setoid_rewriting. *)
  Lemma refine_if_If {A}
    : forall (c : bool) (t e : Comp A),
      refine (if c then t else e)
             (If c Then t Else e).
  Proof.
    reflexivity.
  Qed.

  (* Nontermination with above lemma? *)
  (*
  Lemma refine_If_if {A}
  : forall (c : bool) (t e : Comp A),
      refine (If c Then t Else e)
             (if c then t else e).
  Proof.
    reflexivity.
  Qed.
   *)

  Lemma refine_if_andb {A}
    : forall (i i' : bool)
             (t e : A),
      (if i then (if i' then t else e) else e) =
      if i && i' then t else e.
  Proof.
    destruct i; destruct i'; reflexivity.
  Qed.

  Lemma refineEquiv_if A :
    forall (f : bool -> Comp A) (b : bool) ta ea,
      refineEquiv (f true) ta
      -> refineEquiv (f false) ea
      -> refineEquiv (f b) (If b Then ta Else ea).
  Proof.
    destruct b; simpl; auto.
  Qed.

  Lemma refine_if_P A :
    forall Pc (Pt Pa : Ensemble A),
      refine { a | (Pc -> Pt a) /\ Pa a}
             (b <- {b | Pc -> b = true};
              If b Then { a | Pt a /\ Pa a}
                 Else { a | Pa a}).
  Proof.
    intros * v Comp_v; computes_to_inv.
    computes_to_econstructor; intuition; subst;
    try destruct v0; simpl in *; eauto; computes_to_inv; intuition.
  Qed.

  Definition decides (b : bool) (P : Prop)
    := If b Then P Else ~ P.

  Lemma refine_iff_P A :
    forall (Pc : Prop) (Pt Pa : Ensemble A),
      refine { a | (Pc -> Pt a) /\ Pa a}
             (b <- {b | decides b Pc};
              If b Then { a | Pt a /\ Pa a}
                 Else { a | Pa a}).
  Proof.
    intros * v Comp_v; computes_to_inv.
    computes_to_econstructor; intuition; subst;
    try destruct v0; simpl in *; eauto; computes_to_inv; intuition.
  Qed.

  Lemma refine_if A :
    forall (c : Comp A) (b : bool) ta ea,
      (b = true -> refine c ta)
      -> (b = false -> refine c ea)
      -> refine c (If b Then ta Else ea).
  Proof.
    destruct b; simpl; auto.
  Qed.

  Lemma refineEquiv_Pick_if {A B}
        (P : A -> B -> Prop) :
    forall (cond : bool) (ta ea : A) ta' ea',
      refineEquiv {b | P ta b} ta'
      -> refineEquiv {b | P ea b} ea'
      -> refineEquiv {b | P (If cond Then ta Else ea) b}
                     (If cond Then ta' Else ea').
  Proof.
    intros; setoid_rewrite refineEquiv_if with
            (f := fun cond : bool => {b : B | P (If cond Then ta Else ea) b}); eauto.
    reflexivity.
  Qed.

  Lemma refineEquiv_if_ret {A}
    : forall (cond : bool) (ta ea : A),
      refineEquiv (ret (If cond Then ta Else ea))
                  (If cond Then ret ta Else ret ea).
  Proof.
    split; destruct cond; reflexivity.
  Qed.

  Definition refine_split_ex A B
             (P : A -> Prop) (P' : A -> B -> Prop)
    : @refine _
              { b | exists a, P a /\ P' a b }%comp
              (a <- { a | P a};
               { b | P' a b })%comp.
  Proof.
    t_refine.
  Qed.

  Definition refineEquiv_pick_ex_computes_to_and A B
             (c : Comp A)
             (P : A -> B -> Prop)
    : refineEquiv { b | exists a, c ↝ a /\ P a b}
                  (a <- c;
                   { b | P a b}).
  Proof. t_refine. Qed.

  Definition refineEquiv_pick_ex_computes_to_bind_and A B D
             (c : B -> Comp A)
             (cB : Comp B)
             (P : A -> D -> Prop)
    : refineEquiv { d | exists a, (b <- cB; c b) ↝ a /\ P a d}
                  (b <- cB;
                   { d | exists a, c b ↝ a /\ P a d}).
  Proof.
    split; intros * v Comp_v; computes_to_inv; destruct_ex; intuition;
    t_refine. Qed.

  Lemma refine_Pick_If_Then_Opt {A B}
    : forall (P : Ensemble B) (c e : Comp A) (t : B -> Comp A),
      (forall b, P b -> refine c (t b))
      -> (refine c e)
      -> refine c
                (b <- {b | forall b',
                          (b = Some b' -> P b')};
                 Ifopt b as b' Then t b' Else e).
  Proof.
    t_refine.
    destruct v0; t_refine.
    eapply H; eauto.
  Qed.

  Lemma refine_Pick_Some_dec {A B}
    : forall (P : Ensemble B) (c e : Comp A) (t : B -> Comp A),
      (forall b, P b -> refine c (t b))
      -> ((forall b, ~ P b) -> refine c e)
      -> refine c
                (b <- {b | forall b',
                          (b = Some b' -> P b')
                          /\ (b = None -> forall b', ~ P b')};
                 Ifopt b as b' Then t b' Else e) .
  Proof.
    unfold refine; intros; computes_to_inv;
    destruct v0; simpl in *; computes_to_inv; eauto.
    eapply H; eauto.
    eapply H1; eauto.
    eapply H0; eauto; intros; eapply H1; eauto.
  Qed.

  Add Morphism
      (decides)
      with signature (eq ==> iff ==> iff)
        as decide_eq_iff_iff_morphism.
  Proof.
    unfold decides; intros b p1 p2 equiv.
    destruct b; simpl; intuition.
  Qed.

  Lemma refine_pick_decides {A}
    : forall (P : Prop) (c e : Comp A) (t : Comp A),
      (P -> refine c t)
      -> (~ P -> refine c e)
      -> refine c
                (b <- {b | decides b P};
                 If b Then t Else e).
  Proof.
    unfold refine; intros; computes_to_inv;
    destruct_ex; split_and; computes_to_inv.
    destruct v0; simpl in *; eauto.
  Qed.

  Lemma refine_pick_decides_branches {A}
        (P : Prop)
        (Q Q' : Ensemble A)
        (q q' : Comp A)
    : (P -> refine {a | Q a} q)
      -> (~ P -> refine {a | Q' a} q')
      -> refine {a | (P -> Q a) /\
                     (~ P -> Q' a)}
                (b <- {b | decides b P};
                 If b Then q Else q').
  Proof.
    intros; eapply refine_pick_decides.
    - unfold refine; intros; computes_to_inv;
      econstructor; intuition; eapply H4; eauto.
    - unfold refine; intros; computes_to_inv;
      econstructor; intuition; eapply H4; eauto.
  Qed.

  Lemma refine_pick_decides' {A}
        (P : Prop)
        (Q Q' : Ensemble A)
    : refine {a | (P -> Q a) /\
                  (~ P -> Q' a)}
             (b <- {b | decides b P};
              If b Then
                 {a | Q a}
                 Else
                 {a | Q' a}).
  Proof.
    eapply refine_pick_decides_branches; reflexivity.
  Qed.

  Global Add Parametric Morphism : decides
      with signature eq ==> iff ==> iff
        as decides_mor.
  Proof.
    repeat intro; split_iff; destruct_head bool; simpl; tauto.
  Qed.

  Lemma refineEquiv_decides_eqb (b b1 b2 : bool)
    : decides b (b1 = b2) <-> b = If b2 Then b1 Else negb b1.
  Proof.
    destruct_head bool; simpl; intuition.
  Qed.

  Lemma refine_If_Then_Else_Bind {A B}
    : forall i (t e : Comp A) (b : A -> Comp B),
      refine (a <- If i Then t Else e; b a)
             (If i Then (a <- t;
                         b a)
                 Else (a <- e;
                       b a)).
  Proof.
    intros; destruct i; simpl; reflexivity.
  Qed.

  Lemma refine_if_Then_Else_Duplicate {A}
    : forall (i : bool) (t e e' : Comp A),
      refine (if i then (if i then t else e') else e)
             (if i then t else e).
  Proof.
    destruct i; simpl; reflexivity.
  Qed.

  Lemma refine_If_Then_Else_Duplicate {A}
    : forall i (t e e' : A),
      If i Then (If i Then t Else e') Else e =
      If i Then t Else e.
  Proof.
    destruct i; simpl; reflexivity.
  Qed.

  Lemma refine_If_Opt_Then_Else_Bind {A B C}
    : forall (i : option A)
             (t : A -> Comp B)
             (e : Comp B)
             (c : B -> Comp C),
      refine (b <- If_Opt_Then_Else i t e; c b)
             (If_Opt_Then_Else i (fun a' =>
                                    b <- t a';
                                  c b)
                               (b <- e;
                                c b)).
  Proof.
    intros; destruct i; simpl; reflexivity.
  Qed.

  Lemma refine_trivial_if_then_else :
    forall x,
      refine
        (If_Then_Else x (ret true) (ret false))
        (ret x).
  Proof.
    destruct x; reflexivity.
  Qed.

  Lemma decides_True :
    refine {b | decides b True}%comp
           (ret true).
  Proof.
    computes_to_econstructor;  computes_to_inv; subst; simpl; auto.
  Qed.

  Lemma decides_True_Pre (Q : Prop) :
    refine {b | Q -> decides b True}%comp
           (ret true).
  Proof.
    computes_to_econstructor;  computes_to_inv; subst; simpl; auto.
  Qed.

  Lemma decides_2_True (A : Type) (B : A -> Type) :
    refine {b' | decides b' (forall a : A, B a -> True)}%comp
           (ret true).
  Proof.
    computes_to_econstructor;  computes_to_inv; subst; simpl; auto.
  Qed.

  Lemma decides_3_True (A B : Type) (C : B -> Type) :
    refine {b' | decides b' (A -> forall b : B, C b -> True)}%comp
           (ret true).
  Proof.
    computes_to_econstructor;  computes_to_inv; subst; simpl; auto.
  Qed.

  Lemma decides_neq (A : Type) (B : Prop) (a : A) :
    refine {b' | decides b' (a <> a -> B)}%comp
           (ret true).
  Proof.
    computes_to_econstructor;  computes_to_inv; subst; simpl; auto.
    congruence.
  Qed.

  Lemma Bind_refine_If_Then_Else {A B}
    : forall i (t e : A -> Comp B) (ca : Comp A),
      refine (a <- ca;
              If i Then t a Else e a)
             (If i Then (a <- ca;
                         t a)
                 Else (a <- ca;
                       e a)).
  Proof.
    intros; destruct i; simpl; reflexivity.
  Qed.

  Lemma Bind_refine_If_Opt_Then_Else {A B C}
    : forall i (t : A -> B -> Comp C) (e : A -> Comp C) (ca : Comp A),
      refine (a <- ca;
              Ifopt i as b Then t a b Else e a)
             (Ifopt i as b Then (a <- ca;
                                 t a b)
                           Else (a <- ca;
                                 e a)).
  Proof.
    intros; destruct i; simpl; reflexivity.
  Qed.

  Lemma refineEquiv_swap_bind {A B C} (c1 : Comp A) (c2 : Comp B) (f : A -> B -> Comp C)
    : refineEquiv (a <- c1; b <- c2; f a b) (b <- c2; a <- c1; f a b).
  Proof.
    split; repeat intro;
    computes_to_inv;
    repeat (econstructor; try eassumption).
  Qed.

  Lemma refine_bind_dedup {A B} (c1 : Comp A) (f : A -> A -> Comp B)
    : refine (a <- c1; b <- c1; f a b) (a <- c1; f a a).
  Proof.
    repeat intro;
    computes_to_inv;
    repeat (econstructor; try eassumption).
  Qed.

  Lemma comp_split_snd {A B} (x : A * B)
    : refineEquiv (ret (snd x))
                  (ab <- ret x;
                   ret (snd ab)).
  Proof.
    autorewrite with refine_monad; reflexivity.
  Qed.

  Lemma refine_skip {A B C} (c : Comp A) (f : A -> Comp B) (dummy : A -> Comp C)
    : refine (Bind c f)
             (a <- c;
              dummy a>>
                    f a).
  Proof.
    repeat first [ intro
                 | computes_to_inv
                 | computes_to_econstructor; eassumption
                 | computes_to_econstructor; try eassumption; [] ].
    computes_to_econstructor; eauto.
  Qed.

  Lemma refine_skip2 {A B} (a : Comp A) (dummy : Comp B)
    : refine a
             (dummy>>
                   a).
  Proof.
    repeat first [ intro
                 | computes_to_inv
                 | assumption
                 | econstructor; eassumption
                 | econstructor; try eassumption; [] ].
    eauto.
  Qed.

  Lemma decides_negb :
    forall b P,
      decides (negb b) P -> decides b (~ P).
  Proof.
    unfold decides; setoid_rewrite if_negb; simpl; intros.
    destruct b; simpl in *; intuition.
  Qed.

  Lemma refine_decide_not :
    forall {A} (P: A -> Prop),
      refine (Pick (fun (b : bool) =>
                      decides b (forall (x: A), ~ P x)))
             (Pick (fun (b : bool) =>
                      decides (negb b) (exists (x: A), P x))).
  Proof.
    unfold refine; intros; computes_to_inv.
    computes_to_constructor.
    rewrite <- not_exists_forall; apply decides_negb;
    assumption.
  Qed.

  Lemma refine_decide_negb :
    forall P,
      refineEquiv (Pick (fun b => decides (negb b) P))
                  (Bind (Pick (fun b => decides b P))
                        (fun b => ret (negb b))).
  Proof.
    unfold refineEquiv, refine; simpl;
    split; intros; computes_to_inv;
    subst; repeat computes_to_econstructor; eauto; rewrite Bool.negb_involutive;
    [ assumption | constructor ].
  Qed.

  Lemma refine_decide_negation :
    forall {A} (P: A -> Prop),
      refine (Pick (fun (b : bool) =>
                      decides b (forall (x: A), ~ P x)))
             (Bind (Pick (fun b => decides b (exists (x: A), P x)))
                   (fun b => ret (negb b))).
  Proof.
    intros;
    rewrite refine_decide_not, refine_decide_negb; reflexivity.
  Qed.

  Lemma eq_ret_compute :
    forall (A: Type) (x y: A), x = y -> ret x ↝ y.
  Proof.
    intros; subst; apply ReturnComputes; trivial.
  Qed.

  Lemma refine_eq_ret :
    forall {A} (a a': A),
      a = a' ->
      refineEquiv  (ret a) (ret a').
  Proof.
    intros; subst; reflexivity.
  Qed.

  Lemma refine_snd :
    forall {A B: Type} (P: B -> Prop),
      refine
        { pair | P (snd pair) }
        (_fst <- Pick (fun (x: A) => True);
         _snd <- Pick (fun (y: B) => P y);
         ret (_fst, _snd)).
  Proof.
    t_refine.
  Qed.

  Lemma refine_let :
    forall {A B : Type} (PA : A -> Prop) (PB : B -> Prop),
      refineEquiv (Pick (fun x: A * B  =>  let (a, b) := x in PA a /\ PB b))
                  (a <- {a | PA a};
                   b <- {b | PB b};
                   ret (a, b)).
  Proof.
    t_refine.
  Qed.

  Lemma refine_ret_eq :
    forall {A: Type} (a: A) b,
      b = ret a -> refine (ret a) (b).
  Proof.
    t_refine.
  Qed.

  Lemma refine_eqA_into_ret :
    forall {A: Type} {eqA: list A -> list A -> Prop},
      Reflexive eqA ->
      forall (comp : Comp (list A)) (impl result: list A),
        comp = ret impl -> (
          comp ↝ result ->
          eqA result impl
        ).
  Proof.
    intros; subst; computes_to_inv; subst; trivial.
  Qed.

  Lemma refine_pick_val' :
    forall {A : Type} (a : A)  (P : A -> Prop),
      P a -> refine (Pick P) (ret a).
  Proof.
    intros; apply refine_pick_val; assumption.
  Qed.

  Lemma refine_If_Then_Else_ret {A} :
    forall i (t e : A),
      refine (@If_Then_Else (Comp A) i (ret t) (ret e))
             (ret (@If_Then_Else A i t e)).
  Proof.
    destruct i; reflexivity.
  Qed.

  Lemma refine_If_Opt_Then_Else_ret {A B} :
    forall i (t : A -> B) (e : B),
      refine (@If_Opt_Then_Else A (Comp B) i (fun a => ret (t a)) (ret e))
             (ret (@If_Opt_Then_Else A B i t e)).
  Proof.
    destruct i; reflexivity.
  Qed.

End general_refine_lemmas.

Tactic Notation "finalize" "refinement" :=
  eexists; solve [ reflexivity | eapply reflexivityT ].

Tactic Notation "refine" "using" constr(refinement_rule) :=
  eapply refinement_step;
  [progress setoid_rewrite refinement_rule; try reflexivity | ].

(* 'Wrapper' tactics for various refinements *)

Tactic Notation "refine" "pick" "pair" :=
  rewrite refine_pick_pick;
  [ |
    let x := fresh in
    let H := fresh in
    intros x H;
      let A := match goal with
                 |- ?A /\ ?B => constr:(A)
               end in
      let B := match goal with
                 |- ?A /\ ?B => constr:(B)
               end in
      match eval pattern (fst x) in A with
        ?A' _ =>
        match eval pattern (snd x) in B with
          ?B' _ =>
          match type of H with
          | ?C _ => unify C (fun x => A' (fst x) /\ B' (snd x));
              exact H
          end
        end
      end ]; rewrite refineEquiv_pick_pair.

Lemma refine_under_bind_both_trans {A B}
  : forall c c' c'' (x y : A -> Comp B),
    refine c c' ->
    (forall a, c ↝ a -> refine (x a) (y a))
    -> refine (a <- c'; y a) c''
    -> refine (a <- c; x a) c''.
Proof. intros; etransitivity; eauto using refine_under_bind_both. Qed.

Lemma refine_If_Then_Else_trans {A}
  : forall (c : bool) (x y z : Comp A),
    refine x y ->
    forall x0 y0 : Comp A,
      refine x0 y0
      -> refine (If c Then y Else y0) z
      -> refine (If c Then x Else x0) z.
Proof. intros; etransitivity; eauto using refine_If_Then_Else. Qed.

Lemma refine_If_Opt_Then_Else_trans {A B}
  : forall (c : option B) (e e' z: Comp A) (t t' : B -> Comp A),
    (forall b, refine (t b) (t' b))
    -> (refine e e')
    -> refine (Ifopt c as b' Then t' b' Else e') z
    -> refine (Ifopt c as b' Then t b' Else e) z.
Proof.
  intros; etransitivity; eauto using refine_If_Opt_Then_Else.
Qed.

Local Ltac t2 p := intros; destruct p; intuition.

Lemma refine_sumbool_match :
  forall `(P : {A} + {~A}) B
         (f f' : A -> Comp B) (g g' : ~A -> Comp B),
       pointwise_relation A    refine f f'
    -> pointwise_relation (~A) refine g g'
    -> refine (match P with
               | left  H => f H
               | right H => g H
               end)
              (match P with
               | left  H => f' H
               | right H => g' H
               end).
Proof. t2 P. Qed.

Lemma refine_sumbool_ret :
  forall `(P : {A} + {~A}) `(f : A -> B) (g : ~A -> B),
    refine (match P with
            | left  H => ret (f H)
            | right H => ret (g H)
            end)
           (ret (match P with
                 | left  H => f H
                 | right H => g H
                 end)).
Proof. t2 P. Qed.

Lemma refine_sumbool_bind :
  forall `(P : {A} + {~A})
         `(f : A -> Comp B) (g : ~A -> Comp B)
         `(h : B -> Comp C),
    refine (x <- match P with
                 | left  H => f H
                 | right H => g H
                 end;
            h x)
           (match P with
            | left  H => x <- f H; h x
            | right H => x <- g H; h x
            end).
Proof. t2 P. Qed.

Lemma refine_bind_sumbool :
  forall `(P : {A} + {~A})
         `(f : C -> A -> Comp B) (g : C -> ~A -> Comp B)
         `(c : Comp C),
    refine (x <- c;
            match P with
            | left  H => f x H
            | right H => g x H
            end)
           (match P with
            | left  H => x <- c; f x H
            | right H => x <- c; g x H
            end).
Proof. t2 P. Qed.

Tactic Notation "refine" "pick" "val" open_constr(v) :=
  let T := type of v in
  rewrite refine_pick_val with
  (A := T)
    (a := v) at 1.

Tactic Notation "refine" "pick" "eq" :=
  match goal with
  | |- context[Pick (fun x => x = _)] =>
    setoid_rewrite refine_pick_eq || setoid_rewrite refineEquiv_pick_eq
  | |- context[Pick (fun x => _ = x)] =>
    setoid_rewrite refine_pick_eq' || setoid_rewrite refineEquiv_pick_eq'
  end.


Require Import Coq.Lists.List.

Lemma nil_neq_app_cons : forall { A: Type } ( x : A ) l1 l2,
    nil <> l1 ++ x :: l2.
Proof.
  intros. induction l1; simpl; intro; inversion H.
Qed.


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

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

Print Assumptions mul.


Require
  MathClasses.misc.JMrelation.
Require Import
  MathClasses.interfaces.abstract_algebra MathClasses.interfaces.functors MathClasses.theory.categories.

Record Object := object
  { obj:> Type
  ; Arrows_inst: Arrows obj
  ; Equiv_inst: ∀ x y: obj, Equiv (x ⟶ y)
  ; CatId_inst: CatId obj
  ; CatComp_inst: CatComp obj
  ; Category_inst: Category obj }.

Arguments object _ {Arrows_inst Equiv_inst CatId_inst CatComp_inst Category_inst}.
#[global]
Existing Instance Arrows_inst.
#[global]
Existing Instance Equiv_inst.
#[global]
Existing Instance CatId_inst.
#[global]
Existing Instance CatComp_inst.
#[global]
Existing Instance Category_inst.

Record Arrow (x y: Object): Type := arrow
  { map_obj:> obj x → obj y
  ; Fmap_inst: Fmap map_obj
  ; Functor_inst: Functor map_obj _ }.

Arguments arrow {x y} _ _ _.
Arguments map_obj {x y} _ _.
#[global]
Existing Instance Fmap_inst.
#[global]
Existing Instance Functor_inst.

#[global]
Instance: Arrows Object := Arrow.
(* Hint Extern 4 (Arrows Object) => exact Arrow: typeclass_instances. *)

Section contents.
  Section more_arrows.
    Context (x y: Object).

    Global Instance e: Equiv (x ⟶ y) := λ a b,
      (∀ v, a v ≡ b v) ∧
      (∀ `(f: v ⟶ w), JMrelation.R (=) (fmap a f) _ (=) (fmap b f)).

    Instance e_refl: Reflexive e.
    Proof.
     intro a. unfold e. intuition.
     apply JMrelation.reflexive, _.
    Qed.

    Instance e_sym: Symmetric e.
    Proof with intuition.
     unfold e. intros ?? [P Q]...
     apply JMrelation.symmetric...
    Qed.

    Instance e_trans: Transitive e.
    Proof with intuition.
     unfold e. intros a b c [P Q] [R S]...
      transitivity (b v)...
     apply JMrelation.transitive with _ (=) (fmap b f)...
    Qed.

    Global Instance: Setoid (x ⟶ y).
    Proof. split; apply _. Qed.
  End more_arrows.

  Global Instance: CatId Object := λ _, arrow id _ _.

  Global Program Instance: CatComp Object := λ _ _ _ x y, arrow (x ∘ y) _ _.

  Global Instance: ∀ x y z: Object, Proper ((=) ==> (=) ==> (=)) ((◎): (y ⟶ z) → (x ⟶ y) → (x ⟶ z)).
  Proof with intuition; try apply _.
   unfold equiv, e.
   intros x y z a b [P Q] c d [R S].
   split; intros.
    change (a (c v) ≡ b (d v)). congruence.
   change (JMrelation.R (=) (fmap a (fmap c f)) _ (=) (fmap b (fmap d f))).
   apply JMrelation.transitive with _ (=) (fmap a (fmap d f))...
   specialize (S _ _ f). revert S.
   generalize (fmap c f) (fmap d f).
   repeat rewrite R.
   intros. apply JMrelation.relate.
   rewrite (JMrelation.unJM _ _ _ _ _ S)... (* <- uses K *)
  Qed.

  Global Instance: Category Object.
  Proof. repeat (split; try apply _); intuition; apply reflexivity. Qed.
End contents.


Require Import ExtLib.Structures.Monads.
Require Import BinPos.

Set Implicit Arguments.
Set Strict Implicit.

Inductive FixResult (T : Type) :=
| Term : T -> FixResult T
| Diverge : FixResult T.

Arguments Diverge {_}.
Arguments Term {_} _.

(** The GFix monad is like monad fix except that
    it encapsulates the "gas" that is used as the measure
 **)
Section gfix.
  (** This is essentially ReaderT (optionT m)) **)
  Inductive GFix (T : Type) : Type := mkGFix
  { runGFix : N -> FixResult T }.

  Global Instance MonadFix_GFix : MonadFix GFix :=
  { mfix := fun T U f v => mkGFix (fun n : N =>
    match n with
      | N0 => Diverge
      | Npos g =>
        let F := fix rec (acc : T -> FixResult U) (gas : positive) (x : T)
          : FixResult U :=
          match gas return FixResult U with
            | xI p =>
              runGFix (f (fun x => mkGFix (fun n => rec (fun x => rec acc p x) p x)) x) n
            | xO p =>
              rec (fun x => rec acc p x) p x
            | xH => runGFix (f (fun x => mkGFix (fun _ => acc x)) x) n
          end
        in F (fun x => runGFix (f (fun _ => mkGFix (fun _ => Diverge)) x) n) g v
      end)
  }.

  Global Instance Monad_GFix : Monad GFix :=
  { ret := fun _ v => mkGFix (fun _ => Term v)
  ; bind := fun _ _ c1 c2 =>
    mkGFix (fun gas =>
      match runGFix c1 gas with
        | Diverge => Diverge
        | Term v => runGFix (c2 v) gas
      end)
  }.

End gfix.

(** Demo
Require Import ExtLib.Data.Monads.IdentityMonad.
Definition foo : nat -> GFix ident nat :=
  mfix (fun recur n =>
    match n with
      | 0 => ret 0
      | S n => recur n
    end).

Eval compute in runGFix (foo 10) 100000000000000000000000.
**)


From discprob.basic Require Import base order nify seq_ext.
From discprob.prob Require Import prob countable finite stochastic_order.
From discprob.monad.finite Require Import monad monad_hoare bst bst_correct permutation.
From discprob.monad.finite Require Import max_triangle.
From discprob.rec Require Import rec_convert quicksort_rec.
From mathcomp Require Import ssreflect ssrbool ssrfun eqtype ssrnat seq div choice fintype.
From mathcomp Require Import tuple finfun bigop prime binomial finset.
From mathcomp Require Import path.
Require Import Reals Fourier Psatz Omega.

Definition h (l: seq nat) :=
  match l with
  | [::] => mret ([::], [::])
  | [::a] => mret ([::], [::])
  | a :: l =>
      p ← draw_next a l;
      mret ([seq i <- a :: l | i < sval p], [seq i <- a :: l | i > sval p])%nat
  end.

Definition rsize (l: seq nat) := INR (size l).

Definition m x :=
  match (Rlt_dec x (4/3)) with
  | left _ => 0
  | _ => x * (3 / 4)
  end.

Lemma uniq_unif n:
  uniq (outcomes (unif n)).
Proof.
  rewrite /unif//=.
  rewrite map_inj_uniq.
  - apply enum_uniq.
  - intros (?&?) (?&?). inversion 1; auto.
    subst; f_equal. apply bool_irrelevance.
Qed.

Lemma uniq_unif2 n:
  uniq ([seq i.2 | i <- outcomes (unif n)]).
Proof.
  rewrite /unif//=.
  rewrite map_inj_in_uniq.
  - apply uniq_unif.
  - intros (?&?) (?&?).
    move /mapP => [? ? Heq1]. inversion Heq1; subst.
    move /mapP => [? ? Heq2]. inversion Heq2; subst.
    rewrite //=. inversion 1. repeat f_equal.
    apply val_inj => //=.
Qed.

Lemma Ex_max_partition_sum l:
  Ex (rvar_comp (rvar_of_ldist (h l)) (λ x, Rmax (rsize (x.1)) (rsize (x.2)))) <=
  match (size l) with
     | O => 0
     | _ => \big[Rplus/0]_(i < size l) (/INR (size l) * (Rmax (INR i) (INR (size l - i - 1))))
  end.
Proof.
  destruct l as [| a l].
  - rewrite Ex_fin_comp.
    rewrite -(big_map _ (λ x, true) (λ a, pr_eq _ a * _ (_ (fst a)) (_ (snd a)))).
    rewrite img_rvar_of_ldist//=/img/undup//= !big_cons ?big_nil.
    rewrite pr_mret_simpl //= /m/rsize//=.
    rewrite Rmax_left //; last (fourier).
    rewrite /INR. ring_simplify. fourier.
  - destruct l as [| b0 l0].
    {
      rewrite Ex_fin_comp.
      rewrite -(big_map _ (λ x, true) (λ a, pr_eq _ a * _ (_ (fst a)) (_ (snd a)))).
      rewrite img_rvar_of_ldist//=/img/undup//= !big_cons ?big_nil.
      rewrite pr_mret_simpl //= /m/rsize//=.
      rewrite Rmax_left //; last (fourier).
      replace (INR 0) with 0 by auto. ring_simplify.
      apply Rle_big0 => i _ //=.
      destruct i as (i&Hlt) => //=. assert (i = O) as -> by (nify; omega).
      replace (1 - 0 - 1)%nat with O by (nify; omega).
      replace (INR 0) with 0 by auto. rewrite Rmax_left; last fourier.
      replace (INR 1) with 1 by auto. fourier.
    }
    remember (b0 :: l0) as l eqn:Heql0.

    replace (size (a :: l)) with (S (size l)) by auto.
    rewrite -(Ex_mbind_mret (h (a :: l))).
    rewrite /h Heql0. rewrite -Heql0. clear l0 Heql0.

    rewrite ldist_assoc. rewrite ldist_assoc.
    rewrite Ex_mbind_ldist2.
    rewrite undup_id; last first.
    { apply uniq_unif2. }
    edestruct (quicksort_cost.perm_eq_bij a) as (f&Hfspec&Hfsize&Hinv&Hinj).
    { rewrite perm_sym. apply (permEl (perm_sort leq (a :: l))). }
    eapply sum_reidx_map_le with
        (h := (λ (H : {x : nat | (x <= size l)%nat}),
               let (x, i) := H in
                (Ordinal (n:=(size l).+1) (m:=f x) (Hfsize x i)))).
    * intros (x, Hle) Hin.
      rewrite unif_pr. rewrite /Rdiv Rmult_1_l. apply Rmult_le_compat_l.
      ** left. apply Rinv_0_lt_compat. apply lt_0_INR. rewrite //=. omega.
      ** rewrite -ldist_assoc. rewrite Ex_mbind_mret.
         apply Ex_bound.
         rewrite ?ldist_left_id.
         assert (Hlt: (x < size (a :: l))%nat) by auto.
         rewrite //= in Hlt.
         apply mspec_mret. rewrite /sval/nat_of_ord.
         remember (a :: l) as l0 eqn:Heql0.
         rewrite -Heql0.
         assert (Hsize: S (size l) = (size l0)) by (rewrite Heql0; auto).
         rewrite /rsize.
         rewrite !Rmax_INR. apply le_INR.
         apply Nat.max_le_compat; rewrite /fst/snd.
         *** apply /leP. rewrite Hfspec.
             specialize (permEl (perm_sort leq l0)) => Hperm.
             rewrite //= -(perm_size (quicksort_cost.perm_filter _ _ Hperm)).
             apply quicksort_cost.lower_split_size_sorted_nth; auto.
             **** apply sort_sorted => ??. apply leq_total.
             **** rewrite (perm_size (permEl (perm_sort leq l0))) /=.
                  apply Hfsize. rewrite Heql0. done.
             **** rewrite Heql0 //=.
         *** apply /leP. rewrite Hfspec.
             assert (Hsize': (size l).+1 = size l0) by rewrite Heql0 //.
             rewrite Hsize'.
             specialize (permEl (perm_sort leq l0)) => Hperm.
             rewrite -(perm_size (permEl (perm_sort leq l0))).
             rewrite -(perm_size (quicksort_cost.perm_filter _ _ Hperm)).
             apply quicksort_cost.upper_split_size_sorted_nth; auto.
             **** apply sort_sorted => ??. apply leq_total.
             **** rewrite (perm_size (permEl (perm_sort leq l0))) /=.
                  apply Hfsize. rewrite Heql0. done.
             **** rewrite Heql0 //=.
      ** apply unif_all.
    * intros (?&?) Hin; split; auto.
      rewrite /index_enum. rewrite -enumT. rewrite mem_enum => //=.
    * intros b Hin _ Hfalse. contradiction Hfalse.
      destruct b as (m, Hle).
      destruct (Hinv m Hle) as (x'&Hle'&Heq).
      exists (exist _ x' Hle'); repeat split; auto.
      ** apply unif_all.
      ** apply ord_inj => //=.
    * apply uniq_unif2.
    * rewrite /index_enum. rewrite -enumT. apply enum_uniq.
    * intros (?&?) (?&?) _ => //=. inversion 1; subst. apply /val_inj/Hinj => //=.
Qed.

Lemma Ex_max_partition l:
  Ex (rvar_comp (rvar_of_ldist (h l)) (λ x, Rmax (rsize (x.1)) (rsize (x.2)))) <=
   m (rsize l).
Proof.
  eapply Rle_trans; [ apply Ex_max_partition_sum | apply quicksort_cost.sum_le_m ].
Qed.

Definition T := λ x, rvar_comp (rvar_of_ldist (add_list_random x Leaf)) (λ x, INR (height x)).
  Definition a x :=
    match Rle_dec x 1 with
      | left _ => 0
      | _ =>
        match Rlt_dec x 2 with
        | left _ => (x - 1)
        | _ => 1
        end
    end.
Definition h' x := rvar_of_ldist (h x).

Lemma rand_tree_non_empty_non_leaf l:
  mspec (rand_tree_rec l) (λ t, l ≠ [::] → t ≠ Leaf).
Proof.
  destruct l; [| destruct l]; first congruence.
  - rewrite rt_unfold. apply mspec_mret; congruence.
  - rewrite rt_unfold. do 3 (tbind (λ x, true) => //; intros). apply mspec_mret; congruence.
Qed.

Lemma filter_lt_gt_non_empty l x:
  (size l > 1)%nat →
  uniq l →
  [seq i <- l | (i < x)%nat] ≠ [::] ∨
  [seq i <- l | (x < i)%nat] ≠ [::].
Proof.
  intros Hsize Huniq.
  induction l as [| a l]; first by auto.
  induction l as [| b l]; first by auto.
  rewrite //=.
  repeat case: ifP;  try (left; congruence); try (right; congruence).
  move /ltP. intros Hlt1.
  move /ltP. intros Hlt2.
  move /ltP. intros Hlt3.
  move /ltP. intros Hlt4.
  assert (a = x).
  { omega. }
  assert (b = x).
  { omega. }
  subst. rewrite //= in Huniq. exfalso. move /andP in Huniq.
  destruct Huniq as (Hfalse&?). move /negP in Hfalse. apply Hfalse.
  by rewrite in_cons eq_refl.
Qed.

Lemma huniq l n:
  uniq l → uniq (fst ((h' l) n)) ∧ uniq (snd ((h' l) n)).
Proof.
  intros Huniq. rewrite //=.
  cut (∀ n, uniq ((fst (rvar_of_ldist (h l) n))) ∧
          uniq ((snd (rvar_of_ldist (h l) n)))); first by rewrite //=.
  intros n'. apply fun_to_mspec.
  rewrite /h.
  destruct l as [| a l].
  { apply mspec_mret => //=. }
  destruct l as [| b l0].
  { apply mspec_mret => //=. }
  remember (b::l0) as l eqn:Heql0.
  clear b l0 Heql0.
  tbind (λ x, true).
  { intros (?&?) => //. }
  intros ??.
  apply mspec_mret.
  split; apply filter_uniq => //.
Qed.

Lemma Trec:
    ∀ x r, uniq x → pr_eq (T x) r =
            \big[Rplus/0]_(x' : imgT (h' x))
                        (pr_eq (h' x) (sval x') *
                         pr_eq (rvar_comp (rvar_pair (T (fst (sval x'))) (T (snd (sval x'))))
                                         (fun xy => Rmax (fst xy) (snd xy))) (r - a (rsize x))).
Proof.
  eapply eq_dist_rec; [ apply huniq | | ]; last first.
  { rewrite /T. intros. apply eq_dist_comp. eapply eq_dist_trans; first apply alr_rand_perm.
    apply eq_dist_sym, rt_gen_perm. done. }
  intros ls r.
  induction ls as [| a0 ls] => Huniq.
  - rewrite /T. rewrite -pr_mbind_mret.
    rewrite -(big_map _ (λ x, true)
     (λ x',
     (pr_eq (h' [::]) x' *
      pr_eq (rvar_comp
               (rvar_pair
                  (rvar_comp (rvar_of_ldist (rand_tree_rec x'.1)) (λ x, INR (height x)))
                  (rvar_comp (rvar_of_ldist (rand_tree_rec x'.2)) (λ x, INR (height x))))
               (λ xy : prod_eqType R_eqType R_eqType, Rmax xy.1 xy.2)) (r - a (rsize [::]))))).
    rewrite img_rvar_of_ldist /h.
    rewrite big_cons big_nil Rplus_0_r.
    rewrite /a/rsize/size.
    replace (INR 0) with 0 by auto.
    destruct (Rlt_dec); last nra.
    specialize (pr_rvar_ldist ((x ← rand_tree_rec [::]; mret (INR (height x))))) => -> //=.
    rewrite big_cons big_nil /=.
    rewrite {2}/pr_eq. rewrite pr_eq_alt_comp.
    destruct (Rle_dec); last nra.

    rewrite -(big_map _ (λ x, true) (λ v, if Rmax (fst v) (snd v) == r - 0 then
                                            pr_eq _ v
                                          else
                                            0)).
    rewrite (perm_big _ (img_pair_rv _ _ _ _)).
    rewrite allpairs_comp.
    rewrite img_rvar_comp map_comp (map_comp height).
    specialize (img_rvar_of_ldist (rand_tree_rec [::])) => ->.
    rewrite ?big_cons ?big_nil ?Rplus_0_r //= ?Rminus_0_r.
    rewrite Rmax_left; last reflexivity.
    case: ifP; last by (intros; field).
    rewrite pr_mret_simpl //=. intros. ring_simplify. rewrite pr_eq_rvar_pair //=.
    rewrite /pr_eq pr_eq_alt_comp.
    rewrite -(big_map _ (λ x, true) (λ v, if INR (height (v)) == 0 then pr_eq _ v else 0)).
    rewrite -(big_map _ (λ x, true) (λ v, if INR (height (v)) == 0 then pr_eq _ v else 0)).
    specialize (img_rvar_of_ldist (rand_tree_rec [::])) => -> //=.
    rewrite ?big_cons ?big_nil //= eq_refl.
    rewrite pr_mret_simpl //=. field.
  - destruct ls as [| b0 ls0].
    { rewrite /T. rewrite -pr_mbind_mret.
    rewrite -(big_map _ (λ x, true)
     (λ x',
     (pr_eq (h' _) x' *
      pr_eq (rvar_comp
               (rvar_pair
                  (rvar_comp (rvar_of_ldist (rand_tree_rec x'.1)) (λ x, INR (height x)))
                  (rvar_comp (rvar_of_ldist (rand_tree_rec x'.2)) (λ x, INR (height x))))
               (λ xy : prod_eqType R_eqType R_eqType, Rmax xy.1 xy.2)) (r - a (rsize _))))).
    rewrite img_rvar_of_ldist /h.
    rewrite big_cons big_nil Rplus_0_r.
    rewrite /a/rsize/size.
    replace (INR 1) with 1 by auto.
    destruct (Rlt_dec); last nra.
    specialize (pr_rvar_ldist ((x ← rand_tree_rec [:: a0]; mret (INR (height x))))) => -> //=.
    rewrite big_cons big_nil /=.
    rewrite /h'/h.
    rewrite {2}/pr_eq. rewrite 1!pr_eq_alt_comp.
    rewrite ?big_cons ?big_nil ?Rplus_0_r ?Rminus_0_r.
    rewrite pr_mret_simpl. ring_simplify.
    destruct (Rle_dec); last nra.
    rewrite /Rminus Ropp_0 Rplus_0_r.
    rewrite -(big_map _ (λ x, true) (λ v, if Rmax (fst v) (snd v) == r then
                                            pr_eq _ v
                                          else
                                            0)).
    rewrite (perm_big _ (img_pair_rv _ _ _ _)).
    rewrite allpairs_comp.
    rewrite img_rvar_comp map_comp (map_comp height).
    specialize (img_rvar_of_ldist (rand_tree_rec [::])) => -> //=.
    rewrite ?big_cons ?big_nil //=. replace (INR 0) with 0 by auto.
    rewrite ?Rplus_0_r.
    rewrite Rmax_left; last reflexivity.
    case: ifP; intros; last field.
    rewrite pr_eq_rvar_pair.
    f_equal.
    * rewrite /pr_eq pr_eq_alt_comp.
      rewrite -(big_map _ (λ x, true) (λ v, if INR (height (v)) == 0 then pr_eq _ v else 0)).
      specialize (img_rvar_of_ldist (rand_tree_rec [::])) => -> //=.
      rewrite ?big_cons ?big_nil //= eq_refl.
      rewrite pr_mret_simpl //=. field.
    }
    remember (b0 :: ls0) as ls eqn:Heql0.
   rewrite /T.
    rewrite rt_unfold.
    transitivity
  (pr_eq (rvar_comp (rvar_of_ldist
         ('(ll, lu) ← h (a0 :: ls);
            lls ← rand_tree_rec ll;
            lus ← rand_tree_rec lu;
            mret (lls, lus)))
          (λ x, 1 + Rmax (INR (height (fst x))) (INR (height (snd x))))) r); last first.
    {
      rewrite /h'.
    rewrite -(big_map _ (λ x, true)
     (λ x',
     (pr_eq (h' _) x' *
      pr_eq (rvar_comp
               (rvar_pair
                  (rvar_comp (rvar_of_ldist (rand_tree_rec x'.1)) (λ x, INR (height x)))
                  (rvar_comp (rvar_of_ldist (rand_tree_rec x'.2)) (λ x, INR (height x))))
               (λ xy : prod_eqType R_eqType R_eqType, Rmax xy.1 xy.2)) (r - a (rsize _))))).
    rewrite -pr_mbind_mret ldist_assoc. rewrite pr_mbind_ldist2.
    rewrite img_rvar_of_ldist.
      *  eapply eq_bigr. intros (ll&lu) _.
        f_equal.
        rewrite pr_mbind_mret. rewrite rvar_pair_comp //= rvar_comp_comp //=.
        transitivity (pr_eq (rvar_comp (rvar_pair (rvar_of_ldist (rand_tree_rec ll))
                                                  (rvar_of_ldist (rand_tree_rec lu)))
                                       (λ x : tree * tree, 1 + Rmax (INR (height x.1))
                                                                    (INR (height x.2)))) r).
        { apply: eq_dist_comp. apply: eq_dist_bind_rvar_pair. }
          rewrite {1}/pr_eq {1}[a in _ = a]/pr_eq.
          rewrite ?pr_eq_alt_comp /=. eapply eq_bigr => x ?.
          apply if_case_match. rewrite ?plus_INR ?Rmax_INR /rsize//=/a//=.
          rewrite Heql0 //=. specialize (S_INR (size ls0)) => /= => ->.
          destruct (Rlt_dec) as [Hlt0|Hnlt].
          { clear -Hlt0. specialize (pos_INR (size ls0)). rewrite //=. intros.  exfalso; nra. }
          destruct (Rle_dec); first nra.
          apply /eqP; apply /eqP. case: ifP; first by (move /eqP => ->; nra).
          move /eqP. intros. nra.
    }
    rewrite Heql0.
    rewrite /h.
    rewrite -?pr_mbind_mret.
    rewrite ?ldist_assoc.
    eapply eq_dist_ldist_bind_ext => ?.
    eapply eq_dist_ldist_bind_ext => pv.
    rewrite ldist_left_id.
    rewrite ?ldist_assoc.
    eapply mspec_eq_dist_ldist_bind_ext; first apply (rand_tree_non_empty_non_leaf); intros tl Htl.
    rewrite ?ldist_assoc.
    eapply mspec_eq_dist_ldist_bind_ext; first apply (rand_tree_non_empty_non_leaf); intros tr Htr.
    apply outcomes_eq_dist => //=.
    do 3 f_equal.
    destruct tl, tr; try (rewrite S_INR Rmax_INR; nra); [].
    exfalso.
    edestruct (filter_lt_gt_non_empty [:: a0, b0 & ls0] (sval pv)).
    ** rewrite //=.
    ** rewrite -Heql0. auto.
    ** eapply Htl => //=.
    ** eapply Htr => //=.
Qed.
Definition k := -/ ln(3/4).

Lemma partition_perm_eq' l:
  ∀ n, perm_eq ([seq i <- l | i < n] ++ [seq i <- l | i == n] ++ [seq i <- l | n < i])%nat l.
Proof.
  induction l as [| a l] => n; first by (rewrite //=).
  rewrite //=. case (ltngtP a n) => Hcmp //=.
  - by rewrite perm_cons.
  - eapply (perm_trans).
    { erewrite perm_cat2l. rewrite perm_catC cat_cons. erewrite perm_cons. done. }
    rewrite -cat_cons perm_catCA cat_cons perm_cons.
    rewrite -perm_catCA.
    eapply (perm_trans); last apply (IHl n).
    rewrite perm_cat2l perm_catC //.
  - rewrite -cat_cons perm_catCA // cat_cons perm_cons.
    rewrite perm_catCA //.
Qed.

Theorem bound x w:
    uniq x →
    rsize x > 1 →
    pr_gt (T x) ((k * ln (rsize x) + 1) + INR w)
       ≤ (rsize x) * (3/4)^w.
Proof.
  specialize (quicksort_rec.recurrence_span2.qs_span_bound _ _ _ _ T h' uniq).
  rewrite /size/rsize => Hrec. eapply Hrec; clear.
  - apply huniq.
  - intros l ?. rewrite //=.
    rewrite /quicksort_rec.recurrence_span2.size //=.
    cut (∀ n, INR (length (fst (rvar_of_ldist (h l) n))) +
              INR (length (snd (rvar_of_ldist (h l) n))) <=
                   INR (length l)); first by rewrite //=.
    intros n'. apply fun_to_mspec.
    rewrite /h.
    destruct l as [| a l].
    { apply mspec_mret => //=. replace (INR 0) with 0 by auto. nra. }
    destruct l as [| b l0].
    { apply mspec_mret => //=. replace (INR 0) with 0 by auto. replace (INR 1) with 1 by auto. nra. }
    remember (b::l0) as l eqn:Heql0.
    clear b l0 Heql0.
    tbind (λ x, sval x \in a :: l).
    { intros (?&?) => //. }
    intros (pv&Hin) _.
    remember (a :: l) as l0 eqn:Heql0.
    apply mspec_mret.
    rewrite //=.
    replace (length l0) with (size l0) by auto.
    assert (Hlt: (ssrnat.ltn O (size [seq i <- l0 | i == pv]))%nat).
    {  rewrite size_filter //=. rewrite -has_count has_pred1 Heql0 in_cons //=. }
    rewrite -plus_INR. apply le_INR. move /ltP in Hlt.
    rewrite -?size_legacy.
    specialize (partition_perm_eq' l0 pv). move /perm_size. rewrite ?size_cat.
    intros Hsize. nify. rewrite //= in Hsize. rewrite //=. omega.
  - rewrite /quicksort_rec.recurrence_span2.size //=.
    intros l ? Hgt. rewrite //=.
    cut (∀ n, INR (length (fst ((rvar_of_ldist (h l) n)))) +
              INR (length (snd ((rvar_of_ldist (h l) n)))) <=
                   INR (length l) - 1); first by rewrite //=.
    intros n'. apply fun_to_mspec.
    rewrite /h.
    destruct l as [| a l].
    { move : Hgt. rewrite //=. replace (INR 0) with 0 by auto. nra. }
    destruct l as [| b l0].
    { move : Hgt. rewrite //=. replace (INR 1) with 1 by auto. nra. }
    remember (b::l0) as l eqn:Heql0.
    clear b l0 Heql0.
    tbind (λ x, sval x \in a :: l).
    { intros (?&?) => //. }
    intros (pv&Hin) _.
    remember (a :: l) as l0 eqn:Heql0.
    apply mspec_mret.
    rewrite //=.
    replace (length l0) with (size l0) by auto.
    assert (Hlt: (ssrnat.ltn O (size [seq i <- l0 | i == pv]))%nat).
    {  rewrite size_filter //=. rewrite -has_count has_pred1 Heql0 in_cons //=. }
      move /ltP in Hlt.
    rewrite -?size_legacy.
    specialize (partition_perm_eq' l0 pv). move /perm_size => <-.
    rewrite ?size_cat ?plus_INR.
    specialize (pos_INR (size [seq i <- l0 | (i < pv)%nat])).
    specialize (pos_INR (size [seq i <- l0 | (pv < i)%nat])).
    rewrite //=. intros.
    assert (Hge1: 1 <= INR (size [seq i <- l0 | i == pv])).
    { replace 1 with (INR 1) by auto. apply le_INR. omega. }
    rewrite //= in Hge1. nra.
  - rewrite /quicksort_rec.recurrence_span2.size/quicksort_rec.recurrence_span2.rec_solution.d //=.
    intros x n Hle.
    replace 1 with (INR 1) in Hle by auto. apply INR_le in Hle.
    cut (INR (height (rvar_of_ldist (add_list_random x Leaf) n)) = 0); first by rewrite //=.
    apply fun_to_mspec.
    destruct x; [| destruct x].
    * rewrite alr_unfold.
      apply mspec_mret => //=.
    * rewrite alr_unfold.
      apply mspec_mret => //=.
    * rewrite //= in Hle. omega.
  - intros. eapply (rec_pr_rec_gt _ rsize _ _ _ _ T h' uniq (λ xy, Rmax (fst xy) (snd xy))); auto.
    intros. apply Trec; auto.
  - rewrite /quicksort_rec.recurrence_span2.size/quicksort_rec.recurrence_span2.rec_solution.d
            /quicksort_rec.recurrence_span2.rec_solution.m
            /quicksort_rec.recurrence_span2.qs_factor.alpha.
    etransitivity; first eapply Ex_max_partition.
    rewrite /m/rsize/size/length. do 2 destruct Rlt_dec; try nra.
Qed.

From Interval Require Import Tactic.
Remark concrete:
  ∀ l, uniq l → rsize l = 10 ^ 10 →
       pr_gt (T l) 300 ≤ 1/(10^15).
Proof.
  intros l Huniq Hsize.
  transitivity (pr_gt (T l) ((k * ln (rsize l) + 1) + 208)).
  - apply Rge_le, pr_gt_contra.
    rewrite Hsize. rewrite /k. interval.
  - replace 208 with (INR 208); last first.
    { vm_compute. nra. }
    etransitivity; first apply bound; auto; try nra; try rewrite Hsize; interval.
Qed.