(*goal
1.搞清楚abs的意思以及用途
2.修正abs——half定义
3.is_AVLpartial是否需要改为prop形式
*)
Require Import Coq.Lists.List.
Require Import Coq.ZArith.ZArith.
Import ListNotations.
Local Open Scope Z.

Definition Key: Type := Z.
Definition Value : Type := Z.
Definition Depth : Type := nat.

Record Node := newnode{
  key_of : Key;
  value_of : Value;
  depth_of : Depth
}.

Inductive tree : Type :=
| E: tree
| T : tree -> Node -> tree -> tree.

Inductive LeftOrRight :=
| L: LeftOrRight
| R: LeftOrRight.

Definition half_tree: Type := (LeftOrRight * Node * tree)%type.

Definition partial_tree: Type := list half_tree.

Fixpoint complete_tree  (l: partial_tree ) (b : tree) : tree :=
match l with
|nil => b
|(L, t, brother) :: l' => complete_tree l' (T brother t b)
|(R, t, brother) :: l' => complete_tree l' (T b t brother)
end.
(*自底向上构建二叉搜索树，L为brother在b树的左侧*)
(* -------------       operations -------------------------------------
insert
delete
look_up *)

Fixpoint general_split_aux (lbool rbool : Node -> bool) (s :tree)  (b: partial_tree):
 partial_tree * tree :=
 match s with
 |E => (b, E)
 |T l t' r =>
    if (lbool t') then
        general_split_aux lbool rbool l ((R, t',r)::b)
      else if (rbool t') then
        general_split_aux lbool rbool r ((L, t',l)::b)
      else
       (b, s)
end.
(*generalsplit的过程是将一个完整的树拆分成一个partialtree的过程，节点满足左边的条件就继续拆分左子树，右子树挂着*)

(* lbool :  22 <? node.key
rbool :  node.key         <? 22 *)
(*                                 此处为insert定义                                             *)
Definition general_split_simpl lbool rbool s := general_split_aux lbool rbool s nil.
Definition insert_split (x: Key)(s: tree) (b : partial_tree)  : list half_tree * tree := general_split_aux (fun k => x <? (key_of k)) (fun  k =>  (key_of k) <? x) s b.
(*找到插入的节点x应该处于的位置，x比当前节点小就往左边走，否则往右边走*)
Definition insert_root (x : Key) (v: Value)(s:tree) : tree :=
  match s with
  |E => T E (newnode x v (S O)) E
  |T l (newnode x' v' d) r=> T l (newnode x v d) r
  end.
Definition depth_of_tree s :=
match s with
| E => O
| T _ n _ => depth_of n
end.
Definition root_of_tree s:=
match s with
|E=> None
|T _ n _=>Some n
end.
Definition LT_AVL a b :=
  match a with
    | None => true
    | Some n => key_of n <? key_of b
  end.
Definition GT_AVL a b :=
  match a with
    | None => true
    | Some n => key_of n >? key_of b
  end.
Definition reversedepth (s:nat*nat):=
match s with
|(a,b)=>(b,a)
end.
(*需要检查这里的定义--------------------------------------------------------------------------------------------------------------*)
Definition left_rotate a oldroot y :=
 match oldroot with
  | None => E
  | Some oldroot => match y with
                    |E => E
                    |T b newroot c => T (T a (newnode (key_of oldroot)(value_of oldroot)(S(depth_of_tree a))) b) (newnode (key_of  newroot)(value_of  newroot) (S (S (depth_of_tree a)))) c
                    end
 end.
Definition right_rotate y oldroot a :=
  match oldroot with
  | None => E
  | Some oldroot =>
    match y with
    |E => E
    |T b newroot c => T  b (newnode (key_of newroot)(value_of newroot) (S(S(depth_of_tree a)))) (T c (newnode (key_of  oldroot)(value_of  oldroot)(S(depth_of_tree a))) a)
    end
   end.
(*新定点高度永远为最外侧没修改的子树的高度+2*)
Definition is_balance_with_n (t1 :tree)(n:nat):=
match Nat.leb (depth_of_tree t1) (S n) with
|true=>match Nat.leb n (S(depth_of_tree t1)) with
        |true=>true
        |false=>false
        end
|false=>false
end.
Inductive isavltree : tree->Prop:=
|avlE : isavltree E
|avlT : forall Ltree n Rtree, isavltree Ltree->isavltree Rtree
            -> Nat.le (depth_of_tree Ltree) (S(depth_of_tree Rtree))->
            Nat.le (depth_of_tree Rtree) (S(depth_of_tree Ltree))->
            LT_AVL (root_of_tree Ltree) n = true ->
            GT_AVL (root_of_tree Rtree) n = true ->
                isavltree (T Ltree n Rtree).
Fixpoint isavltree_bool (t:tree):bool:=
match t with
|E=>true
|T Lefttree n Righttree=>if Nat.leb (depth_of_tree Lefttree) (S (depth_of_tree Righttree))
                         then if Nat.leb (depth_of_tree Righttree) (S (depth_of_tree Lefttree))
                              then match isavltree_bool Lefttree with
                                   |true=>match isavltree_bool Righttree with
                                          |true=>true
                                          |false=>false
                                          end
                                   |false=>false
                                   end
                              else false
                         else false
end.

(* is_avl_partial中的n是split得到的subtree的高度 *)
Inductive is_avl_partial : partial_tree -> tree -> Prop :=
  | is_avl_partial_nil : forall t , isavltree t -> is_avl_partial nil t
  | is_avl_partial_consL : forall father brother l' t,
     isavltree (T brother father t) ->
     is_avl_partial l' (T brother father t) ->
     is_avl_partial ((L,father,brother) :: l') t
  | is_avl_partial_consR : forall father brother l' t,
     isavltree (T t father brother) ->
     is_avl_partial l' (T t father brother) ->
     is_avl_partial ((R,father,brother) :: l') t.

Definition true_insert_balance (LRf:LeftOrRight)(f:Node)(b st:tree):=
match LRf with
|L=>match st with
    |E=>E
    |T RL Rnode RR =>if Nat.leb(depth_of_tree RL) (depth_of_tree RR)
                     then left_rotate b (Some f) st
                     else left_rotate b (Some f) (right_rotate RL (root_of_tree st) RR)
                     end
|R=>match st with
    |E=>E
    |T LR Lnode LL =>if Nat.leb(depth_of_tree LR) (depth_of_tree LL)
                     then right_rotate st (Some f) b
                     else right_rotate (left_rotate LL (root_of_tree st) LR) (Some f) b
                     end
end.
(*改变高度只在普通的balance操作以及rotate操作进行的时候进行*)
(*修改balance为一层一层的*)
Fixpoint insert_balance  (l: partial_tree) (s: tree) : tree  :=
match l with
(*插入节点为根节点*)
|nil=>s
|(LRfather,father,brother)::l'=>
          if(Nat.leb (depth_of_tree s) (depth_of father))
          then if (Nat.ltb (depth_of_tree s) (depth_of father))
               then
              match LRfather with
                   |L=>insert_balance l' (T brother father s)
                   |R=>insert_balance l' (T s father brother)
                   end
               else
              match LRfather with
                   |L=>insert_balance l' (T brother (newnode (key_of father)(value_of father)(S(depth_of father))) s)
                   |R=>insert_balance l' (T s (newnode (key_of father)(value_of father) (S(depth_of father))) brother)
                   end
          else insert_balance l'  ( true_insert_balance LRfather father brother s )
end.
Definition insert' x v s :=let (h,b) := insert_split x s nil in (h, insert_root x v b).
(*没有进行balance操作*)

Definition insert (k : Key)(v : Value)(t : tree) : tree :=
  let (l,s) := insert' k v t in insert_balance l s.

Definition not_empty (t:tree): bool:=
match t with
|E=>false
|T _ _ _=>true
end.

Definition getltree (t:tree): tree:=
match t with
|E=>E
|T l n r=>l
end.

Definition getrtree (t:tree): tree:=
match t with
|E=>E
|T l n r=>r
end.
(* Lemma split_prop : forall k s l t,t=snd(insert_split k s l)->
(not_empty t=true)->isavltree (T (getltree t) (newnode k (value_of (root_of_tree t)) (depth_of (root_of_tree t))) (getrtree t)).
 *)
Definition fit_tree x sub : Prop :=
  match sub with
    | E => True
    | T Ltree n Rtree => match (root_of_tree Ltree) , (root_of_tree Rtree) with
       | None , None => True
       | None , Some r => key_of r > x
       | Some l , None => key_of l < x
       | Some l , Some r => key_of l < x /\ key_of r > x
       end
end.

Lemma general_split_aux_prop : forall f1 f2 t l1 l2,
  (fst (general_split_aux f1 f2 t (l1 ++ l2)), snd (general_split_aux f1 f2 t (l1 ++ l2))) =
  (fst (general_split_aux f1 f2 t l1) ++ l2, (snd (general_split_aux f1 f2 t l1))).
Proof.
intros. generalize dependent l1. generalize dependent l2.
induction t.
- intros. auto.
- intros. destruct (f1 n) eqn:EL.
+ unfold general_split_aux. rewrite EL. specialize (IHt1 l2 ((R, n, t2) :: l1) ). fold general_split_aux. apply IHt1.
+ unfold general_split_aux. destruct (f2 n) eqn:ER.
* rewrite EL. specialize (IHt2 l2 ((L, n, t1) :: l1) ). fold general_split_aux. apply IHt2.
* rewrite EL. auto.
Qed.

Lemma fit_tree_split : forall k t , isavltree t -> fit_tree k (snd (insert_split k t [])).
Proof.
intros.
induction t.
- simpl. exact I.
- inversion H. apply IHt1 in H3. apply IHt2 in H4. unfold insert_split. destruct (k<?key_of n) eqn:EL.
+ unfold general_split_aux. rewrite EL . pose proof(general_split_aux_prop (fun k0 : Node => k <? key_of k0)(fun k0 : Node => key_of k0 <? k) t1 [] [(R,n,t2)]). inversion H9. fold general_split_aux. (* rewrite H12. rewrite之后直接用insertsplit 是
的结论就好了 *) admit.
+ destruct (key_of n <? k) eqn:ER.
* admit.
* inversion H. unfold general_split_aux. rewrite EL. rewrite ER. simpl snd. assert(Heq:(key_of n = k)). admit. simpl. destruct (root_of_tree t1) eqn:E1.
** destruct (root_of_tree t2) eqn:E2.
*** inversion H16. inversion H17. rewrite Heq in H19,H20. admit.
*** inversion H16. rewrite Heq in H19. admit.
** destruct (root_of_tree t2) eqn:E2.
*** inversion H17. rewrite Heq in H19. admit.
*** exact I.
Admitted.

Lemma node_tree_height : forall l r n,isavltree (T l n r)->(Nat.lt(depth_of_tree l)  (depth_of n )) /\(Nat.lt(depth_of_tree r)(depth_of n )).
Proof.
intros.
Admitted.

Lemma insert_balance_prop : forall par sub , is_avl_partial par sub -> isavltree (insert_balance par sub).
Proof.
intros. generalize dependent sub.
induction par.
- intros. simpl. inversion H. apply H0.
- intros. destruct a. destruct p. inversion H.
+ inversion H5. simpl. apply leb_correct in H12. apply leb_correct in H13. pose proof(node_tree_height t sub n). pose proof H5. apply H16 in H5. destruct H5. pose proof H18. rewrite<- Nat.ltb_lt in H18. pose proof H18. rewrite Nat.ltb_lt in H20. unfold lt in H20. apply le_S in H20. assert(Hle :(Nat.le (S (depth_of_tree sub)) (S (depth_of n)))->(Nat.le(depth_of_tree sub)(depth_of n))). admit. apply Hle in H20.
rewrite<- Nat.leb_le in H20. rewrite H20. rewrite<-Nat.ltb_lt in H19. rewrite H19. specialize (IHpar (T t n sub)). apply IHpar in H6. apply H6.
+ inversion H5. simpl. apply leb_correct in H12. apply leb_correct in H13. pose proof(node_tree_height sub t n). pose proof H5. apply H16 in H5. destruct H5. pose proof H5. rewrite<- Nat.ltb_lt in H5. pose proof H5. rewrite Nat.ltb_lt in H20. unfold lt in H20. apply le_S in H20. assert(Hle :(Nat.le (S (depth_of_tree sub)) (S (depth_of n)))->(Nat.le(depth_of_tree sub)(depth_of n))). admit. apply Hle in H20.
rewrite<- Nat.leb_le in H20. rewrite H20. rewrite<-Nat.ltb_lt in H19. rewrite H19. specialize (IHpar (T sub n t)). apply IHpar in H6. apply H6.
Admitted.

Lemma fit_tree_prop_need1 : forall x v d x' v' Ltree Rtree,
isavltree
  (T Ltree {| key_of := x'; value_of := v'; depth_of := d |} Rtree) ->
fit_tree x (T Ltree {| key_of := x'; value_of := v'; depth_of := d |} Rtree) ->
isavltree
  (T Ltree {| key_of := x; value_of := v; depth_of := d |} Rtree).
Proof.
  intros. inversion H.
  apply avlT ; auto ; clear- H0; simpl in H0 ; destruct (root_of_tree Ltree) , (root_of_tree Rtree) ; simpl in * ; auto;
    try (apply Z.ltb_lt ; apply H0) ; try (apply Z.gtb_lt ; apply Z.gt_lt ; apply H0).
Qed.

Lemma insert_prop_need : forall par sub x v , is_avl_partial par sub -> fit_tree x sub ->
  isavltree (insert_balance par (insert_root x v sub)).
Proof.
  intros.
  destruct sub ; intros.
  - simpl.
  pose proof (insert_balance_prop par E). pose proof H.
   apply H1 in H. induction par.
   + simpl. apply avlT; auto.
   + destruct a. destruct p.
    admit.
  - simpl insert_root. destruct n.
    apply insert_balance_prop.
    inversion H ; subst.
    + apply is_avl_partial_nil. apply (fit_tree_prop_need1 _ _ _ _ _ _ _ H1 H0); auto.
    + apply is_avl_partial_consL ; auto.
      inversion H ; subst.
      * inversion H1 ; subst ; apply avlT ; auto.
        -- apply (fit_tree_prop_need1 _ _ _ _ _ _ _ H9 H0).
        -- clear - H0.
           simpl in H0.
           admit.
      *
Admitted.

Lemma complete_tree_prop : forall x s,
  s = complete_tree (fst (insert_split x s nil))  (snd (insert_split x s nil)).
Proof.
intros.
induction s.
- auto.
- unfold insert_split. destruct (x<?key_of n) eqn:E1.
+ unfold general_split_aux. rewrite E1. assert(HR:[(R,n,s2)]=[]++[(R,n,s2)] ). auto. unfold half_tree. rewrite-> HR. fold general_split_aux. (*  rewrite->(general_split_aux_prop (fun k : Node => x <? key_of k) (fun k : Node => key_of k <? x) s1 [] [(R,n,s2)]). *) Admitted.
(* rewrite之后snd部分直接变成IHs1中的snd，第一部分也是这样，但多了一个halftree衔接在后面，直接对completetree进行展开，先替换s1在替换complete回去 *)
Lemma is_avl_partial_add_back : forall par sub par',
  is_avl_partial par sub ->
  is_avl_partial par' (complete_tree par sub) ->
  is_avl_partial (par ++ par') sub.
Proof.
intros.
generalize dependent sub.
induction par.
- simpl in *. auto.
- intros. inversion H.
+ destruct a. destruct p. simpl in H0. specialize (IHpar (T brother father sub)). pose proof( IHpar  H5).  inversion H1.  rewrite<- H8 in H0. rewrite <-H9 in H0. rewrite <-H10 in H0. apply H6 in H0. assert ((((L, n, t0) :: par) ++ par')=((L, n, t0) :: (par ++ par'))). reflexivity. unfold half_tree. rewrite H7. rewrite <-H9. rewrite<- H10. (* inversion H0.
* apply is_avl_partial_consL. apply H12. rewrite<-H11 in H0. apply H0.
* apply is_avl_partial_consL. inversion H12.   *)
Admitted.

(* 可以得到性质：insert_balance par (insert_balance Half sub)就直接证完了 *)
Lemma insert_prop_need2 : forall x s , isavltree s -> is_avl_partial  (fst (insert_split x s nil))  (snd (insert_split x s nil)).
Proof.
intros.
induction H.
- simpl. apply is_avl_partial_nil. apply avlE.
- unfold insert_split in *.
  destruct (x<?key_of n) eqn:C1 ; destruct (key_of n <? x) eqn : C2;
  simpl ; try rewrite C1 ; try rewrite C2 ; simpl.
  + pose proof (general_split_aux_prop (fun k : Node => x <? key_of k) (fun k : Node => key_of k <? x) Ltree [] [(R,n,Rtree)]).
    unfold partial_tree in H5.
    inversion H5.
    unfold half_tree at 2 4.
    rewrite -> H7,  H8.
    apply is_avl_partial_add_back ; auto. clear H5 H7 H8.

    pose proof (complete_tree_prop x Ltree).
    unfold insert_split in H5.
    rewrite <- H5. clear H5.
    apply is_avl_partial_consR ; auto.
    * apply avlT ; auto.
    * apply is_avl_partial_nil. apply avlT ; auto.
  + pose proof (general_split_aux_prop (fun k : Node => x <? key_of k) (fun k : Node => key_of k <? x) Ltree [] [(R,n,Rtree)]).
    unfold partial_tree in H5.
    inversion H5.
    unfold half_tree at 2 4.
    rewrite -> H7,  H8.
    apply is_avl_partial_add_back ; auto. clear H5 H7 H8.

    pose proof (complete_tree_prop x Ltree).
    unfold insert_split in H5.
    rewrite <- H5. clear H5.
    apply is_avl_partial_consR ; auto.
    * apply avlT ; auto.
    * apply is_avl_partial_nil. apply avlT ; auto.
  + pose proof (general_split_aux_prop (fun k : Node => x <? key_of k) (fun k : Node => key_of k <? x) Rtree [] [(L,n,Ltree)]).
    unfold partial_tree in H5.
    inversion H5.
    unfold half_tree at 2 4.
    rewrite -> H7,  H8.
    apply is_avl_partial_add_back ; auto. clear H5 H7 H8.
    pose proof (complete_tree_prop x Rtree).
    unfold insert_split in H5.
    rewrite <- H5. clear H5.
    apply is_avl_partial_consL ; auto.
    * apply avlT ; auto.
    * apply is_avl_partial_nil. apply avlT ; auto.
  + apply is_avl_partial_nil. apply avlT ; auto.
Qed.

Theorem insert_avl_correct :forall t k v,isavltree t->isavltree (insert k v t).
Proof.
intros.
pose proof (insert_prop_need2 k t H).
revert dependent k.
induction H ; unfold insert in *; simpl in * ; intros.
- apply avlT ;  auto ; apply avlE.
- unfold insert' in *.
  rewrite (surjective_pairing (insert_split k (T Ltree n Rtree) [])).
  simpl.
  apply insert_prop_need ; auto.
  apply fit_tree_split ; auto.
  apply avlT ; auto.
Qed.

(* Inductive insert (k: Key)(v:Value)(t:tree) : tree->Prop :=
|insert_intro : forall (l : list half_tree) (s : tree) (b : tree),
(l,s)=insert' k v t->b=insert_balance l s->insert k v t b.
 *)
Lemma derived_from_is_avl_partial : forall (a:half_tree)(l:partial_tree)(n:nat),
is_avl_partial (a :: l) (S n) = true->depth_of_tree( snd a )=n\/depth_of_tree( snd a )=S n\/depth_of_tree( snd a )=S(S n).
Proof.
intros.
simpl in H. destruct a. destruct p. simpl. destruct (is_balance_with_n t (S n)) eqn:H'.
- unfold is_balance_with_n in H'. destruct (depth_of_tree t <=? S (S n))%nat eqn:H''.
+ destruct  (S n <=? S (depth_of_tree t))%nat eqn:H'''.
* simpl in H'''. Print Nat.leb_le. apply Nat.leb_le in H''. apply Nat.leb_le in H'''.
inversion H''.
** right. right. reflexivity.
** inversion H1.
*** right. left. reflexivity.
*** inversion H3.
**** left. reflexivity.
**** apply Nat.ltb_ge in H4. assert(~(m1<depth_of_tree t)%nat). apply Nat.ltb_nlt in H4.
apply H4. unfold lt in H6. rewrite H5 in H6. unfold not in H6. apply H6 in H'''. inversion H'''.
* simpl in H'''. inversion H'.
+ inversion H'.
- inversion H.
Qed.
Lemma derived_from_is_avl_partial' : forall (a:half_tree) (l:partial_tree)(n:nat),
is_avl_partial (a::l) n=true->Nat.le (depth_of_tree (snd a)) (S n).
Proof.
intros.
destruct n eqn:E.
- destruct a. destruct p. simpl. simpl in H. destruct (is_balance_with_n t 0) eqn:E'.
+ unfold is_balance_with_n  in E'. destruct (depth_of_tree t <=? 1)%nat eqn:E''.
* rewrite Nat.leb_le in E''. apply E''.
* inversion E. inversion E'.
+ inversion H.
- pose proof(derived_from_is_avl_partial a l n0). apply H0 in H. destruct H.
+ rewrite H. apply le_S. apply le_S. apply le_n.
+ destruct H. rewrite H. apply le_S. apply le_n. rewrite H. apply le_n.
Qed.
Lemma length_after_insert : forall (s:tree)(x:Key)(v:Value),
depth_of_tree (insert_root x v s)=depth_of_tree s\/depth_of_tree (insert_root x v s)=S(depth_of_tree s).
Proof.
intros.
destruct s.
- simpl. left. reflexivity.
- destruct n. simpl. left. reflexivity.
Qed.
Lemma is_partial_next : forall (l:partial_tree) (a:half_tree)(n:nat),
is_avl_partial (a::l) n=true->is_avl_partial l (S n)=true.
Proof.
intros.
Admitted.
(*以下证明(* (* (* (* balance操作可以保证在子树中插入节点不会影响整棵树的avl性质 *) *) *) *)*)
Theorem avlpartial_correctness :    forall (l:partial_tree) (s:tree)(s':tree)(x:Key) (v:Value)(n:nat),
isavltree s->isavltree s'->depth_of_tree s'=n\/depth_of_tree s'=(S n)->depth_of_tree s=n
-> is_avl_partial l n=true->isavltree (insert_balance l s').
Proof.
intros. generalize dependent n. generalize dependent s. generalize dependent s'.
induction  l.
- intros. simpl. apply H0.
- intros. destruct a. destruct p. simpl. destruct (depth_of_tree s' <=? S (depth_of_tree t))%nat eqn:E.
+ destruct l0 eqn:E'. (*不需要平衡下的左右之分*)
* (*左侧*) specialize (IHl (T t n0 s')). assert(H4:isavltree(T t n0 s')).
 pose proof H3 as E4 . simpl in H3. destruct (is_balance_with_n t n) eqn:E1.
** destruct (isavltree_bool t) eqn:E2.
*** pose proof E2. apply isavltree_bool_prop_equal in E2. apply avlT. apply E2. apply H0.
pose proof (derived_from_is_avl_partial' (L,n0,t),l,n). apply X in E4. simpl in E4. destruct H1.
**** rewrite<- H1 in E4. apply E4.
**** rewrite<- H1 in E4. apply le_S. apply E4.
**** apply Nat.leb_le in E. apply E.
*** inversion H3.
** inversion H3.
** specialize( IHl H4). specialize (IHl (T t n0 s)). assert(H5:isavltree(T t n0 s)).
pose proof H3 as E4 . simpl in H3. destruct (is_balance_with_n t n) eqn:E1.
*** destruct (isavltree_bool t) eqn:E2.
**** pose proof E2. apply isavltree_bool_prop_equal in E2. apply avlT. apply E2. apply H.
pose proof (derived_from_is_avl_partial' (L,n0,t),l,n). apply X in E4. simpl in E4. destruct H1.
***** rewrite<- H2 in E4. apply E4.
***** rewrite<- H2 in E4. apply E4.
***** destruct H1. rewrite H2. rewrite <-H1. apply Nat.leb_le in E. apply E.
assert(ss': Nat.le (depth_of_tree s)(depth_of_tree s')). rewrite H1. rewrite H2. apply le_S. apply le_n.
apply Nat.leb_le in E. pose proof( Nat.le_trans (depth_of_tree s)(depth_of_tree s')(S (depth_of_tree t))).
apply H6. apply ss'. apply E.
**** inversion H3.
*** inversion H3.
*** specialize(IHl H5).
specialize(IHl (S n)).

* (*右侧*)

* admit.

+ (*需要平衡之下的平衡后的情况*)

(*                            此处为delete定义                               *)
Definition delete_split (x:Key) (s:tree ) (b:partial_tree):list half_tree*tree:= insert_split x s b.
Definition minimum_split (s: tree) (b: list half_tree):  list half_tree
:= fst (general_split_aux (fun k - => true) (fun k =>  false) s b).
(*fst作用为取出二元组中的第一个元素*)
(*minimum_split将位置改至最小值左下方的空树*)
Definition Up_split (b: list half_tree) : list half_tree * tree :=
match b with
| nil => (nil , E)
|(L, n, r)::l => (l , T r n E)
|(R ,n, r)::l => (l , T E n r)
end.
(*upsplit的作用是取出右子树最小元素*)
Definition delete_root (s: tree) (h : list half_tree): list half_tree * tree :=
match s with
|E => (h, E)
|T E n E => (h, E)
|T E n (T rightsonL rightson rightsonR) => (h, T rightsonL rightson rightsonR)
|T (T leftsonL leftson leftsonR) n E => (h, T leftsonL leftson leftsonR)
|T l n r => match Up_split (minimum_split r nil) with
                  |(restr, T E successor x) =>
                             ( ( restr ++ ( (L,  (newnode (key_of successor) (value_of successor) (depth_of n))    , l)::h )),x)
                  |(_ , _) => (h, E)
                  end
end.
Fixpoint delete_balance  (l:partial_tree) (s:tree) : partial_tree*tree:=
match l with
(*平衡的节点为根节点*)
|nil=>
    match s with
    |E=>(nil,E)
    |T Ltree n Rtree=>if Nat.ltb (S(depth_of_tree Ltree)) (depth_of_tree Rtree)
              then match Rtree with
              |E=>(l,s)
              |(T RL nR RR)=> if leb (depth_of_tree RL) (depth_of_tree  RR)
                              then (nil,left_rotate Ltree n Rtree)
                              else (nil,left_rotate  Ltree n (right_rotate RL nR RR) )
              end
              else if Nat.ltb (S(depth_of_tree Rtree))  (depth_of_tree Ltree)
              then match Ltree with
              |E=>(l,s)
              |(T LL nL LR)=>if leb (depth_of_tree LR) (depth_of_tree LL)
                              then (nil,right_rotate Ltree n Rtree)
                              else (nil,right_rotate  (left_rotate LL nL LR)  n Rtree)
              end
              else (l,s)
    end
(*平衡的节点不是根节点*)
|(LRfather,father,brother)::l'=>
           if leb (depth_of_tree brother) (S(depth_of_tree s))
                      then match LRfather with
                      |L=>delete_balance l' (T brother father s)
                      |R=>delete_balance l' (T s father brother)
                      end
                      else match LRfather with
                        |L=>
                              match brother with
                              |E=>(l,s)
                              |(T Lofbrother nb Rofbrother) =>
                                      if Nat.ltb (depth_of_tree s)(depth_of_tree Lofbrother)
                                        then delete_balance l' (right_rotate brother father s)
                                        else delete_balance l'(  right_rotate    (left_rotate Lofbrother nb Rofbrother)  (father)(s))
                              end
                        |R=>
                              match brother with
                              |E=>(l,s)
                              |(T Lofbrother nb Rofbrother) =>
                                      if Nat.ltb (depth_of_tree s) (depth_of_tree Rofbrother)
                                        then delete_balance l' (left_rotate s father brother)
                                        else   delete_balance l' ( left_rotate s father  (right_rotate Lofbrother nb Rofbrother) )
                               end
                            end

end.
Definition true_delete (x:Key)(s:tree) := let (h,s):= delete_split x s nil in  (delete_root s h).
(*                              此处为lookup定义                                   *)
Fixpoint lookup (x:Key)(t:tree): option Value:=
match t with
|E=>None
|T lefttree n righttree=>if x<?(key_of n) then
                                          match (lookup x lefttree) with
                                          |None=>None
                                          |Some v'=>(Some  v')
                                          end
                                          else if (key_of n)<?x then
                                          match (lookup x righttree) with
                                          |None=>None
                                          |Some v'=>(Some  v')
                                          end
                                          else (Some (value_of n))
end.

(* Variable t: Node.
Check (key_of t).
Check (value_of t). *)

(*------------  证明AVL性质 --------------------------------------------
 SearchTree
 depth
 abs*)
Definition relate_map := Key -> option Value .
Definition v_update (k:Key)(v: Value )(m: relate_map):relate_map :=
 fun x =>  if  k =?x then Some v else m x.
(*相当于向原来的map中加入新的匹配，如果是新加入的点就单独考虑他对应的value值，否则按之前的map考虑*)
Definition relate_default : relate_map  := fun  x => None.
Definition combine (m1 m2: relate_map ) : relate_map :=
  fun x => match m1 x , m2 x with
          |None, Some v => Some v
          |Some v, None => Some v
          |_ ,_ => None
          end.
Inductive Abs : tree -> relate_map -> Prop :=
|Abs_E : Abs E (relate_default)
|Abs_T: forall a b left right n ,
  Abs left a ->
  Abs right b ->
  Abs (T  left n right)  (v_update (key_of n) (value_of n) (combine  a b) ).
Inductive Abs_half : list half_tree -> relate_map -> Prop :=
|Abs_nil : Abs_half nil relate_default
|Abs_cons : forall l lmap LOR n ktree tmap,
   Abs_half l lmap -> Abs ktree tmap -> Abs_half ((LOR,n,ktree)::l) (combine lmap (v_update (key_of n) (value_of n) tmap)).
(*有问题*)

(*searchtree *)
Definition optionZ_lt (ok1 ok2: option Key): Prop :=
  match ok1, ok2 with
  | Some k1, Some k2 => k1 < k2
  | _, _ => True
  end.
Inductive SearchTree : option Key -> tree -> option Key -> Prop :=
| ST_E : forall lo hi, optionZ_lt lo hi -> SearchTree lo E hi
| ST_T: forall lo l n r hi,
    SearchTree lo l (Some (key_of n)) ->
    SearchTree (Some (key_of n)) r hi ->
    SearchTree lo (T l n r) hi.
Theorem insert_avl : forall k v t finaltree, isavltree t ->insert k v t finaltree-> isavltree finaltree .
Proof.
intros.
inversion H0. subst.
Admitted.

Inductive SearchTree_half_in:
  option Key -> partial_tree -> option Key -> Prop :=
| ST_in_nil:
    forall lo hi, optionZ_lt lo hi -> SearchTree_half_in lo nil hi
| ST_in_cons_L:
    forall lo hi h l n,
      SearchTree_half_in lo h hi ->
      SearchTree lo l (Some (key_of n)) ->
      SearchTree_half_in (Some (key_of n)) ((L, n, l) :: h) hi
| ST_in_cons_R:
    forall lo hi h r n,
      SearchTree_half_in lo h hi ->
      SearchTree (Some (key_of n)) r hi ->
      SearchTree_half_in lo ((R, n, r) :: h) (Some (key_of n)).
Inductive SearchTree_half_out:
  option Key -> partial_tree -> option Key ->  Prop :=
| ST_out_nil:
    forall lo hi, optionZ_lt lo hi -> SearchTree_half_out lo nil hi
| ST_out_cons_L:
    forall lo hi h l n,
      SearchTree_half_out lo h hi ->
      SearchTree lo l (Some (key_of n)) ->
      optionZ_lt (Some (key_of n)) hi ->
      SearchTree_half_out lo ((L, n, l) :: h) hi
| ST_out_cons_R:
    forall lo hi h r n,
      SearchTree_half_out lo h hi ->
      SearchTree (Some (key_of n)) r hi ->
      optionZ_lt lo (Some (key_of n)) ->
      SearchTree_half_out lo ((R, n, r) :: h) hi.

(*depth *)

(*abs*)

Definition relate_single (k: Key) (v: Value): relate_map :=
  fun x =>
    if Z.eq_dec x k then Some v else None.

Inductive Abs' : tree -> relate_map -> Prop :=
| Abs_E' : Abs' E relate_default
| Abs_T': forall l n r lm rm,
    Abs' l lm ->
    Abs' r rm ->
    Abs'
      (T l n r)
      (combine lm
         (combine (relate_single (key_of n) (value_of n)) rm)).
Inductive Abs_half : partial_tree -> relate_map -> Prop :=
| Abs_half_nil : Abs_half nil relate_default
| Abs_half_cons: forall LR n t h m1 m2,
    Abs t m1 ->
    Abs_half h m2 ->
    Abs_half
      ((LR, n, t) :: h)
      (combine m1
         (combine (relate_single (key_of n) (value_of n)) m2)).


(**
 Verified SAR-BP: A verified C implementation of SAR backprojection
 with a certified absolute error bound.

 Version 1.0 (2015-12-04)

 Copyright (C) 2015 Reservoir Labs Inc.
 All rights reserved.

 This file is free software. You can redistribute it and/or modify it
 under the terms of the GNU General Public License as published by the
 Free Software Foundation, either version 3 of the License (GNU GPL
 v3), or (at your option) any later version.  A verbatim copy of the
 GNU GPL v3 is included in gpl-3.0.txt.

 This file is distributed in the hope that it will be useful, but
 WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See LICENSE for
 more details about the use and redistribution of this file and the
 whole Verified SAR-BP library.

 This work is sponsored in part by DARPA MTO as part of the Power
 Efficiency Revolution for Embedded Computing Technologies (PERFECT)
 program (issued by DARPA/CMO under Contract No: HR0011-12-C-0123). The
 views and conclusions contained in this work are those of the authors
 and should not be interpreted as representing the official policies,
 either expressly or implied, of the DARPA or the
 U.S. Government. Distribution Statement "A" (Approved for Public
 Release, Distribution Unlimited.)

 If you are using or modifying Verified SAR-BP in your work, please
 consider citing the following paper:

 Tahina Ramananandro, Paul Mountcastle, Benoit Meister and Richard
 Lethin.
 A Unified Coq Framework for Verifying C Programs with Floating-Point
 Computations.
 In CPP (5th ACM/SIGPLAN conference on Certified Programs and Proofs)
 2016.

 Verified SAR-BP derives from prior work listed in ACKS along with
 their copyright and licensing information.

 Verified SAR-BP requires third-party libraries listed in ACKS along
 with their copyright information.
*)
(**
Author: Tahina Ramananandro <ramananandro@reservoir.com>

Final proof, main theorem.
Submitted to CPP 2016
*)

Require Export
SARBackProjSourceSpec
SARBackProjSourceLink SARBackProjSourceOpt.

(**
This file summarizes the main result on our
implementation of SAR backprojection:

For every program linked with our implementation,
SAR backprojection introduces at most `pixel_error'
total rounding and approximation error for each pixel.

The hypotheses and correctness statement are specified in
SARSpecs.v

The real-number algorithm for SAR backprojection is
specified in SARBackProj.v

The code is defined in SARBackProjSource.v,
which was obtained from backprojection.c
using CompCert clightgen.

The numerical bounds on the input data are defined in SARBounds.v,
which was obtained from a pass on the PERFECT data suite (
http://hpc.pnl.gov/PERFECT/ )

The implementation error bound is *computed* in
SARBackProjSourceOpt.v, along with the proof. Since the same proof
scripts are used for all image sizes (small, medium and large), no
computed numerical values will textually appear in them.

Among the proofs that are not too hard to read:
- SARBackProjSourceNormOpt.v: proof of C implementation
  of square root and norm
- SARBackProjSourceSinOpt.v: proof of C implementation
  of sine
**)

Theorem sar_backprojection_correct:
  forall (ge: Clight.genv)
         `(LINK: LinkingHypotheses ge),
    sar_backprojection_spec
      ge
      (Clight.Internal
         (f_sar_backprojection2 N_RANGE_UPSAMPLED (Z.of_nat BP_NPIX_X) (Z.of_nat BP_NPIX_Y) (Z.of_nat N_PULSES)))
      pixel_error.
Proof.
  exact SARBackProjSourceOpt.sar_backprojection_correct.
Qed.

Eval unfold sar_backprojection_spec, pixel_error in
 $( let t := type of sar_backprojection_correct in exact t )$ .

Print Assumptions sar_backprojection_correct.


(* Celsius project *)
(* Clément Blaudeau - Lamp@EPFL & Inria 2020-2022 *)
(* ------------------------------------------------------------------------ *)
(* This file describes the typing rules *)

From Celsius Require Export Helpers.

(* ------------------------------------------------------------------------ *)
(** ** Mode lattice and subtyping *)

Inductive S_Mode: Mode -> Mode -> Prop :=
| s_mode_refl : forall μ, μ ⊑ μ
| s_mode_hot: forall μ, hot ⊑ μ
| s_mode_warm: forall Ω, warm ⊑ cool Ω
| s_mode_cool: forall Ω1 Ω2, Ω2 <= Ω1 -> cool Ω1 ⊑ cool Ω2
| s_mode_cold: forall μ, μ ⊑ cold
where "m1 ⊑ m2" := (S_Mode m1 m2).
Global Hint Constructors S_Mode: typ.

(* We then have basic results *)

Lemma s_mode_trans: forall μ1 μ2 μ3, μ1 ⊑ μ2 -> μ2 ⊑ μ3 -> μ1 ⊑ μ3.
Proof.
  intros.
  inversion H; steps;
    inversion H0; eauto with typ lia.
Qed.
Global Hint Resolve s_mode_trans: typ.

(* Subtyping *)

Inductive S_Typ : Tpe -> Tpe -> Prop :=
| s_typ_mode C μ1 μ2:
  μ1 ⊑ μ2 ->
  (C, μ1) <: (C, μ2)
where "T1 <: T2" := (S_Typ T1 T2).
Global Hint Resolve s_typ_mode: typ.

Ltac S_Typ:= repeat match goal with | H: _ <: _ |- _ => inverts H end.
Global Hint Extern 1 => S_Typ: typ. (* We want to invert subtyping as it is trivial *)

(* Subtyping between lists of types *)

Definition S_Typs := Forall2 (fun T1 T2 => S_Typ T1 T2).

(* Basic results *)

Lemma s_typ_refl:
  forall T, T <: T.
Proof.
  destruct T, m; eauto with typ.
Qed.
Global Hint Resolve s_typ_refl: typ.

Lemma s_typ_trans:
  forall T1 T2 T3, T1 <: T2 -> T2 <: T3 -> T1 <: T3.
Proof.
  intros; inversion H; inversion H0; steps; eauto with typ.
Qed.
Global Hint Resolve s_typ_trans: typ.

Lemma S_typs_refl:
  forall Ts, S_Typs Ts Ts.
Proof.
  intros. unfold S_Typs.
  induction Ts; eauto 3 using Forall2_cons with typ.
Qed.
Global Hint Resolve S_typs_refl: typ.

(* We then define predicate to check that all locations of a list are hot *)

Definition P_hot (T: Tpe) :=
  match T with
  | (_, hot) => True
  | _ => False
  end.
Definition P_hots := Forall P_hot.

Lemma P_hots_In:
  forall Ts T, P_hots Ts -> List.In T Ts -> P_hot T.
Proof.
  unfold P_hots; intros; eapply Forall_forall; eauto.
Qed.

(* ------------------------------------------------------------------------ *)
(** ** Expression Typing *)

Inductive expr_typing : EnvTyping -> Tpe -> Expr -> Tpe -> Prop :=

(* We have ambient subtyping *)
| t_sub:
    forall Γ Tthis e T T',
      (Γ, Tthis) ⊢ e : T' ->
      (T' <: T) ->
      (Γ, Tthis) ⊢ e : T

(* Lookup of a variable *)
| t_var:
    forall Γ Tthis x T,
      typeLookup Γ x = Some T ->
      (Γ, Tthis) ⊢ (e_var x) : T

(* Type of current object *)
| t_this:
    forall Γ Tthis,
      (Γ, Tthis) ⊢ e_this : Tthis

(* Accessing a field of an hot object gives an hot object *)
| t_selhot:
    forall Γ Tthis e f C D μ,
      (Γ, Tthis) ⊢ e : (D, hot) ->
      (fieldType D f = Some (C, μ)) ->
      (Γ, Tthis) ⊢ (e_fld e f) : (C, hot)

(* In other cases, accessing a defined fields give back its base type*)
| t_selcool:
    forall Γ Tthis e f T D Ω,
      (Γ, Tthis) ⊢ e : (D, cool Ω) ->
      f < Ω ->
      (fieldType D f = Some T) ->
      (Γ, Tthis) ⊢ (e_fld e f) : T

(* A newly created instance with hot arguments is also hot *)
| t_new_hot:
    forall Γ Tthis C args argsTs Args Flds Mtds,
      ct C = class Args Flds Mtds ->
      (Γ, Tthis) ⊩ args : argsTs ->
      P_hots argsTs ->
      S_Typs argsTs Args ->
      (Γ, Tthis) ⊢ (e_new C args) : (C, hot)

(* At least, a newly created instance is warm*)
| t_new:
    forall Γ Tthis C args Args Flds Mtds,
      ct C = class Args Flds Mtds ->
      (Γ, Tthis) ⊩ args : Args ->
      (Γ, Tthis) ⊢ (e_new C args) : (C, warm)

(* We restrict assignment to hot values to preserve monotonicity*)
| t_block:
    forall Γ Tthis T e1 f e2 e3 C μ,
      (Γ, Tthis) ⊢ (e_fld e1 f) : (C, μ) ->
      (Γ, Tthis) ⊢ e2 : (C, hot) ->
      (Γ, Tthis) ⊢ e3 : T ->
      (Γ, Tthis) ⊢ (e_asgn e1 f e2 e3) : T

(* Calling a method on an hot object with hot arguments produces a hot object*)
| t_call_hot:
    forall Γ Tthis e m args argTs paramTs D body μ0 μ_r C,
      (Γ, Tthis) ⊢ e : (C, hot) ->
      methodInfo C m = Some (μ0, paramTs, (D, μ_r), body) ->
      (Γ, Tthis) ⊩ args : argTs ->
      P_hots argTs ->
      S_Typs argTs paramTs ->
      (Γ, Tthis) ⊢ (e_mtd e m args) : (D, hot)

(* Otherwise, it produces the base type *)
| t_call:
    forall Γ Tthis e m args argsT T e__m μ0 μ__m C,
      (Γ, Tthis) ⊢ e : (C, μ0) ->
      μ0 ⊑ μ__m ->
      methodInfo C m = Some (μ__m, argsT, T, e__m) ->
      (Γ, Tthis) ⊩ args : argsT ->
      (Γ, Tthis) ⊢ (e_mtd e m args) : T

where  "( Γ , Tthis )  ⊢ e : T" := (expr_typing Γ Tthis e T)

with expr_list_typing: EnvTyping -> Tpe -> (list Expr) -> (list Tpe) -> Prop :=
| t_exprs_nil: forall Γ Tthis, expr_list_typing Γ Tthis [] []
| t_exprs_cons: forall Γ Tthis Ts es T e,
    (Γ, Tthis) ⊩ es : Ts ->
    (Γ, Tthis) ⊢ e : T ->
    (Γ, Tthis) ⊩ (e::es) : (T::Ts)
where  "( Γ , Tthis )  ⊩ es : Ts" := (expr_list_typing Γ Tthis es Ts).

Global Hint Constructors expr_typing: typ.

(* We add this lemma as if it were a rule (actually a consequence of t_selcool) *)

Lemma t_selwarm:
    forall Γ Tthis e f U D,
      (Γ, Tthis) ⊢ e : (D, warm) ->
      (fieldType D f = Some U) ->
      (Γ, Tthis) ⊢ (e_fld e f) : U.
Proof.
  eauto with typ.
Qed.

(* ------------------------------------------------------------------------ *)
(** ** Field typing *)

Definition T_Field (Γ:EnvTyping) T f :=
  match f with
  | field U e => (Γ, T) ⊢ e : U
  end.

Inductive T_Fields: EnvTyping -> Tpe -> (list Field) -> Prop :=
| t_fields_nil: forall Γ Tthis, T_Fields Γ Tthis []
| t_fields_cool_cons: forall Γ C Ω f fs,
    T_Field Γ (C, cool Ω) f ->
    T_Fields Γ (C, cool (S Ω)) fs ->
    T_Fields Γ (C, cool Ω) (f::fs).

Lemma T_Fields_In:
  forall Γ C Flds1 f Flds2 Ω,
    T_Fields Γ (C, cool Ω) (Flds1 ++ f :: Flds2) ->
    T_Field Γ (C, cool (Ω+dom Flds1)) f.
Proof with (eauto using T_Fields with typ lia).
  induction  Flds1; intros.
  - simpl in H. simpl. inverts H. rewrite <- plus_n_O...
  - simpl in H; simpl. inverts H. rewrite <- plus_Snm_nSm ...
Qed.

(* ------------------------------------------------------------------------ *)
(** ** Class typing *)

Definition T_Classes := forall C,
    match (ct C) with
    | class Args Flds Mtds =>
        T_Fields Args (C, cool 0) Flds /\
          (forall m μ ρ T e, Mtds m = Some (method μ ρ T e) -> (ρ, (C, μ)) ⊢ e : T)
    end.

(* ------------------------------------------------------------------------ *)
(** ** Program typing *)

Definition T_Prog :=
  T_Classes /\
  match EntryClass with
  | class nil nil Mtds => exists T e, Mtds main = Some (method hot nil T e)
  | _ => False
  end.

(* ------------------------------------------------------------------------ *)
(** * Inversion lemmas *)
(* Due to the ambient subtyping, we cannot invert typing judgments directly. We thus have to show
basic inversion lemmas. *)

Lemma t_this_inv:
  forall Γ U T,
    (Γ, U) ⊢ e_this : T -> U <: T.
Proof.
  intros. remember e_this as e.
  induction H; steps; eauto with typ.
Qed.

Lemma t_fld_inv:
  forall Γ T__this e f C μ,
    ((Γ, T__this) ⊢ (e_fld e f) : (C, μ)) ->
    exists D μ__e μ__f,
      ((Γ, T__this) ⊢ e : (D, μ__e)) /\
        fieldType D f = Some (C, μ__f) /\
        ((μ__e = hot) \/ (exists Ω, μ__e = cool Ω /\ f < Ω /\ μ__f ⊑ μ)).
Proof with (eauto with typ lia).
  introv H.
  remember (e_fld e f) as E.
  remember (C,μ) as T.
  gen C μ e f.
  induction H; steps...
  all: try destruct U' as [?C ?μ]; inverts H0.
  all: try specialize (IHT_Expr C μ0 eq_refl _ _ eq_refl)
    as (?D & ?μ__e & ?C & ?μ__f & ? & ? & [|[|]]); steps...
  all: repeat eexists...
  all: right; eexists...
Qed.

Lemma t_mtd_inv: forall Γ T__this e m args T,
    ((Γ, T__this) ⊢ (e_mtd e m args) : T) ->
    exists C D e__m argTs paramTs μ__m μ0 μ__r,
      ((Γ, T__this) ⊢ e : (C, μ0)) /\
        methodInfo C m = Some (μ__m, paramTs, (D, μ__r),  e__m) /\
        μ0 ⊑ μ__m /\
        ((Γ, T__this) ⊩ args : argTs) /\
        S_Typs argTs paramTs /\
        ((D, hot) <: T) /\
        ((D, μ__r) <: T \/ P_hots argTs /\ μ0 = hot).
Proof with (eauto with typ lia).
  introv H.
  remember (e_mtd e m args) as E.
  induction H; steps;
    try solve [repeat eexists; eauto with typ].
Qed.

Lemma t_new_inv:
  forall Γ T__this C args T,
    ((Γ, T__this) ⊢ (e_new C args) : T) ->
    exists Args Flds Mtds argsTs μ,
      T = (C, μ) /\
        ct C = class Args Flds Mtds /\
        ((Γ, T__this) ⊩ args : argsTs) /\
        S_Typs argsTs Args /\
        (warm ⊑ μ \/ (P_hots argsTs)).
Proof with (eauto with typ lia).
  introv H.
  remember (e_new C args) as E.
  induction H; subst; try discriminate...
  - S_Typ.
    destruct IHexpr_typing as (Args & Flds & Mtds & argsTs & μ0 & ? & ? & ? & ? & ?)...
    inverts H0; steps;
    repeat eexists; eauto with typ.
  - repeat eexists; steps...
  - inverts HeqE.
    exists Args Flds Mtds Args warm; splits ...
Qed.

Lemma t_asgn_inv:
  forall Γ T__this e1 f e2 e3 C μ,
    ((Γ, T__this) ⊢ (e_asgn e1 f e2 e3) : (C, μ)) ->
    exists D μ1 μ',
      ((Γ, T__this) ⊢ (e_fld e1 f) : (D, μ1)) /\
        ((Γ, T__this) ⊢ e2 : (D, hot)) /\
        (μ' ⊑ μ) /\
        ((Γ, T__this) ⊢ e3 : (C, μ')).
Proof with (eauto 3 with typ lia).
  introv H.
  remember (e_asgn e1 f e2 e3) as E.
  remember (C,μ) as T.
  gen C μ e1 f e2 e3.
  induction H; steps.
  - exists D, μ1, μ'; splits...
  - exists C, μ, μ0; splits...
Qed.

(* ------------------------------------------------------------------------ *)
(** ** Induction principle for typing *)

Section typing_ind.
  Variable P : forall Γ Tthis e U, ((Γ, Tthis) ⊢ e : U) -> Prop.
  Variable Pl : forall Γ Tthis el Ul, ((Γ, Tthis) ⊩ el : Ul) -> Prop.

  Variable P_sub: forall Γ Tthis e U U' H__U' H__sub,
      P Γ Tthis e U' H__U' ->
      P Γ Tthis e U (t_sub Γ Tthis e U U' H__U' H__sub).

  Variable P_var: forall Γ Tthis x U H__lkp,
      P Γ Tthis (e_var x) U (t_var Γ Tthis x U H__lkp).

  Variable P_this: forall Γ Tthis,
      P Γ Tthis e_this Tthis (t_this Γ Tthis).

  Variable P_selhot: forall Γ Tthis e f C D μ H__D H__field,
      P Γ Tthis e (D, hot) H__D ->
      P Γ Tthis (e_fld e f) (C, hot) (t_selhot Γ Tthis e f C D μ H__D H__field).

  Variable P_selcool: forall Γ Tthis e f U D Ω H__D H__f H__field,
      P Γ Tthis e (D, cool Ω) H__D ->
      P Γ Tthis (e_fld e f) U (t_selcool Γ Tthis e f U D Ω H__D H__f H__field).

  Variable P_new: forall Γ Tthis C args Args Flds Mtds H__ct H__args,
      Pl Γ Tthis args Args H__args ->
      P Γ Tthis (e_new C args) (C, warm) (t_new Γ Tthis C args Args Flds Mtds H__ct H__args).

  Variable P_new_hot: forall Γ Tthis C args argsTs Args Flds Mtds H__ct H__args H__hots H__subs,
      Pl Γ Tthis args argsTs H__args ->
      P Γ Tthis (e_new C args) (C, hot) (t_new_hot Γ Tthis C args argsTs Args Flds Mtds H__ct H__args H__hots H__subs).

  Variable P_block: forall Γ Tthis U e1 f e2 e3 C μ H__fld H__e2 H__e3,
      P Γ Tthis (e_fld e1 f) (C, μ) H__fld ->
      P Γ Tthis e2 (C, hot) H__e2 ->
      P Γ Tthis e3 U H__e3 ->
      P Γ Tthis (e_asgn e1 f e2 e3) U (t_block Γ Tthis U e1 f e2 e3 C μ H__fld H__e2 H__e3).

  Variable P_call: forall Γ Tthis e m args argsT U e__m μ0 μ__m C H__e H__sub H__mtd H__args,
      P Γ Tthis e (C, μ0) H__e ->
      Pl Γ Tthis args argsT H__args ->
      (* P argsT (C, μ__m) e__m U (typable_method m C μ__m argsT U e__m H__mtd) ->*)
      P Γ Tthis (e_mtd e m args) U (t_call Γ Tthis e m args argsT U e__m μ0 μ__m C H__e H__sub H__mtd H__args).

  Variable P_call_hot: forall Γ Tthis e m args argTs paramTs D body μ0 μ_r C H__e H__mtd H__args H__hots H__subs,
      P Γ Tthis e (C, hot) H__e ->
      Pl Γ Tthis args argTs H__args ->
      P Γ Tthis (e_mtd e m args) (D, hot) (t_call_hot Γ Tthis e m args argTs paramTs D body μ0 μ_r C H__e H__mtd H__args H__hots H__subs).

  Variable Pl_nil: forall Γ Tthis,
      Pl Γ Tthis [] [] (t_exprs_nil Γ Tthis).

  Variable Pl_cons: forall Γ Tthis U Ul e el H__el H__e,
      P Γ Tthis e U H__e ->
      Pl Γ Tthis el Ul H__el ->
      Pl Γ Tthis (e::el) (U::Ul) (t_exprs_cons Γ Tthis Ul el U e H__el H__e).

  Fixpoint typing_ind Γ Tthis e U (typing: (Γ, Tthis) ⊢ e : U) : P Γ Tthis e U typing :=
    match typing with
    | t_sub Γ Tthis e U U' H__U' H__sub => P_sub Γ Tthis e U U' H__U' H__sub (typing_ind Γ Tthis e U' H__U')
    | t_var Γ Tthis x U H__lkp => P_var Γ Tthis x U H__lkp
    | t_this Γ Tthis => P_this Γ Tthis
    | t_selhot Γ Tthis e f C D μ H__D H__field =>
        P_selhot Γ Tthis e f C D μ H__D H__field
                 (typing_ind Γ Tthis e (D, hot) H__D)
    | t_selcool Γ Tthis e f U D Ω H__D H__f H__field =>
        P_selcool Γ Tthis e f U D Ω H__D H__f H__field
                  (typing_ind Γ Tthis e (D, cool Ω) H__D)
    | t_new Γ Tthis C args Args Flds Mtds H__ct H__args =>
        P_new Γ Tthis C args Args Flds Mtds H__ct H__args
              (typing_list_ind Γ Tthis args Args H__args)
    | t_new_hot Γ Tthis C args argsTs Args Flds Mtds H__ct H__args H__hots H__subs =>
        P_new_hot Γ Tthis C args argsTs Args Flds Mtds H__ct H__args H__hots H__subs
                  (typing_list_ind Γ Tthis args argsTs H__args)
    | t_block Γ Tthis U e1 f e2 e3 C μ H__fld H__e2 H__e3 =>
        P_block Γ Tthis U e1 f e2 e3 C μ H__fld H__e2 H__e3
                (typing_ind Γ Tthis (e_fld e1 f) (C, μ) H__fld)
                (typing_ind Γ Tthis e2 (C, hot) H__e2)
                (typing_ind Γ Tthis e3 U H__e3)
    | t_call Γ Tthis e m args argsT U e__m μ0 μ__m C H__e H__sub H__mtd H__args =>
        P_call Γ Tthis e m args argsT U e__m μ0 μ__m C H__e H__sub H__mtd H__args
               (typing_ind Γ Tthis e (C, μ0) H__e)
               (typing_list_ind Γ Tthis args argsT H__args)
    | t_call_hot Γ Tthis e m args argTs paramTs D body μ0 μ_r C H__e H__mtd H__args H__hots H__subs =>
        P_call_hot Γ Tthis e m args argTs paramTs D body μ0 μ_r C H__e H__mtd H__args H__hots H__subs
                   (typing_ind Γ Tthis e (C, hot) H__e)
                   (typing_list_ind Γ Tthis args argTs H__args)
    end
  with typing_list_ind Γ Tthis el Ul (typing_list: (Γ, Tthis) ⊩ el : Ul) : Pl Γ Tthis el Ul typing_list :=
         match typing_list with
         | t_exprs_nil Γ Tthis => Pl_nil Γ Tthis
         | t_exprs_cons Γ Tthis Ul el U e H__el H__e =>
             Pl_cons Γ Tthis U Ul e el H__el H__e
                     (typing_ind Γ Tthis e U H__e)
                     (typing_list_ind Γ Tthis el Ul H__el)
         end.

End typing_ind.


Require Export Iron.Language.SystemF2Store.TyJudge.
Require Export Iron.Language.SystemF2Store.Exp.
Require Export Iron.Language.SystemF2Store.StoreValue.
Require Export Iron.Language.SystemF2Store.StoreBind.

(********************************************************************)
(* Store typing models the store.
   All types in the store typing have a corresponding binding in
   the store *)
Definition STOREM (ds: defs) (st: stenv) (s: store)
 := length st = length s.
Hint Unfold STOREM.

(********************************************************************)
(* Well typed store. *)
Definition STORET (ds: defs) (st: stenv) (ss: store)
 := forall l dcObj svFields
 ,  get l ss = Some (SObj dcObj svFields)
 -> (exists tcObj tsParam tsFields
    ,  get l st  = Some (makeTApps (TCon tcObj) tsParam)
    /\ hasDef ds (DefData dcObj tsFields tcObj)
    /\ Forall2 (TYPE ds nil nil st)
               (map expOfSValue svFields)
               (map (substTTs 0 tsParam) tsFields)).
Hint Unfold STORET.

(********************************************************************)
(* Well formed store.
   Store is well formed under some data type definitions and a
   store typing. *)
Definition WfS (ds: defs) (se: stenv) (ss: store)
 := DEFSOK ds
 /\ Forall closedT se
 /\ STOREM ds se ss
 /\ STORET ds se ss.
Hint Unfold WfS.

Lemma WfS_defsok
 : forall ds se ss, WfS ds se ss -> DEFSOK ds.
Proof. intros. inverts H. tauto. Qed.
Hint Resolve WfS_defsok.

Lemma WfS_closedT
 : forall ds se ss, WfS ds se ss -> Forall closedT se.
Proof. intros. inverts H. tauto. Qed.
Hint Resolve WfS_closedT.

(********************************************************************)
(* Lemmas about object fields *)

(* For each object in the store, the number of fields
   is the same as what is predicted by the data definition *)
Lemma storet_field_lengths
 :  forall ds dc se s l svs tsFields tObj
 ,  STORET ds se s
 -> get l s   = Some (SObj dc svs)
 -> hasDef ds (DefData dc tsFields tObj)
 -> length svs = length tsFields.
Proof.
 intros.
 unfold STORET in *.
 spec H l dc svs. rip.
  destruct H as [tcObj].
  destruct H as [tsParam].
  destruct H as [tsFields'].
  rip.
  defs_merge.
  have  ( length (map expOfSValue svs)
        = length (map (substTTs 0 tsParam) tsFields)).
  rewrite map_length in H.
  rewrite map_length in H.
 auto.
Qed.
Hint Resolve storet_field_lengths.

(* If we can get an object from the store,
   then we can also get any of the fields specified by its data def *)
Lemma storet_field_has
 :  forall ds s se svs tcObj tsFields i l dc
 ,  STORET ds se s
 -> get l s   = Some (SObj dc svs)
 -> hasDef ds (DefData dc tsFields tcObj)
 -> i < length tsFields
 -> exists svField vField
        ,  get i svs = Some svField
        /\ svalueOf vField svField.
Proof.
 intros.
 have (length svs = length tsFields).
 have (exists svField, get i svs = Some svField)
  by  (eapply get_length_less_exists; rs).
 shift svField.
 have (exists v, svalueOf v svField). eauto.
Qed.

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

(* If we can get an objects field,
   then that field has the type determined by the data def and
   store environment *)
Lemma storet_field_type
 :  forall ds se s l i svs vField svField tsParam tcObj tField tsFields dc
 ,  STORET ds se s
 -> get l s          = Some (SObj dc svs)
 -> get l se         = Some (makeTApps (TCon tcObj) tsParam)
 -> get i svs        = Some svField
 -> svalueOf vField svField
 -> hasDef ds (DefData dc tsFields tcObj)
 -> get i tsFields   = Some tField
 -> TYPE ds nil nil se vField (substTTs 0 tsParam tField).
Proof.
 intros.
 have (STORET ds se s).
 unfold STORET in H.
  spec H l dc svs. rip.
  destruct H as [tcObj'].
  destruct H as [tsParam'].
  destruct H as [tsFields'].
  rip. defs_merge.

 have (exists vsFields, Forall2 svalueOf vsFields svs) as HFs.
 destruct HFs as [vsFields].

 have (get i (map expOfSValue svs) = Some vField)
  by  (rw (vField = expOfSValue svField); eauto).

 have (get i (map (substTTs 0 tsParam) tsFields)
       = Some (substTTs 0 tsParam tField)).

 eapply Forall2_get_get_same; eauto.
Qed.

(* If we replace a field in a well typed store with one of the same
   type then the store is still well typed *)
Lemma storet_replace_field
 :  forall ds se s tField tField' vField1 svField1 vField2 svField2 svs i l dc tcObj
           tsFields tsParam
 ,  STORET ds se s
 -> get l s          = Some (SObj dc svs)
 -> get l se         = Some (makeTApps (TCon tcObj) tsParam)
 -> hasDef ds (DefData dc tsFields tcObj)
 -> get i tsFields   = Some tField'
 -> tField           = substTTs 0 tsParam tField'
 -> TYPE ds nil nil se vField1 tField -> svalueOf vField1 svField1
 -> TYPE ds nil nil se vField2 tField -> svalueOf vField2 svField2
 -> get i svs        = Some svField1
 -> STORET ds se (replace l (SObj dc (replace i svField2 svs)) s).
Proof.
 intros.
 unfold STORET in *. rip.

 have (l0 = l \/ l0 <> l) as HL.
 destruct HL.

 Case "l0 = l".
  subst.
  assert (dcObj = dc).
   erewrite replace_get_eq in H10.
   inverts H10. auto. eauto.

  subst.
  spec H H0.
  destruct H as [tcObj'].
  destruct H as [tsParam'].
  destruct H as [tsFields'].
  rip.
  exists tcObj.
  exists tsParam.
  exists tsFields.
  rip. defs_merge.

  have (svFields = replace i svField2 svs)
   by  (erewrite replace_get_eq in H10; eauto;
        inverts H10; auto; subst).

  eapply Forall2_map.
  eapply Forall2_map' in H12.

  have (vField1 = expOfSValue svField1). subst.
  have (vField2 = expOfSValue svField2). subst.

  have (tsFields = replace i tField' tsFields) as HR
   by  (symmetry; eauto). rewrite HR. clear HR.

  eapply Forall2_replace; eauto.

 Case "l0 <> l".
  rewrite replace_get_neq in H10; auto.
Qed.

(********************************************************************)
(* When we extend the store and store typing with a new binding,
   then the resulting store is still well formed. *)
Lemma store_extended_wellformed
 :  forall ds se ss bo to dc svs tc tsParam
 ,  WfS ds se ss
 -> bo = SObj dc svs
 -> to = makeTApps (TCon tc) tsParam
 -> TYPEB ds se bo to
 -> WfS ds (to <: se) (bo <: ss).
Proof.
 intros ds se ss bo to dc svs tc tsParam HW HBO HTO HT.
 subst.

 inverts HT.
 assert (tc0 = tc /\ tsParam0 = tsParam).
  eapply makeTApps_eq_params. eauto. rip. clear H1.

 (* Extended store typing is still closed *)
 assert (Forall closedT (makeTApps (TCon tc) tsParam <: se)).
  assert (Forall closedT tsParam).
   unfold closedT.
   rrwrite (0 = length (@nil ki)).
   eapply kind_wfT_Forall2.
   eauto.
  have (closedT (makeTApps (TCon tc) tsParam))
   by  (eapply makeTApps_wfT; eauto).
  eauto.

 (* Extended store typing models extended store *)
 assert (STOREM ds (makeTApps (TCon tc) tsParam <: se) (SObj dc svs <: ss)).
  unfold STOREM.

  have (length ss = length se)
   by  (unfold WfS in *; burn).

  repeat (rewrite length_simpl_snoc).
  eauto.

 (* Extended store is well typed under extended store typing *)
 assert (STORET ds (makeTApps (TCon tc) tsParam <: se) (SObj dc svs <: ss)).
   inverts HW. rip.
   unfold STORET in *.
   intros.

   assert (l <= length ss) as HL.
    assert (l < length (SObj dc svs <: ss)).
     eauto. rr. omega.

   inverts HL.

   SCase "l = length ss".
    assert (dcObj = dc /\ svFields = svs).
     rr. split; congruence. rip.

    exists tc tsParam tsFields. rip.
    have (length ss = length se) as HL.
     rewrite HL. rr. auto.
    have (xs = map expOfSValue svs) as HV.
    rewrite <- HV.
    eapply (Forall2_impl (TYPE ds nil nil se)); eauto.

   SCase "l < length s".
    assert (get l ss = Some (SObj dcObj svFields)) as HG.
     have (l < length ss).
     rewrite get_length_less_snoc in H9; auto.

    lets D: H10 l dcObj svFields HG. clear H10.
     destruct D as [tcObj].       exists tcObj.
     destruct H10 as [tsParam'].  exists tsParam'.
     destruct H10 as [tsFields']. exists tsFields'.
     rip.

    eapply (Forall2_impl (TYPE ds nil nil se)); eauto.

 (* Build WfS out of previous assertions *)
 inverts HW.
 auto.
Qed.
Hint Resolve store_extended_wellformed.

(********************************************************************)
(* Lemmas about existence of store components from WfS *)

Lemma store_has_sbind_for_stenv
 :  forall ds se ss l tObj
 ,  WfS ds se ss
 -> get l se = Some tObj
 -> (exists dc svs, get l ss = Some (SObj dc svs)).
Proof.
 intros.
 inverts H. rip.
 have (length se = length ss).
 have (exists sb, get l ss = Some sb).
 dest sb. destruct sb. eauto.
Qed.
Hint Resolve store_has_sbind_for_stenv.

Lemma store_has_sbind_for_XLoc
 :  forall ds ke te se ss l tObj
 ,  WfS ds se ss
 -> TYPE ds ke te se (XLoc l) tObj
 -> (exists dc svs, get l ss = Some (SObj dc svs)).
Proof.
 intros.
 inverts keep H. rip.
 inverts_type.
 eauto.
Qed.
Hint Resolve store_has_sbind_for_XLoc.

(* If we have a well typed case match on a store location containing
   some data object, then there is a case alternative corresponding to
   that object's data constructor. *)
Lemma store_has_sbind_for_XLoc_and_alt
 :  forall ds se ss l alts t
 ,  WfS ds se ss
 -> TYPE ds nil nil se (XCase (XLoc l) alts) t
 -> (exists dc, (exists svs, get    l  ss   = Some (SObj dc svs))
            /\  (exists x,   getAlt dc alts = Some (AAlt dc x))).
Proof.
 intros.
 inverts H0.

 have (exists dc svs, get l ss = Some (SObj dc svs)).
  shift dc. split.
  shift svs. auto.
  dest svs.

 eapply getAlt_exists.
 inverts_type.
 nforall.
 inverts H. rip.
 unfold STORET in *.
  spec H12 l dc svs. rip.
  destruct H12 as [tcObj'].
  dest tsParam.
  dest tsFields.
  rip.
  rewrite H7 in H12. inverts H12.
  erewrite getCtorOfType_makeTApps with (tc := tcObj') in H5; eauto.
    inverts H5.
  erewrite getCtorOfType_makeTApps with (tc := tcObj) in H9; eauto.
Qed.
Hint Resolve store_has_sbind_for_XLoc_and_alt.


(**

Instantiate the hypotheses of InitialHSS with Lam for HSET.

Written by: Anders Mörtberg, 2016

version for simplified notion of HSS by Ralph Matthes (2022, 2023)
the file is identical to the homonymous file in the parent directory, except for importing files from the present directory

*)

Require Import UniMath.Foundations.PartD.

Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Functors.
Local Open Scope cat.
Require Import UniMath.CategoryTheory.FunctorCategory.
Require Import UniMath.CategoryTheory.whiskering.
Require Import UniMath.CategoryTheory.limits.binproducts.
Require Import UniMath.CategoryTheory.limits.bincoproducts.
Require Import UniMath.CategoryTheory.limits.terminal.
Require Import UniMath.CategoryTheory.limits.initial.
Require Import UniMath.CategoryTheory.FunctorAlgebras.
Require Import UniMath.CategoryTheory.PointedFunctors.
Require Import UniMath.CategoryTheory.PrecategoryBinProduct.
Require Import UniMath.SubstitutionSystems.Signatures.
Require Import UniMath.CategoryTheory.Monads.Monads.
Require Import UniMath.SubstitutionSystems.BinSumOfSignatures.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.SubstitutionSystems.
Require Import UniMath.SubstitutionSystems.LamSignature.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.LiftingInitial_alt.
Require Import UniMath.SubstitutionSystems.SimplifiedHSS.MonadsFromSubstitutionSystems.
Require Import UniMath.SubstitutionSystems.Notation.
Local Open Scope subsys.

Require Import UniMath.CategoryTheory.limits.graphs.colimits.
Require Import UniMath.CategoryTheory.exponentials.
Require Import UniMath.CategoryTheory.categories.HSET.Core.
Require Import UniMath.CategoryTheory.categories.HSET.Limits.
Require Import UniMath.CategoryTheory.categories.HSET.Colimits.
Require Import UniMath.CategoryTheory.categories.HSET.Structures.
Require Import UniMath.CategoryTheory.Chains.All.

Section LamHSET.

Let Lam_S : Signature HSET _ _ :=
  Lam_Sig HSET TerminalHSET BinCoproductsHSET BinProductsHSET.

Local Notation "'EndHSET'":= ([HSET, HSET]) .

Local Lemma is_omega_cocont_Lam_S : is_omega_cocont Lam_S.
Proof.
apply is_omega_cocont_Lam.
* apply is_omega_cocont_constprod_functor1.
  apply Exponentials_functor_HSET.
* apply ColimsHSET_of_shape.
Defined.

Lemma Lam_Initial_HSET : Initial (category_FunctorAlg (Id_H _ BinCoproductsHSET Lam_S)).
Proof.
use colimAlgInitial.
- apply (Initial_functor_precat _ _ InitialHSET).
- unfold Id_H, Const_plus_H.
  apply is_omega_cocont_BinCoproduct_of_functors.
  + apply is_omega_cocont_constant_functor.
  + apply is_omega_cocont_Lam_S.
- apply ColimsFunctorCategory_of_shape; apply ColimsHSET_of_shape.
Defined.

(* Lemma KanExt_HSET : ∏ Z : precategory_Ptd HSET has_homsets_HSET, *)
(*    RightKanExtension.GlobalRightKanExtensionExists HSET HSET *)
(*      (U Z) HSET has_homsets_HSET has_homsets_HSET. *)
(* Proof. *)
(* intro Z. *)
(* apply RightKanExtension_from_limits, LimsHSET. *)
(* Defined. *)

Definition LamHSS_Initial_HSET : Initial (hss_category BinCoproductsHSET Lam_S).
Proof.
apply InitialHSS.
- apply InitialHSET.
- apply ColimsHSET_of_shape.
- apply is_omega_cocont_Lam_S.
Defined.

Definition LamMonad : Monad HSET.
Proof.
use Monad_from_hss.
- apply BinCoproductsHSET.
- apply Lam_S.
- apply LamHSS_Initial_HSET.
Defined.

End LamHSET.


(* ========================================================================= *)
(** * Biequivalence between Monads and Ktriples                              *)
(* ========================================================================= *)

Require Import UniMath.Foundations.All.
Require Import UniMath.MoreFoundations.All.
Require Import UniMath.CategoryTheory.Core.Prelude.
Require Import UniMath.CategoryTheory.FunctorCategory.
Require Import UniMath.CategoryTheory.Core.Isos.
Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Univalence.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.Core.NaturalTransformations.
Require Import UniMath.CategoryTheory.whiskering.
Require Import UniMath.CategoryTheory.opp_precat.
Require Import UniMath.CategoryTheory.DisplayedCats.Core.
Require Import UniMath.Bicategories.Core.Examples.BicatOfUnivCats.
Require Import UniMath.Bicategories.Core.Bicat. Import Bicat.Notations.
Require Import UniMath.Bicategories.Morphisms.Adjunctions.
Require Import UniMath.Bicategories.Core.AdjointUnique.
Require Import UniMath.Bicategories.Core.Invertible_2cells.
Require Import UniMath.Bicategories.Core.Univalence.

Require Import UniMath.Bicategories.PseudoFunctors.Display.PseudoFunctorBicat.
Require Import UniMath.Bicategories.PseudoFunctors.PseudoFunctor.
Import PseudoFunctor.Notations.

Require Import UniMath.Bicategories.DisplayedBicats.DispBicat. Import DispBicat.Notations.
Require Import UniMath.Bicategories.DisplayedBicats.DispAdjunctions.
Require Import UniMath.Bicategories.DisplayedBicats.DispInvertibles.
Require Import UniMath.Bicategories.DisplayedBicats.DispUnivalence.

Require Import UniMath.Bicategories.PseudoFunctors.Examples.Identity.
Require Import UniMath.Bicategories.PseudoFunctors.Examples.Composition.
Require Import UniMath.Bicategories.PseudoFunctors.Examples.Projection.
Require Import UniMath.Bicategories.Transformations.PseudoTransformation.
Require Import UniMath.Bicategories.Transformations.Examples.AlgebraMap.
Require Import UniMath.Bicategories.Modifications.Modification.
Require Import UniMath.Bicategories.PseudoFunctors.Biequivalence.

Require Import UniMath.Bicategories.Core.Unitors.
Require Import UniMath.Bicategories.DisplayedBicats.DispPseudofunctor.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.Algebras.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.Add2Cell.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.Prod.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.FullSub.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.Monads.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.KleisliTriple.
Require Import UniMath.Bicategories.DisplayedBicats.DispBuilders.
Require Import UniMath.Bicategories.DisplayedBicats.DispTransformation.
Require Import UniMath.Bicategories.DisplayedBicats.DispModification.
Require Import UniMath.Bicategories.DisplayedBicats.DispBiequivalence.
Require Import UniMath.Bicategories.Transformations.Examples.Unitality.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.KleisliTriple.

Require Import UniMath.CategoryTheory.Equivalences.CompositesAndInverses.

Local Open Scope cat.
Local Open Scope bicategory_scope.

(* ------------------------------------------------------------------------- *)
(** ** Miscellanea                                                           *)
(* ------------------------------------------------------------------------- *)

Definition isaprop_eq_2cell
           {B : bicat}
           {a b : B}
           {f g : a --> b}
           (x y : f ==> g)
  : isaprop (x = y).
Proof.
  apply B.
Defined.

(* ------------------------------------------------------------------------- *)
(** ** Basic monadic constructions built from a Ktriple.                     *)
(* ------------------------------------------------------------------------- *)

Section Monad_of_Kleisli_Data.

Context {x : category} (k : kleisli_triple x).

Local Lemma unit_kleisli_natural
  : is_nat_trans (functor_identity_data x)
                 (functor_data_of_kleisli_triple k)
                 (unit_kt k).
Proof.
  intros a b f. cbn.
  refine (!_).
  apply (unit_bind k).
Qed.

Definition unit_kleisli
  : functor_identity x ⟹ functor_of_kleisli_triple k
  := make_nat_trans (functor_identity x)
                    (functor_of_kleisli_triple k)
                    (unit_kt k)
                    unit_kleisli_natural.

Local Lemma mu_kleisli_natural
  : is_nat_trans (functor_composite_data (functor_of_kleisli_triple k)
                                         (functor_of_kleisli_triple k))
                 (functor_of_kleisli_triple k)
                 (λ a, bind_kt k (identity (k a))).
Proof.
  intros a b f. cbn.
  do 2 rewrite (bind_bind k).
  apply maponpaths.
  rewrite assoc'.
  rewrite (unit_bind k).
  rewrite id_left.
  apply id_right.
Qed.

Definition mu_kleisli
  : functor_of_kleisli_triple k ∙ functor_of_kleisli_triple k
    ⟹
    functor_of_kleisli_triple k
  := make_nat_trans (functor_of_kleisli_triple k ∙ functor_of_kleisli_triple k)
                    (functor_of_kleisli_triple k)
                    _
                    mu_kleisli_natural.
End Monad_of_Kleisli_Data.

Section Monad_of_Kleisli_Data.

Context {x : univalent_category} (k : kleisli_triple x).

Definition unit_mu_kleisli
  : monad bicat_of_univ_cats x.
Proof.
  use make_cat_monad.
  - exact (functor_of_kleisli_triple k).
  - exact (unit_kleisli k).
  - exact (mu_kleisli k).
  - abstract
      (cbn; intros;
       rewrite (bind_bind k);
       rewrite assoc';
       rewrite (unit_bind k);
       rewrite id_right;
       apply (bind_unit k)).
  - abstract
      (cbn; intros;
       apply (unit_bind k)).
  - abstract
      (cbn; intros;
       do 2 rewrite (bind_bind k);
       rewrite id_left;
       apply maponpaths;
       rewrite assoc';
       rewrite (unit_bind k);
       rewrite id_right; apply idpath).
Defined.

End Monad_of_Kleisli_Data.

(* ------------------------------------------------------------------------- *)
(** ** Pseudofunctor Ktriples to Monads.                                     *)
(* ------------------------------------------------------------------------- *)

Definition functor_of_kleisli_comm
      {x y : univalent_category}
      {f : x ⟶ y}
      {kx : kleisli_triple x}
      {ky : kleisli_triple y}
      (kf : kleisli_triple_on_functor kx ky f)
  : (functor_of_kleisli_triple kx ∙ f)
    ⟹
    (f ∙ functor_of_kleisli_triple ky).
Proof.
  use make_nat_trans.
  - exact (λ a, inv_from_z_iso (pr1 kf a)).
  - abstract
      (intros a b p;
       cbn;
       pose (pr22 kf) as H;
       cbn in H;
       rewrite H;
       rewrite !assoc';
       apply maponpaths;
       rewrite z_iso_inv_after_z_iso;
       rewrite id_right;
       apply maponpaths;
       rewrite functor_comp;
       rewrite !assoc';
       apply maponpaths;
       rewrite (pr12 kf);
       rewrite assoc';
       rewrite z_iso_inv_after_z_iso;
       apply id_right).
Defined.

Lemma functor_of_kleisli_comm_nat_z_iso
      {x y : univalent_category}
      {f : x ⟶ y}
      {kx : kleisli_triple x}
      {ky : kleisli_triple y}
      (kf : kleisli_triple_on_functor kx ky f)
  : is_nat_z_iso (functor_of_kleisli_comm kf).
Proof.
  intro a.
  apply is_z_iso_inv_from_z_iso.
Defined.

Definition functor_of_kleisli_z_iso
           {x y : univalent_category}
           {f : x ⟶ y}
           {kx : kleisli_triple x}
           {ky : kleisli_triple y}
           (kf : kleisli_triple_on_functor kx ky f)
  : nat_z_iso (functor_of_kleisli_triple kx ∙ f)
            (f ∙ functor_of_kleisli_triple ky).
Proof.
  use make_nat_z_iso.
  - exact (functor_of_kleisli_comm kf).
  - exact (functor_of_kleisli_comm_nat_z_iso kf).
Defined.

Definition unit_mu_kleisli_functor
           {C D : univalent_category}
           {F : C ⟶ D}
           {KC : kleisli_triple_disp_bicat C}
           {KD : kleisli_triple_disp_bicat D}
           (KF : KC -->[F] KD)
  : unit_mu_kleisli KC -->[ F] unit_mu_kleisli KD.
Proof.
  use make_cat_monad_mor ; cbn.
  - exact (functor_of_kleisli_z_iso KF).
  - abstract
      (intros X ; cbn ;
       rewrite (kleisli_triple_on_functor_unit_kt KF);
       rewrite assoc';
       rewrite z_iso_inv_after_z_iso;
       rewrite id_right;
       apply idpath).
  - abstract
      (intros X ; cbn ;
       rewrite (kleisli_triple_on_functor_bind_kt KF);
       rewrite !assoc';
       apply maponpaths;
       rewrite (bind_bind KD);
       rewrite !assoc';
       rewrite (unit_bind KD), id_right;
       rewrite z_iso_inv_after_z_iso, id_right;
       rewrite functor_id, id_left;
       apply idpath).
Defined.

Definition Ktriple_to_Monad
  : disp_psfunctor kleisli_triple_disp_bicat
                   (monad bicat_of_univ_cats)
                   (id_psfunctor bicat_of_univ_cats).
Proof.
  use make_disp_psfunctor.
  - apply disp_2cells_isaprop_monad.
    apply univalent_cat_is_univalent_2.
  - exact (disp_locally_groupoid_monad
             bicat_of_univ_cats
             univalent_cat_is_univalent_2).
  - exact @unit_mu_kleisli.
  - exact @unit_mu_kleisli_functor.
  - abstract
      (cbn; intros x y f g α kx ky kf kg e;
       refine ((_,, (tt,, tt)),, tt);
       use nat_trans_eq; try apply homset_property;
       cbn; intro a;
       apply pathsinv0;
       apply z_iso_inv_on_left;
       rewrite assoc';
       rewrite <- e;
       rewrite assoc;
       rewrite z_iso_after_z_iso_inv;
       apply pathsinv0;
       apply id_left).
  - abstract
      (intros x kx;
       refine ((_,, (tt,, tt)),, tt);
       unfold alg_disp_cat_2cell;
       use nat_trans_eq; try apply homset_property;
       intro a; cbn;
       rewrite !id_left;
       rewrite (bind_bind kx);
       rewrite (unit_bind kx);
       apply pathsinv0;
       apply (bind_unit kx)).
  - abstract
      (simpl;
       intros x y z f g kx ky kz kf kg;
       refine ((_,, (tt,, tt)),, tt);
       use nat_trans_eq; try apply homset_property;
       intro a; cbn;
       rewrite !id_left;
       rewrite !id_right;
       rewrite assoc';
       rewrite (bind_bind kz);
       rewrite (unit_bind kz);
       rewrite (bind_unit kz);
       rewrite id_right;
       apply idpath).
Defined.

(* ------------------------------------------------------------------------- *)
(*  Kleisli_of_Monad                                                         *)
(* ------------------------------------------------------------------------- *)

Definition Monad_to_Ktriple_data {x : univalent_category}
           (m : monad bicat_of_univ_cats x)
  : kleisli_triple_disp_bicat (id_psfunctor bicat_of_univ_cats x).
Proof.
  use make_kleisli_triple.
  - apply m.
  - exact (pr1 (monad_unit m)).
  - exact (λ _ _ F, monad_bind m F).
  - intros A. apply (cat_monad_unit_bind m).
  - simpl. intros. apply (cat_monad_bind_unit m).
  - simpl. intros. apply (cat_monad_bind_bind m).
Defined.

(* NB: We need to take the inverse to match the definition used to build the biequivalence. *)
Definition monad_mor_natural_pointwise
           {C₁ C₂ : univalent_category}
           {F : C₁ ⟶ C₂}
           {M₁ : monad bicat_of_univ_cats C₁}
           {M₂ : monad bicat_of_univ_cats C₂}
           (FF : M₁ -->[F] M₂)
           (X : C₁)
  : z_iso ((monad_endo M₂ : C₂ ⟶ C₂) (F X)) (F ((monad_endo M₁ : C₁ ⟶ C₁) X))
  := CompositesAndInverses.nat_z_iso_to_pointwise_z_iso
       (nat_z_iso_inv (monad_mor_nat_z_iso FF)) X.

Lemma inv_monad_mor_natural_pointwise
      {C₁ C₂ : univalent_category}
      {F : C₁ ⟶ C₂}
      {M₁ : monad bicat_of_univ_cats C₁}
      {M₂ : monad bicat_of_univ_cats C₂}
      (FF : M₁ -->[F] M₂)
      (X : C₁)
  : inv_from_z_iso (monad_mor_natural_pointwise FF X)
    =
    CompositesAndInverses.nat_z_iso_to_pointwise_z_iso (monad_mor_nat_z_iso FF) X.
Proof.
  refine (!_).
  apply inv_z_iso_unique'.
  unfold precomp_with.
  cbn.
  apply (nat_trans_eq_pointwise (vcomp_linv (monad_mor_natural FF)) X).
Qed.

Definition Monad_to_Ktriple_functor
           {x y : univalent_category}
           {f : bicat_of_univ_cats ⟦ x, y ⟧}
           {mx : (monad bicat_of_univ_cats) x}
           {my : (monad bicat_of_univ_cats) y}
           (mf : mx -->[ f] my)
  : Monad_to_Ktriple_data mx -->[ f ] Monad_to_Ktriple_data my.
Proof.
  use make_kleisli_triple_on_functor.
  - exact (monad_mor_natural_pointwise mf).
  - abstract (
        refine (λ (X : x), _); simpl;
        pose (nat_trans_eq_pointwise (monad_mor_unit mf) X) as mf_unit;
        cbn in mf_unit;
        do 2 rewrite id_left in mf_unit;
        etrans; [ apply pathsinv0 | apply maponpaths_2; exact mf_unit ];
        etrans; [ rewrite assoc' | apply id_right ];
      apply maponpaths;
        exact (z_iso_inv_after_z_iso (nat_z_iso_pointwise_z_iso (invertible_2cell_to_nat_z_iso _ _ (monad_mor_natural mf)) X))
      ).
  - abstract (
        refine (λ (X Y : x) (p : x ⟦ X, pr1 (Monad_to_Ktriple_data mx) Y ⟧), _);
        unfold Monad_to_Ktriple_data, bind_kt; simpl;
        etrans; [ apply (monad_mor_bind_alt mf) | idtac ];
        do 2 rewrite (inv_monad_mor_natural_pointwise mf);
        do 2 rewrite assoc';
        apply idpath
      ).
Defined.

(** the extra identities in the statement are no longer of use for [Monad_to_Ktriple] *)
Definition Monad_to_Ktriple_2cell
  : ∏ (x y : univalent_category)
      (f g : x ⟶ y)
      (α : prebicat_cells bicat_of_univ_cats f g)
      (mx : (monad bicat_of_univ_cats) x) (my : (monad bicat_of_univ_cats) y)
      (mf : mx -->[ f] my)
      (mg : mx -->[ g] my),
    mf ==>[ α] mg
    →
    ∏ X,
    pr1 ((pr2 (monad_mor_natural mf)) ^-1) X
        · id₁ (f (pr1 (monad_endo mx) X))
        · pr1 α ((pr111 (pr1 mx)) X) =
    monad_bind my (pr1 α X · pr1 (monad_unit my) (g X))
               · (pr1 ((pr2 (monad_mor_natural mg)) ^-1) X
                      · id₁ (g (pr1 (monad_endo mx) X))).
Proof.
  intros x y f g α mx my mf mg mα.
  refine (λ X: (x:univalent_category), _).
  rewrite !id_right.
  pose (nat_trans_eq_pointwise (pr11 mα) X) as d.
  pose (maponpaths (λ z, z · pr1 ((pr2 (monad_mor_natural mg)) ^-1) X) d) as p₁.
  cbn in p₁.
  rewrite assoc' in p₁.
  pose (maponpaths (λ z, pr1 α (pr1 (pr11 mx) X) · z)
                   (!(nat_trans_eq_pointwise
                        (vcomp_rinv (monad_mor_natural mg))
                        X))) as p₂.
  pose (!(id_right _) @ p₂ @ p₁) as r.
  refine (maponpaths (λ z, _ · z) r @ _).
  clear d p₁ p₂ r.
  rewrite !assoc.
  apply maponpaths_2.
  pose (maponpaths (λ z, z · # (pr111 my) (pr1 α X))
                   (!(nat_trans_eq_pointwise
                        (vcomp_linv (monad_mor_natural mf))
                        X))) as p.
  pose (!(id_left _) @ p) as r.
  refine (!r @ _).
  clear p r.
  apply cat_monad_map_as_bind.
Qed.

Definition Monad_to_Ktriple_identitor
  : ∏ (x : bicat_of_univ_cats) (xx : (monad bicat_of_univ_cats) x),
    (id_disp (Monad_to_Ktriple_data xx))
      ==>[ psfunctor_id (id_psfunctor bicat_of_univ_cats) x]
      Monad_to_Ktriple_functor (id_disp xx).
Proof.
  intros x mx X; cbn.
  do 2 rewrite id_left ; do 2 rewrite id_right.
  rewrite (functor_id (pr11 mx)), id_right.
  refine (!_).
  apply (cat_monad_unit_bind mx).
Qed.

Definition Monad_to_Ktriple_compositor
  : ∏ (x y z : univalent_category)
      (f : x ⟶ y) (g : y ⟶ z)
      (xx : (monad bicat_of_univ_cats) x)
      (yy : (monad bicat_of_univ_cats) y)
      (zz : (monad bicat_of_univ_cats) z)
      (ff : xx -->[ f] yy) (gg : yy -->[ g] zz),
    (Monad_to_Ktriple_functor ff;; Monad_to_Ktriple_functor gg)
      ==>[ id₂ _]
      Monad_to_Ktriple_functor (ff;; gg).
Proof.
  intros x y z f g mx my mz mf mg.
  refine (λ X : pr1 x, _).
  cbn.
  etrans.
  { apply id_right. }
  refine (!_).
  etrans.
  {
    etrans.
    {
      apply maponpaths_2.
      etrans.
      {
        apply maponpaths.
        apply id_left.
      }
      exact (cat_monad_unit_bind mz).
    }
    refine (id_left _ @ _).
    apply maponpaths.
    refine (id_left _ @ _).
    apply maponpaths.
    refine (id_left _ @ _).
    apply id_right.
  }
  etrans.
  {
    etrans.
    {
      apply maponpaths_2.
      exact (functor_id (pr11 mz) (g(f X))).
    }
    apply id_left.
  }
  apply idpath.
Qed. (* 32 seconds on my computer *)

Definition Monad_to_Ktriple
  : disp_psfunctor (monad bicat_of_univ_cats)
                   kleisli_triple_disp_bicat
                   (id_psfunctor bicat_of_univ_cats).
Proof.
  use make_disp_psfunctor.
  - exact disp_2cells_isaprop_kleisli.
  - exact disp_locally_groupoid_kleisli.
  - exact @Monad_to_Ktriple_data.
  - exact @Monad_to_Ktriple_functor.
  - intros x y f g α mx my mf mg mα X.
    assert (Hyp := Monad_to_Ktriple_2cell x y f g α mx my mf mg mα X).
    cbn.
    do 2 rewrite id_right in Hyp.
    assumption.
  - exact Monad_to_Ktriple_identitor.
  - exact Monad_to_Ktriple_compositor.
Defined.

Lemma bind_kt_monad_to_kleisli
      {x : univalent_category}
      (k : kleisli_triple x)
      {a b : x}
      (f : x ⟦ a, k b ⟧)
  : bind_kt (Monad_to_Ktriple_data (unit_mu_kleisli k)) f = bind_kt k f.
Proof.
  unfold bind_kt at 1; simpl.
  unfold monad_bind; simpl.
  rewrite (bind_bind k).
  apply maponpaths.
  etrans; [ idtac | apply id_right ].
  rewrite assoc'.
  apply maponpaths.
  apply (unit_bind k).
Qed.

Definition Monad_biequiv_Ktriple_unit
  : disp_pstrans
      (disp_pseudo_comp
         (id_psfunctor bicat_of_univ_cats) (id_psfunctor bicat_of_univ_cats)
         (monad bicat_of_univ_cats)
         kleisli_triple_disp_bicat
         (monad bicat_of_univ_cats)
         Monad_to_Ktriple
         Ktriple_to_Monad)
      (disp_pseudo_id (monad bicat_of_univ_cats))
      (lunitor_pstrans (id_psfunctor bicat_of_univ_cats)).
Proof.
  use make_disp_pstrans.
  - exact (disp_2cells_isaprop_monad
             bicat_of_univ_cats
             univalent_cat_is_univalent_2).
  - exact (disp_locally_groupoid_monad
             bicat_of_univ_cats
             univalent_cat_is_univalent_2).
  - intros.
    use make_cat_monad_mor.
    + simpl.
      cbn.
      use make_nat_z_iso.
      * use make_nat_trans.
        ** intro z. apply identity.
        ** abstract
             (intros z t f ; cbn;
              rewrite id_left, id_right;
              unfold monad_bind;
              rewrite (functor_comp (monad_endo xx : _ ⟶ _));
              rewrite assoc';
              etrans;
              [ apply maponpaths; apply (cat_monad_ημ xx)
              | apply id_right ]).
      * intros z. apply identity_is_z_iso.
    + intros z.
      apply id_right.
    + abstract (
          simpl;
          intros X;
          rewrite id_left;
          apply id_right).
  - abstract (
        intros;
        use make_cat_monad_cell;
        simpl;
        intros X;
        rewrite !(functor_id ((monad_endo yy) : _ ⟶ _));
        rewrite (functor_id f);
        rewrite !id_left;
        rewrite !(functor_id ((monad_endo yy) : _ ⟶ _));
        rewrite !id_right;
        apply pathsinv0;
        apply inv_z_iso_unique';
        unfold precomp_with;
        simpl;
        exact (z_iso_after_z_iso_inv (nat_z_iso_pointwise_z_iso (invertible_2cell_to_nat_z_iso _ _ (monad_mor_natural ff)) X))
      ).
Defined.

Definition Monad_bequiv_Ktriple_counit
  : disp_pstrans
      (disp_pseudo_comp
         (id_psfunctor bicat_of_univ_cats) (id_psfunctor bicat_of_univ_cats)
         kleisli_triple_disp_bicat
         (monad bicat_of_univ_cats)
         kleisli_triple_disp_bicat
         Ktriple_to_Monad Monad_to_Ktriple)
      (disp_pseudo_id kleisli_triple_disp_bicat)
      (lunitor_pstrans (id_psfunctor bicat_of_univ_cats)).
Proof.
  use make_disp_pstrans.
  - exact disp_2cells_isaprop_kleisli.
  - exact disp_locally_groupoid_kleisli.
  - refine (λ (x : univalent_category) (kx : kleisli_triple x), _).
    use make_kleisli_triple_on_functor.
    + exact (λ X, identity_z_iso (kx X)).
    + abstract (
          refine (λ A : x, _);
          apply pathsinv0;
          apply id_right).
    + abstract (
          refine (λ (A B : pr1 x) (f : pr1 x ⟦ A, pr1 kx B ⟧), _);
          simpl;
          rewrite id_right;
          etrans; [ apply bind_kt_monad_to_kleisli | idtac ];
          etrans;
          [ pose (kleisli_triple_on_functor_bind_kt
                    (kleisli_triple_on_identity_functor
                       kx)
                    _ _ f
                 ) as H;
            simpl in H;
            exact H
          | idtac ];
          apply id_right
        ).
  - abstract (
        refine (λ (x y : univalent_category)
                  (f : pr1 x ⟶ pr1 y)
                  (kx : kleisli_triple x)
                  (ky : kleisli_triple y)
                  (kf : kleisli_triple_on_functor kx ky f)
                  (X : x),
                _);
        simpl;
        apply pathsinv0;
        etrans;
        [ apply maponpaths_2;
          do 2 rewrite id_left;
          apply (bind_unit ky)
        | idtac ];
        etrans; [ apply id_left | idtac];
        etrans; [ apply id_left | idtac];
        apply pathsinv0;
        simpl;
        etrans;
        [ rewrite assoc';
          apply maponpaths;
          rewrite functor_id;
          etrans; [ apply id_left | idtac];
          apply id_left
        | idtac ];
          etrans; [ apply id_right | idtac];
        apply idpath
      ).
Defined.

Definition Monad_biequiv_Ktriple_unit_counit
  : is_disp_biequivalence_unit_counit
      (monad bicat_of_univ_cats)
      kleisli_triple_disp_bicat
      (id_is_biequivalence _) Monad_to_Ktriple Ktriple_to_Monad.
Proof.
  split.
  - exact Monad_biequiv_Ktriple_unit.
  - exact Monad_bequiv_Ktriple_counit.
Defined.

Definition Monad_biequiv_Ktriple_unit_inv
  : disp_pstrans
      (disp_pseudo_id (monad bicat_of_univ_cats))
      (disp_pseudo_comp
         (id_psfunctor bicat_of_univ_cats) (id_psfunctor bicat_of_univ_cats)
         (monad bicat_of_univ_cats)
         kleisli_triple_disp_bicat
         (monad bicat_of_univ_cats)
         Monad_to_Ktriple
         Ktriple_to_Monad)
      (linvunitor_pstrans (id_psfunctor bicat_of_univ_cats)).
Proof.
  use make_disp_pstrans.
  - exact (disp_2cells_isaprop_monad
             bicat_of_univ_cats
             univalent_cat_is_univalent_2).
  - exact (disp_locally_groupoid_monad
             bicat_of_univ_cats
             univalent_cat_is_univalent_2).
  - intros.
    use make_cat_monad_mor.
    + simpl.
      cbn.
      use make_nat_z_iso.
      * use make_nat_trans.
        ** intro z. apply identity.
        ** abstract
             (intros z t f ; cbn ;
              rewrite id_left, id_right ;
              apply cat_monad_map_as_bind).
      * intros z. apply identity_is_z_iso.
    + intros z.
      apply id_right.
    + abstract
        (simpl ;
         intros X ;
         rewrite !id_left, id_right ;
         rewrite bind_unit ;
         rewrite id_left ;
         cbn ;
         unfold monad_bind ;
         rewrite functor_id ;
         rewrite id_left ;
         apply idpath).
  - abstract
      (intros ;
       use make_cat_monad_cell ;
       simpl ;
       intro z ;
       rewrite !id_left ;
       rewrite !id_right ;
       rewrite (functor_id f) ;
       rewrite id_left ;
       cbn ;
       rewrite <- assoc ;
       apply maponpaths ;
       refine (!_) ;
       refine (maponpaths (λ q, _ · q) (cat_monad_unit_bind _) @ _) ;
       apply id_right).
Defined.

Definition Monad_biequiv_Ktriple_counit_inv
  : disp_pstrans
      (disp_pseudo_id kleisli_triple_disp_bicat)
      (disp_pseudo_comp
         (id_psfunctor bicat_of_univ_cats) (id_psfunctor bicat_of_univ_cats)
         kleisli_triple_disp_bicat
         (monad bicat_of_univ_cats)
         kleisli_triple_disp_bicat
         Ktriple_to_Monad Monad_to_Ktriple)
      (linvunitor_pstrans (id_psfunctor bicat_of_univ_cats)).
Proof.
  use make_disp_pstrans.
  - exact disp_2cells_isaprop_kleisli.
  - exact disp_locally_groupoid_kleisli.
  - refine (λ (x : univalent_category) (kx : kleisli_triple x), _).
    use make_kleisli_triple_on_functor.
    + exact (λ X, identity_z_iso (kx X)).
    + abstract (
          refine (λ A : x, _);
          apply pathsinv0;
          apply id_right).
    + abstract
        (intros A B f ;
         simpl ;
         rewrite id_left, id_right ;
         refine (!_) ;
         etrans ; [ apply bind_kt_monad_to_kleisli | ] ;
         apply maponpaths ;
         apply id_right).
  - abstract
      ( intros x y f kx ky kf z ;
        simpl ;
        cbn;
        rewrite !id_left, !id_right;
        rewrite functor_id, id_right;
        etrans;
        [| apply cancel_postcomposition, pathsinv0,
          (cat_monad_unit_bind (unit_mu_kleisli ky))];
        apply pathsinv0, id_left
      ).
Defined.

Definition Monad_disp_biequiv_Ktriple
  : disp_is_biequivalence_data
      (monad bicat_of_univ_cats)
      kleisli_triple_disp_bicat
      (id_is_biequivalence _)
      Monad_biequiv_Ktriple_unit_counit.
Proof.
  simple refine (_ ,, _ ,, ((_ ,, _) ,, (_ ,, _))).
  - exact Monad_biequiv_Ktriple_unit_inv.
  - exact Monad_biequiv_Ktriple_counit_inv.
  - use make_disp_invmodification.
    + exact (disp_2cells_isaprop_monad
               bicat_of_univ_cats
               univalent_cat_is_univalent_2).
    + exact (disp_locally_groupoid_monad
               bicat_of_univ_cats
               univalent_cat_is_univalent_2).
    + abstract
        (intros x xx ;
         use make_cat_monad_cell ;
         intros z ;
         simpl ;
         rewrite !id_left ;
         rewrite (functor_id (pr11 xx)), (functor_id (monad_endo xx)) ;
         exact (!(id_left _))).
  - use make_disp_invmodification.
    + exact (disp_2cells_isaprop_monad
               bicat_of_univ_cats
               univalent_cat_is_univalent_2).
    + exact (disp_locally_groupoid_monad
               bicat_of_univ_cats
               univalent_cat_is_univalent_2).
    + abstract
        (intros x xx ;
         use make_cat_monad_cell ;
         intros z ;
         simpl ;
         rewrite !id_left ;
         refine (!_) ;
         refine (bind_bind
                   (Monad_to_Ktriple_data xx)
                   (unit_kt (Monad_to_Ktriple_data xx) z)
                   (unit_kt (Monad_to_Ktriple_data xx) z)
                 @ _) ;
         apply maponpaths ;
         apply unit_bind).
  - use make_disp_invmodification.
    + exact disp_2cells_isaprop_kleisli.
    + exact disp_locally_groupoid_kleisli.
    + abstract
        (intros x xx z ;
         simpl ;
         rewrite !id_left, id_right ;
         rewrite bind_unit ;
         apply idpath).
  - use make_disp_invmodification.
    + exact disp_2cells_isaprop_kleisli.
    + exact disp_locally_groupoid_kleisli.
    + abstract
        (intros x xx z ;
         simpl ;
         rewrite !id_left, id_right ;
         rewrite bind_unit ;
         apply idpath).
Defined.

Definition Monad_to_Ktriple_psfunctor
  : psfunctor
      (total_bicat (monad bicat_of_univ_cats))
      (total_bicat kleisli_triple_disp_bicat)
  := total_psfunctor
       (monad bicat_of_univ_cats)
       kleisli_triple_disp_bicat
       (id_psfunctor bicat_of_univ_cats)
       Monad_to_Ktriple.

Definition Monad_biequiv_Ktriple
  : is_biequivalence Monad_to_Ktriple_psfunctor
  := total_is_biequivalence
       _
       _
       _
       Monad_disp_biequiv_Ktriple.


Require Import Bmc.KInduction.
Require Import Bmc.Example.

Require Import SMTC.Tactic.
Require Import SMTC.Integers.

(*Definition k_ind_step_ni (T : trans)
  (P : prop) (k: nat) : Prop :=
  forall ss : sseq,
  ~(loop_free T ss 0 k /\ safety_nth_offset P ss 0 k /\ ~P ss.[k]).

Definition k_induction_post_ni (I : prop) (T : trans) (P : prop) (k: nat) : Prop :=
  k_ind_step_ni T P k /\ safety_nth_ni I T P k.
*)

Set SMT Solver "z3".
Set SMT Debug.

Axiom by_smt : forall P : Prop, P.

Goal k_induction_post_conj ex1_I ex1_T ex1_P 3.
Proof.
  unfold ex1_I, ex1_T, ex1_P.
  unfold k_induction_post_conj.
  unfold k_ind_step_conj.
  unfold safety_nth_offset.
  unfold safety_nth_conj.
  unfold prop_nth_init_conj.
  unfold path.
  unfold sseq.
  unfold nth.
  unfold state.
  repeat rewrite -> Nat.add_0_l.
  repeat rewrite -> Nat.add_0_r.

  split.
  - intros.
    smt solve; apply by_smt.

  - repeat split.
    + intros.
      smt solve; apply by_smt.
    + intros.
      smt solve; apply by_smt.
    + intros.
      smt solve; apply by_smt.
    + intros.
      smt solve; apply by_smt.
Qed.

Goal k_induction_post_conj ex2_I ex2_T ex2_P 2.
Proof.
  unfold ex2_I, ex2_T, ex2_P.
  unfold k_induction_post_conj, k_ind_step_conj, safety_nth_offset, safety_nth_conj, prop_nth_init_conj, path, sseq, nth, state.
  repeat rewrite -> Nat.add_0_l;
  repeat rewrite -> Nat.add_0_r.
  split.
  intros; smt solve; apply by_smt.
  repeat split; intros; smt solve; apply by_smt.
Qed.

Goal k_induction_post_conj ex3_I ex3_T ex3_P 5.
Proof.
  unfold ex3_I, ex3_T, ex3_P.
  unfold k_induction_post_conj, k_ind_step_conj, safety_nth_offset, safety_nth_conj, prop_nth_init_conj, path, sseq, nth, state.
  repeat rewrite -> Nat.add_0_l;
  repeat rewrite -> Nat.add_0_r.
  split.
  intros; smt solve; apply by_smt.
  repeat split; intros; smt solve; apply by_smt.
Qed.

(* eof *)


(* Lemmas used by the hmac proofs that are independent of C/CompCert/VST*)

Require Import compcert.lib.Integers.
Require Import compcert.lib.Coqlib.
Require Import List. Import ListNotations.
Require Import sha.general_lemmas.
Require Import sha.hmac_pure_lemmas.
Require Import sha.SHA256.
Require Import sha.pure_lemmas.     (* sha *)
Require Import sha.spec_sha.
(*Require Import sublist.*)

Require Import sha.HMAC_functional_prog.
Require Import sha.HMAC256_functional_prog.
(*
Lemma SF_ByteRepr x: isbyteZ x ->
                     HP.HMAC_SHA256.sixtyfour x =
                     map Byte.unsigned (HP.HMAC_SHA256.sixtyfour (Byte.repr x)).
Proof. intros. unfold HP.HMAC_SHA256.sixtyfour.
 rewrite map_list_repeat.
 rewrite Byte.unsigned_repr; trivial. destruct H.
 assert (BMU: Byte.max_unsigned = 255). reflexivity. omega.
Qed.
*)
Lemma str_to_Z_length: forall k,
      String.length k = length (str_to_Z k).
Proof. intros. induction k; simpl; auto. Qed.

Lemma first64_sixtyfour {A} (a:A):
      firstn 64 (HMAC_SHA256.sixtyfour a) = HMAC_SHA256.sixtyfour a.
Proof. apply firstn_precise. simpl. reflexivity. Qed.

Lemma length_SF {A} (x:A): length (HMAC_SHA256.sixtyfour x) = 64%nat.
Proof.
  unfold HMAC_SHA256.sixtyfour.
  rewrite length_list_repeat; trivial.
Qed.

Lemma Zlength_mkArgZ k pad: Zlength (HMAC_SHA256.mkArgZ k pad) = Z.of_nat (min (length k) 64).
Proof. intros. repeat rewrite Zlength_correct.
   unfold HMAC_SHA256.mkArgZ, HMAC_SHA256.mkArg, HMAC_SHA256.sixtyfour.
   repeat rewrite map_length.
   rewrite combine_length, length_list_repeat . trivial.
Qed.

Lemma zeropad_isbyteZ: forall l, Forall isbyteZ l -> Forall isbyteZ (HMAC_SHA256.zeroPad l).
Proof. unfold isbyteZ, HMAC_SHA256.zeroPad; intros.
  rewrite Forall_forall in *. intros.
  apply in_app_or in H0.
  destruct H0. apply (H _ H0).
  apply in_list_repeat in H0. (*apply In_Nlist in H0.*) subst; omega.
Qed.

Lemma nth_zeropad_left {d d'}: forall l i (I: 0<= i < Zlength l),
      nth (Z.to_nat i) (HMAC_SHA256.zeroPad l) d = nth (Z.to_nat i) l d'.
Proof. unfold HMAC_SHA256.zeroPad. intros.
  destruct I.
  apply Z2Nat.inj_lt in H0; try omega.
  rewrite Zlength_correct, Nat2Z.id in H0.
  rewrite app_nth1; trivial.
  apply nth_indep; trivial.
Qed.

Lemma mkKey_left {d d'}: forall l (L: false = (Zlength l >? 64))
        i (I: 0<= i < Zlength l),
      nth (Z.to_nat i) (HMAC_SHA256.mkKey l) d = nth (Z.to_nat i) l d'.
Proof.
  unfold HMAC_SHA256.mkKey. intros. simpl. rewrite <- L.
  apply nth_zeropad_left; trivial.
Qed.

Lemma nth_zeropad_right {d} l i (I: Zlength l <= i < 64):
      nth (Z.to_nat i) (HMAC_SHA256.zeroPad l) d = Z0.
Proof. unfold HMAC_SHA256.zeroPad.
  specialize (Zlength_nonneg l). intros.
  rewrite app_nth2. rewrite nth_list_repeat'. trivial.
  apply minus_lt_compat_r. destruct I. apply Z2Nat.inj_lt in H1. apply H1.
  omega.
  omega.
  destruct I. apply Z2Nat.inj_le in H0; trivial.
    rewrite Zlength_correct, Nat2Z.id in H0. apply H0.
    omega.
  destruct I. apply Z2Nat.inj_le in H0; trivial.
    rewrite Zlength_correct, Nat2Z.id in H0. apply H0.
    omega.
Qed.

Lemma mkKey_right {d}: forall l (L: false = (Zlength l >? 64))
        i (I: Zlength l <= i < 64),
      nth (Z.to_nat i) (HMAC_SHA256.mkKey l) d = Z0.
Proof.
  unfold HMAC_SHA256.mkKey. intros. simpl. rewrite <- L.
  apply nth_zeropad_right; trivial.
Qed.

Lemma zeroPad_BlockSize: forall k, (length k <= SHA256.BlockSize)%nat ->
  length (HMAC_SHA256.zeroPad k) = SHA256.BlockSize%nat.
Proof. unfold HMAC_SHA256.zeroPad. intros. rewrite app_length, (*length_Nlist*) length_list_repeat. omega.
Qed.

Lemma length_SHA256': forall l,
  length (functional_prog.SHA_256' l) = SHA256.DigestLength.
Proof. intros. unfold functional_prog.SHA_256'. simpl.
  rewrite length_intlist_to_Zlist, functional_prog.length_process_msg. reflexivity.
Qed.

Lemma mkKey_length l: length (HMAC_SHA256.mkKey l) = SHA256.BlockSize.
Proof. intros. unfold HMAC_SHA256.mkKey.
  remember (Zlength l >? Z.of_nat SHA256.BlockSize) as z.
  destruct z. apply zeroPad_BlockSize.
    rewrite length_SHA256'.
    apply Nat2Z.inj_le. simpl. omega.
  apply zeroPad_BlockSize.
    rewrite Zlength_correct in Heqz.
    apply Nat2Z.inj_le.
    specialize (Zgt_cases (Z.of_nat (Datatypes.length l)) (Z.of_nat SHA256.BlockSize)). rewrite <- Heqz. trivial.
Qed.

Lemma mkKey_atBlockSize s: length s = SHA256.BlockSize%nat ->
      HMAC_SHA256.mkKey s = s.
Proof. unfold HMAC_SHA256.mkKey. rewrite Zlength_correct.
  intros. rewrite H.
  simpl. unfold HMAC_SHA256.zeroPad. rewrite H. simpl.
  apply app_nil_r.
Qed.

Lemma isbyteZ_mkKey: forall l, Forall isbyteZ l -> Forall isbyteZ (HMAC_SHA256.mkKey l).
Proof. intros.
  unfold HMAC_SHA256.mkKey.
  remember (Zlength l >? Z.of_nat SHA256.BlockSize).
  destruct b.
    apply zeropad_isbyteZ. apply SHA256.Hash_isbyteZ.
    apply zeropad_isbyteZ; trivial.
Qed.

Lemma isbyte_hmac ipad opad m k:
   Forall isbyteZ (HMAC_SHA256.HMAC ipad opad m k).
Proof. apply SHA256.Hash_isbyteZ.  Qed.
(*
Lemma updAbs_len: forall L s t, update_abs L s t -> s256a_len t = s256a_len s + Zlength L * 8.
Proof. intros. inversion H; clear H.
  unfold s256a_len, bitlength. simpl.
  rewrite Zlength_app. subst.
  repeat rewrite Z.mul_add_distr_r.
  repeat rewrite <- Z.add_assoc.
  assert (Zlength blocks * WORD * 8 + Zlength newfrag * 8 =
          Zlength oldfrag * 8 + Zlength L * 8).
    assert (Zlength (oldfrag ++ L) = Zlength (intlist_to_Zlist blocks ++ newfrag)).
      rewrite H4; trivial.
    clear H4. rewrite Zlength_app in H.
              rewrite <- (Z.mul_add_distr_r (Zlength oldfrag)), H. clear H.
              rewrite Zlength_app, Zlength_intlist_to_Zlist.
(*              rewrite (Z.mul_comm WORD). *) rewrite Z.mul_add_distr_r. trivial.
  rewrite H; trivial.
Qed.*)

Lemma HMAC_length d k: length (HMAC256 d k) = 32%nat.
Proof.
  unfold HMAC_SHA256.HMAC, HMAC_SHA256.OUTER.
  apply length_SHA256'.
Qed.
Lemma HMAC_Zlength d k: Zlength (HMAC256 d k) = 32.
Proof.
  rewrite Zlength_correct, HMAC_length. reflexivity.
Qed.


From prelude Require Export countable co_pset.
From algebra Require Export base.

Definition namespace := list positive.
Definition nroot : namespace := nil.
Definition ndot `{Countable A} (N : namespace) (x : A) : namespace :=
  encode x :: N.
Coercion nclose (N : namespace) : coPset := coPset_suffixes (encode N).

Infix ".@" := ndot (at level 19, left associativity) : C_scope.
Notation "(.@)" := ndot (only parsing) : C_scope.

Instance ndot_inj `{Countable A} : Inj2 (=) (=) (=) (@ndot A _ _).
Proof. by intros N1 x1 N2 x2 ?; simplify_eq. Qed.
Lemma nclose_nroot : nclose nroot = ⊤.
Proof. by apply (sig_eq_pi _). Qed.
Lemma encode_nclose N : encode N ∈ nclose N.
Proof. by apply elem_coPset_suffixes; exists xH; rewrite (left_id_L _ _). Qed.
Lemma nclose_subseteq `{Countable A} N x : nclose (N .@ x) ⊆ nclose N.
Proof.
  intros p; rewrite /nclose !elem_coPset_suffixes; intros [q ->].
  destruct (list_encode_suffix N (N .@ x)) as [q' ?]; [by exists [encode x]|].
  by exists (q ++ q')%positive; rewrite <-(assoc_L _); f_equal.
Qed.
Lemma ndot_nclose `{Countable A} N x : encode (N .@ x) ∈ nclose N.
Proof. apply nclose_subseteq with x, encode_nclose. Qed.

Instance ndisjoint : Disjoint namespace := λ N1 N2,
  ∃ N1' N2', N1' `suffix_of` N1 ∧ N2' `suffix_of` N2 ∧
             length N1' = length N2' ∧ N1' ≠ N2'.
Typeclasses Opaque ndisjoint.

Section ndisjoint.
  Context `{Countable A}.
  Implicit Types x y : A.

  Global Instance ndisjoint_comm : Comm iff ndisjoint.
  Proof. intros N1 N2. rewrite /disjoint /ndisjoint; naive_solver. Qed.

  Lemma ndot_ne_disjoint N (x y : A) : x ≠ y → N .@ x ⊥ N .@ y.
  Proof. intros Hxy. exists (N .@ x), (N .@ y); naive_solver. Qed.

  Lemma ndot_preserve_disjoint_l N1 N2 x : N1 ⊥ N2 → N1 .@ x ⊥ N2.
  Proof.
    intros (N1' & N2' & Hpr1 & Hpr2 & Hl & Hne). exists N1', N2'.
    split_and?; try done; []. by apply suffix_of_cons_r.
  Qed.

  Lemma ndot_preserve_disjoint_r N1 N2 x : N1 ⊥ N2 → N1 ⊥ N2 .@ x .
  Proof. rewrite ![N1 ⊥ _]comm. apply ndot_preserve_disjoint_l. Qed.

  Lemma ndisj_disjoint N1 N2 : N1 ⊥ N2 → nclose N1 ∩ nclose N2 = ∅.
  Proof.
    intros (N1' & N2' & [N1'' ->] & [N2'' ->] & Hl & Hne).
    apply elem_of_equiv_empty_L=> p; unfold nclose.
    rewrite elem_of_intersection !elem_coPset_suffixes; intros [[q ->] [q' Hq]].
    rewrite !list_encode_app !assoc in Hq.
    by eapply Hne, list_encode_suffix_eq.
  Qed.
End ndisjoint.

(* This tactic solves goals about inclusion and disjointness
   of masks (i.e., coPsets) with set_solver, taking
   disjointness of namespaces into account. *)
(* TODO: This tactic is by far now yet as powerful as it should be.
   For example, given N1 ⊥ N2, it should be able to solve
   nclose (ndot N1 x) ∩ N2 ≡ ∅. It should also solve
   (ndot N x) ∩ (ndot N y) ≡ ∅ if x ≠ y is in the context or
   follows from [discriminate]. *)
Ltac set_solver_ndisj :=
  repeat match goal with
         (* TODO: Restrict these to have type namespace *)
         | [ H : (?N1 ⊥ ?N2) |-_ ] => apply ndisj_disjoint in H
         end;
  set_solver.
(* TODO: restrict this to match only if this is ⊆ of coPset *)
Hint Extern 500 (_ ⊆ _) => set_solver_ndisj : ndisj.
(* The hope is that registering these will suffice to solve most goals
   of the form N1 ⊥ N2.
   TODO: Can this prove x ≠ y if discriminate can? *)
Hint Resolve ndot_ne_disjoint : ndisj.
Hint Resolve ndot_preserve_disjoint_l : ndisj.
Hint Resolve ndot_preserve_disjoint_r : ndisj.