Require Import Verdi.Verdi.
Require Import Verdi.HandlerMonad.
Require Import Verdi.NameOverlay.

Require Import NameAdjacency.
Require Import AggregationDefinitions.
Require Import TreeAux.

Require Import Sumbool.

Require Import mathcomp.ssreflect.ssreflect.
Require Import mathcomp.ssreflect.ssrbool.
Require Import mathcomp.ssreflect.eqtype.
Require Import mathcomp.ssreflect.fintype.
Require Import mathcomp.ssreflect.finset.
Require Import mathcomp.fingroup.fingroup.

Require Import commfingroup.

Require String.

Local Arguments update {_} {_} _ _ _ _ _ : simpl never.

Set Implicit Arguments.

Module TreeAggregation (Import NT : NameType)
 (NOT : NameOrderedType NT) (NSet : MSetInterface.S with Module E := NOT)
 (NOTC : NameOrderedTypeCompat NT) (NMap : FMapInterface.S with Module E := NOTC)
 (Import RNT : RootNameType NT) (Import CFG : CommutativeFinGroup)
 (Import ANT : AdjacentNameType NT) (Import A : Adjacency NT NOT NSet ANT)
 (Import TA : TAux NT NOT NSet NOTC NMap)
 (Import AD : ADefs NT NOT NSet NOTC NMap CFG).

Import GroupScope.

Inductive Msg : Type :=
| Aggregate : m -> Msg
| Fail : Msg
| Level : option lv -> Msg.

Definition Msg_eq_dec : forall x y : Msg, {x = y} + {x <> y}.
decide equality; first exact: m_eq_dec.
case: o; case: o0.
- move => n m.
  case (lv_eq_dec n m) => H_dec; first by rewrite H_dec; left.
  right.
  move => H_eq.
  injection H_eq => H_eq'.
  by rewrite H_eq' in H_dec.
- by right.
- by right.
- by left.
Defined.

Inductive Input : Type :=
| Local : m -> Input
| SendAggregate : Input
| AggregateRequest : String.string -> Input
| LevelRequest : String.string -> Input
| Broadcast : Input.

Definition Input_eq_dec : forall x y : Input, {x = y} + {x <> y}.
decide equality; auto using String.string_dec, m_eq_dec.
Defined.

Inductive Output : Type :=
| AggregateResponse : String.string -> m -> Output
| LevelResponse : String.string -> option lv -> Output.

Definition Output_eq_dec : forall x y : Output, {x = y} + {x <> y}.
decide equality; auto using String.string_dec; first exact: m_eq_dec.
case: o; case: o0.
- move => lv0 lv1.
  case (lv_eq_dec lv0 lv1) => H_dec; first by rewrite H_dec; left.
  right.
  move => H_eq.
  injection H_eq => H_eq'.
  by rewrite H_eq' in H_dec.
- by right.
- by right.
- by left.
Defined.

Record Data :=  mkData {
  local : m ;
  aggregate : m ;
  adjacent : NS ;
  balance : NM ;
  broadcast : bool ;
  levels : NL
}.

Definition InitData (n : name) :=
if root_dec n then
  {| local := 1 ;
     aggregate := 1 ;
     adjacent := adjacency n nodes ;
     balance := init_map (adjacency n nodes) ;
     broadcast := true ;
     levels := NMap.empty lv |}
else
  {| local := 1 ;
     aggregate := 1 ;
     adjacent := adjacency n nodes ;
     balance := init_map (adjacency n nodes) ;
     broadcast := false ;
     levels := NMap.empty lv |}.

Definition Handler (S : Type) := GenHandler (name * Msg) S Output unit.

Definition RootNetHandler (src : name) (msg : Msg) : Handler Data :=
st <- get ;;
match msg with
| Aggregate m_msg =>
  match NMap.find src st.(balance) with
  | None => nop
  | Some m_src =>
    put {| local := st.(local) ;
           aggregate := st.(aggregate) * m_msg ;
           adjacent := st.(adjacent) ;
           balance := NMap.add src (m_src * (m_msg)^-1) st.(balance) ;
           broadcast := st.(broadcast) ;
           levels := st.(levels) |}
  end
| Level _ => nop
| Fail =>
  match NMap.find src st.(balance) with
  | Some m_bal =>
    put {| local := st.(local) ;
           aggregate := st.(aggregate) * m_bal ;
           adjacent := NSet.remove src st.(adjacent) ;
           balance := NMap.remove src st.(balance) ;
           broadcast := st.(broadcast) ;
           levels := st.(levels) |}
  | None =>
    put {| local := st.(local) ;
           aggregate := st.(aggregate) ;
           adjacent := NSet.remove src st.(adjacent) ;
           balance := st.(balance) ;
           broadcast := st.(broadcast) ;
           levels := st.(levels) |}
  end
end.

Definition NonRootNetHandler (me src: name) (msg : Msg) : Handler Data :=
st <- get ;;
match msg with
| Aggregate m_msg =>
  match NMap.find src st.(balance) with
  | None => nop
  | Some m_src =>
    put {| local := st.(local) ;
           aggregate := st.(aggregate) * m_msg ;
           adjacent := st.(adjacent) ;
           balance := NMap.add src (m_src * (m_msg)^-1) st.(balance) ;
           broadcast := st.(broadcast) ;
           levels := st.(levels) |}
  end
| Level None =>
  if olv_eq_dec (level st.(adjacent) st.(levels)) (level st.(adjacent) (NMap.remove src st.(levels))) then
    put {| local := st.(local) ;
           aggregate := st.(aggregate) ;
           adjacent := st.(adjacent) ;
           balance := st.(balance) ;
           broadcast := st.(broadcast) ;
           levels := NMap.remove src st.(levels) |}
  else
    put {| local := st.(local) ;
           aggregate := st.(aggregate) ;
           adjacent := st.(adjacent) ;
           balance := st.(balance) ;
           broadcast := true ;
           levels := NMap.remove src st.(levels) |}
| Level (Some lv') =>
  if olv_eq_dec (level st.(adjacent) st.(levels)) (level st.(adjacent) (NMap.add src lv' st.(levels))) then
    put {| local := st.(local) ;
           aggregate := st.(aggregate) ;
           adjacent := st.(adjacent) ;
           balance := st.(balance) ;
           broadcast := st.(broadcast) ;
           levels := NMap.add src lv' st.(levels) |}
  else
    put {| local := st.(local) ;
           aggregate := st.(aggregate) ;
           adjacent := st.(adjacent) ;
           balance := st.(balance) ;
           broadcast := true ;
           levels := NMap.add src lv' st.(levels) |}
| Fail =>
  match NMap.find src st.(balance) with
  | Some m_bal =>
    if olv_eq_dec (level st.(adjacent) st.(levels)) (level (NSet.remove src st.(adjacent)) (NMap.remove src st.(levels))) then
      put {| local := st.(local) ;
             aggregate := st.(aggregate) * m_bal ;
             adjacent := NSet.remove src st.(adjacent) ;
             balance := NMap.remove src st.(balance) ;
             broadcast := st.(broadcast) ;
             levels := NMap.remove src st.(levels) |}
    else
      put {| local := st.(local) ;
             aggregate := st.(aggregate) * m_bal ;
             adjacent := NSet.remove src st.(adjacent) ;
             balance := NMap.remove src st.(balance) ;
             broadcast := true ;
             levels := NMap.remove src st.(levels) |}
  | None =>
    if olv_eq_dec (level st.(adjacent) st.(levels)) (level (NSet.remove src st.(adjacent)) (NMap.remove src st.(levels))) then
      put {| local := st.(local) ;
             aggregate := st.(aggregate) ;
             adjacent := NSet.remove src st.(adjacent) ;
             balance := st.(balance) ;
             broadcast := st.(broadcast) ;
             levels := NMap.remove src st.(levels) |}
    else
      put {| local := st.(local) ;
             aggregate := st.(aggregate) ;
             adjacent := NSet.remove src st.(adjacent) ;
             balance := st.(balance) ;
             broadcast := true ;
             levels := NMap.remove src st.(levels) |}
  end
end.

Definition NetHandler (me src : name) (msg : Msg) : Handler Data :=
if root_dec me then RootNetHandler src msg
else NonRootNetHandler me src msg.

Definition send_level_fold (lvo : option lv) (n : name) (res : Handler Data) : Handler Data :=
send (n, Level lvo) ;; res.

Definition send_level_adjacent (lvo : option lv) (fs : NS) : Handler Data :=
NSet.fold (send_level_fold lvo) fs nop.

Definition RootIOHandler (i : Input) : Handler Data :=
st <- get ;;
match i with
| Local m_msg =>
  put {| local := m_msg;
         aggregate := st.(aggregate) * m_msg * st.(local)^-1;
         adjacent := st.(adjacent);
         balance := st.(balance);
         broadcast := st.(broadcast);
         levels := st.(levels) |}
| SendAggregate => nop
| AggregateRequest client_id =>
  write_output (AggregateResponse client_id st.(aggregate))
| Broadcast =>
  when st.(broadcast)
  (send_level_adjacent (Some 0) st.(adjacent) ;;
   put {| local := st.(local);
          aggregate := st.(aggregate);
          adjacent := st.(adjacent);
          balance := st.(balance);
          broadcast := false;
          levels := st.(levels) |})
| LevelRequest client_id =>
  write_output (LevelResponse client_id (Some 0))
end.

Definition NonRootIOHandler (i : Input) : Handler Data :=
st <- get ;;
match i with
| Local m_msg =>
  put {| local := m_msg;
         aggregate := st.(aggregate) * m_msg * st.(local)^-1;
         adjacent := st.(adjacent);
         balance := st.(balance);
         broadcast := st.(broadcast);
         levels := st.(levels) |}
| SendAggregate =>
  when (sumbool_not _ _ (m_eq_dec st.(aggregate) 1))
  (match parent st.(adjacent) st.(levels) with
  | None => nop
  | Some dst =>
    match NMap.find dst st.(balance) with
    | None => nop
    | Some m_dst =>
      send (dst, (Aggregate st.(aggregate))) ;;
      put {| local := st.(local);
             aggregate := 1;
             adjacent := st.(adjacent);
             balance := NMap.add dst (m_dst * st.(aggregate)) st.(balance);
             broadcast := st.(broadcast);
             levels := st.(levels) |}
    end
  end)
| AggregateRequest client_id =>
  write_output (AggregateResponse client_id st.(aggregate))
| Broadcast =>
  when st.(broadcast)
  (send_level_adjacent (level st.(adjacent) st.(levels)) st.(adjacent) ;;
  put {| local := st.(local);
         aggregate := st.(aggregate);
         adjacent := st.(adjacent);
         balance := st.(balance);
         broadcast := false;
         levels := st.(levels) |})
| LevelRequest client_id =>
  write_output (LevelResponse client_id (level st.(adjacent) st.(levels)))
end.

Definition IOHandler (me : name) (i : Input) : Handler Data :=
if root_dec me then RootIOHandler i
else NonRootIOHandler i.

Instance TreeAggregation_BaseParams : BaseParams :=
  {
    data := Data;
    input := Input;
    output := Output
  }.

Instance TreeAggregation_MultiParams : MultiParams TreeAggregation_BaseParams :=
  {
    name := name ;
    msg  := Msg ;
    msg_eq_dec := Msg_eq_dec ;
    name_eq_dec := name_eq_dec ;
    nodes := nodes ;
    all_names_nodes := all_names_nodes ;
    no_dup_nodes := no_dup_nodes ;
    init_handlers := InitData ;
    net_handlers := fun dst src msg s =>
                      runGenHandler_ignore s (NetHandler dst src msg) ;
    input_handlers := fun nm msg s =>
                        runGenHandler_ignore s (IOHandler nm msg)
  }.

Instance TreeAggregation_NameOverlayParams : NameOverlayParams TreeAggregation_MultiParams :=
  {
    adjacent_to := adjacent_to ;
    adjacent_to_dec := adjacent_to_dec ;
    adjacent_to_symmetric := adjacent_to_symmetric ;
    adjacent_to_irreflexive := adjacent_to_irreflexive
  }.

Instance TreeAggregation_FailMsgParams : FailMsgParams TreeAggregation_MultiParams :=
  {
    msg_fail := Fail
  }.

Lemma net_handlers_NetHandler :
  forall dst src m st os st' ms,
    net_handlers dst src m st = (os, st', ms) ->
    NetHandler dst src m st = (tt, os, st', ms).
Proof.
intros.
simpl in *.
monad_unfold.
repeat break_let.
find_inversion.
destruct u. auto.
Qed.

Lemma NetHandler_cases :
  forall dst src msg st out st' ms,
    NetHandler dst src msg st = (tt, out, st', ms) ->
    (exists m_msg m_src, msg = Aggregate m_msg /\ NMap.find src st.(balance) = Some m_src /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) * m_msg /\
     st'.(adjacent) = st.(adjacent) /\
     st'.(balance) = NMap.add src (m_src * (m_msg)^-1) st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = st.(levels) /\
     out = [] /\ ms = []) \/
    (exists m_msg, msg = Aggregate m_msg /\ NMap.find src st.(balance) = None /\
     st' = st /\
     out = [] /\ ms = []) \/
    (root dst /\ msg = Fail /\ exists m_bal, NMap.find src st.(balance) = Some m_bal /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) * m_bal /\
     st'.(adjacent) = NSet.remove src st.(adjacent) /\
     st'.(balance) = NMap.remove src st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ msg = Fail /\
     exists m_bal, NMap.find src st.(balance) = Some m_bal /\
     level st.(adjacent) st.(levels) = level (NSet.remove src st.(adjacent)) (NMap.remove src st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) * m_bal /\
     st'.(adjacent) = NSet.remove src st.(adjacent) /\
     st'.(balance) = NMap.remove src st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = NMap.remove src st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ msg = Fail /\
     exists m_bal, NMap.find src st.(balance) = Some m_bal /\
     level st.(adjacent) st.(levels) <> level (NSet.remove src st.(adjacent)) (NMap.remove src st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) * m_bal /\
     st'.(adjacent) = NSet.remove src st.(adjacent) /\
     st'.(balance) = NMap.remove src st.(balance) /\
     st'.(broadcast) = true /\
     st'.(levels) = NMap.remove src st.(levels) /\
     out = [] /\ ms = []) \/
    (root dst /\ msg = Fail /\ NMap.find src st.(balance) = None /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = NSet.remove src st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ msg = Fail /\ NMap.find src st.(balance) = None /\
     level st.(adjacent) st.(levels) = level (NSet.remove src st.(adjacent)) (NMap.remove src st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = NSet.remove src st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = NMap.remove src st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ msg = Fail /\ NMap.find src st.(balance) = None /\
     level st.(adjacent) st.(levels) <> level (NSet.remove src st.(adjacent)) (NMap.remove src st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = NSet.remove src st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = true /\
     st'.(levels) = NMap.remove src st.(levels) /\
     out = [] /\ ms = []) \/
    (root dst /\ exists lvo, msg = Level lvo /\
     st' = st /\
     out = [] /\ ms = []) \/
    (~ root dst /\ exists lv_msg, msg = Level (Some lv_msg) /\
     level st.(adjacent) st.(levels) = level st.(adjacent) (NMap.add src lv_msg st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = NMap.add src lv_msg st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ exists lv_msg, msg = Level (Some lv_msg) /\
     level st.(adjacent) st.(levels) <> level st.(adjacent) (NMap.add src lv_msg st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = true /\
     st'.(levels) = NMap.add src lv_msg st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ msg = Level None /\
     level st.(adjacent) st.(levels) = level st.(adjacent) (NMap.remove src st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = st.(broadcast) /\
     st'.(levels) = NMap.remove src st.(levels) /\
     out = [] /\ ms = []) \/
    (~ root dst /\ msg = Level None /\
     level st.(adjacent) st.(levels) <> level st.(adjacent) (NMap.remove src st.(levels)) /\
     st'.(local) = st.(local) /\
     st'.(aggregate) = st.(aggregate) /\
     st'.(adjacent) = st.(adjacent) /\
     st'.(balance) = st.(balance) /\
     st'.(broadcast) = true /\
     st'.(levels) = NMap.remove src st.(levels) /\
     out = [] /\ ms = []).
Proof.
move => dst src msg st out st' ms.
rewrite /NetHandler /RootNetHandler /NonRootNetHandler.
case: msg => [m_msg||olv_msg]; monad_unfold; case root_dec => /= H_dec; repeat break_let; repeat find_injection; try (break_match; repeat find_injection); try (break_if; repeat find_injection).
- by left; exists m_msg, a.
- by right; left; exists m_msg.
- by left; exists m_msg, a.
- by right; left; exists m_msg.
- right; right; left.
  repeat split => //.
  by exists a.
- by right; right; right; right; right; left.
- right; right; right; left.
  repeat split => //.
  exists a.
  move: Heqb.
  by case: olv_eq_dec.
- right; right; right; right; left.
  repeat split => //.
  exists a.
  move: Heqb.
  by case: olv_eq_dec.
- right; right; right; right; right; right; left.
  move: Heqb.
  by case: olv_eq_dec.
- right; right; right; right; right; right; right; left.
  move: Heqb.
  by case: olv_eq_dec.
- move => H_eq.
  find_injection.
  right; right; right; right; right; right; right; right; left.
  split => //.
  by exists olv_msg.
- right; right; right; right; right; right; right; right; right; left.
  split => //.
  exists l1.
  move: Heqb.
  by case: olv_eq_dec.
- right; right; right; right; right; right; right; right; right; right; left.
  split => //.
  exists l1.
  move: Heqb.
  by case: olv_eq_dec.
- right; right; right; right; right; right; right; right; right; right; right; left.
  move: Heqb.
  by case: olv_eq_dec.
- right; right; right; right; right; right; right; right; right; right; right; right.
  move: Heqb.
  by case: olv_eq_dec.
Qed.

Lemma input_handlers_IOHandler :
  forall h i d os d' ms,
    input_handlers h i d = (os, d', ms) ->
    IOHandler h i d = (tt, os, d', ms).
Proof.
intros.
simpl in *.
monad_unfold.
repeat break_let.
find_inversion.
destruct u. auto.
Qed.

Lemma send_level_fold_app :
  forall ns st olv nm,
snd (fold_left
       (fun (a : Handler Data) (e : NSet.elt) => send_level_fold olv e a) ns
       (fun s : Data => (tt, [], s, nm)) st) =
snd (fold_left
       (fun (a : Handler Data) (e : NSet.elt) => send_level_fold olv e a) ns
       (fun s : Data => (tt, [], s, [])) st) ++ nm.
Proof.
elim => //=.
move => n ns IH st olv nm.
rewrite {1}/send_level_fold /=.
monad_unfold.
rewrite /=.
rewrite IH.
rewrite app_assoc.
by rewrite -IH.
Qed.

Instance TreeAggregation_TreeMsg : TreeMsg :=
  {
    tree_msg := Msg ;
    tree_level := Level
  }.

Lemma send_level_adjacent_fst_eq :
forall fs olv st,
  snd (send_level_adjacent olv fs st) = level_adjacent olv fs.
Proof.
move => fs olv st.
rewrite /send_level_adjacent /level_adjacent.
rewrite 2!NSet.fold_spec.
move: olv st.
elim: NSet.elements => [|n ns IH] //=.
move => olv st.
rewrite {2}/level_fold {2}/send_level_fold.
rewrite (@fold_left_level_fold_eq TreeAggregation_TreeMsg).
have IH' := IH olv st.
rewrite -IH'.
monad_unfold.
by rewrite -send_level_fold_app.
Qed.

Lemma fst_fst_fst_tt_send_level_fold :
forall ns nm olv st,
fst
  (fst
     (fst
        (fold_left
           (fun (a : Handler Data) (e : NSet.elt) =>
              send_level_fold olv e a) ns
           (fun s : Data => (tt, [], s, nm)) st))) = tt.
Proof.
elim => //=.
move => n ns IH nm olv st.
by rewrite IH.
Qed.

Lemma send_level_adjacent_fst_fst_eq :
forall fs olv st,
  fst (fst (fst (send_level_adjacent olv fs st))) = tt.
Proof.
move => fs olv st.
rewrite /send_level_adjacent.
rewrite NSet.fold_spec.
by rewrite fst_fst_fst_tt_send_level_fold.
Qed.

Lemma snd_fst_fst_out_send_level_fold :
forall ns nm olv st,
snd
  (fst
     (fst
        (fold_left
           (fun (a : Handler Data) (e : NSet.elt) =>
              send_level_fold olv e a) ns
           (fun s : Data => (tt, [], s, nm)) st))) = [].
Proof.
elim => //=.
move => n ns IH nm olv st.
by rewrite IH.
Qed.

Lemma snd_fst_st_send_level_fold :
forall ns nm olv st,
snd (fst
        (fold_left
           (fun (a : Handler Data) (e : NSet.elt) =>
              send_level_fold olv e a) ns
           (fun s : Data => (tt, [], s, nm)) st)) = st.
Proof.
elim => //=.
move => n ns IH nm olv st.
by rewrite IH.
Qed.

Lemma send_level_adjacent_snd_fst_fst :
forall fs olv st,
  snd (fst (fst (send_level_adjacent olv fs st))) = [].
Proof.
move => fs olv st.
rewrite /send_level_adjacent.
rewrite NSet.fold_spec.
by rewrite snd_fst_fst_out_send_level_fold.
Qed.

Lemma send_level_adjacent_snd_fst :
forall fs olv st,
  snd (fst (send_level_adjacent olv fs st)) = st.
Proof.
move => fs olv st.
rewrite /send_level_adjacent.
rewrite NSet.fold_spec.
by rewrite snd_fst_st_send_level_fold.
Qed.

Lemma send_level_adjacent_eq :
  forall fs olv st,
  send_level_adjacent olv fs st = (tt, [], st, level_adjacent olv fs).
Proof.
move => fs olv st.
case H_eq: send_level_adjacent => [[[u o] s b]].
have H_eq'_1 := send_level_adjacent_fst_fst_eq fs olv st.
rewrite H_eq /= in H_eq'_1.
have H_eq'_2 := send_level_adjacent_fst_eq fs olv st.
rewrite H_eq /= in H_eq'_2.
have H_eq'_3 := send_level_adjacent_snd_fst_fst fs olv st.
rewrite H_eq /= in H_eq'_3.
have H_eq'_4 := send_level_adjacent_snd_fst fs olv st.
rewrite H_eq /= in H_eq'_4.
by rewrite H_eq'_1 H_eq'_2 H_eq'_3 H_eq'_4.
Qed.

Lemma IOHandler_cases :
  forall h i st u out st' ms,
      IOHandler h i st = (u, out, st', ms) ->
      (exists m_msg, i = Local m_msg /\
         st'.(local) = m_msg /\
         st'.(aggregate) = st.(aggregate) * m_msg * st.(local)^-1 /\
         st'.(adjacent) = st.(adjacent) /\
         st'.(balance) = st.(balance) /\
         st'.(broadcast) = st.(broadcast) /\
         st'.(levels) = st.(levels) /\
         out = [] /\ ms = []) \/
      (root h /\ i = SendAggregate /\
         st' = st /\
         out = [] /\ ms = []) \/
      (~ root h /\ i = SendAggregate /\
       st.(aggregate) <> 1 /\
       exists dst m_dst, parent st.(adjacent) st.(levels) = Some dst /\ NMap.find dst st.(balance) = Some m_dst /\
       st'.(local) = st.(local) /\
       st'.(aggregate) = 1 /\
       st'.(adjacent) = st.(adjacent) /\
       st'.(balance) = NMap.add dst (m_dst * st.(aggregate)) st.(balance) /\
       st'.(broadcast) = st.(broadcast) /\
       st'.(levels) = st.(levels) /\
       out = [] /\ ms = [(dst, Aggregate st.(aggregate))]) \/
      (~ root h /\ i = SendAggregate /\
       st.(aggregate) = 1 /\
       st' = st /\
       out = [] /\ ms = []) \/
      (~ root h /\ i = SendAggregate /\
       st.(aggregate) <> 1 /\
       parent st.(adjacent) st.(levels) = None /\
       st' = st /\
       out = [] /\ ms = []) \/
      (~ root h /\ i = SendAggregate /\
       st.(aggregate) <> 1 /\
       exists dst, parent st.(adjacent) st.(levels) = Some dst /\ NMap.find dst st.(balance) = None /\
       st' = st /\
       out = [] /\ ms = []) \/
      (exists client_id, i = AggregateRequest client_id /\
       st' = st /\
       out = [AggregateResponse client_id (aggregate st)] /\ ms = []) \/
      (root h /\ i = Broadcast /\ st.(broadcast) = true /\
       st'.(local) = st.(local) /\
       st'.(aggregate) = st.(aggregate) /\
       st'.(adjacent) = st.(adjacent) /\
       st'.(balance) = st.(balance) /\
       st'.(broadcast) = false /\
       st'.(levels) = st.(levels) /\
       out = [] /\ ms = level_adjacent (Some 0) st.(adjacent)) \/
      (~ root h /\ i = Broadcast /\ st.(broadcast) = true /\
       st'.(local) = st.(local) /\
       st'.(aggregate) = st.(aggregate) /\
       st'.(adjacent) = st.(adjacent) /\
       st'.(balance) = st.(balance) /\
       st'.(broadcast) = false /\
       st'.(levels) = st.(levels) /\
       out = [] /\ ms = level_adjacent (level st.(adjacent) st.(levels)) st.(adjacent)) \/
      (i = Broadcast /\ st.(broadcast) = false /\
       st' = st /\
       out = [] /\ ms = []) \/
      (root h /\ exists client_id, i = LevelRequest client_id /\
       st' = st /\
       out = [LevelResponse client_id (Some 0)] /\ ms = []) \/
      (~ root h /\ exists client_id, i = LevelRequest client_id /\
       st' = st /\
       out = [LevelResponse client_id (level st.(adjacent) st.(levels))] /\ ms = []).
Proof.
move => h i st u out st' ms.
rewrite /IOHandler /RootIOHandler /NonRootIOHandler.
case: i => [m_msg||client_id|client_id|]; monad_unfold; case root_dec => /= H_dec H_eq; repeat break_let; repeat find_injection; repeat break_match; repeat break_let; repeat find_injection.
- by left; exists m_msg.
- by left; exists m_msg.
- by right; left.
- unfold sumbool_not in *.
  break_match => //.
  right; right; left.
  repeat split => //.
  by exists n, a.
- unfold sumbool_not in *.
  break_match => //.
  right; right; right; right; right; left.
  repeat split => //.
  by exists n.
- unfold sumbool_not in *.
  break_match => //.
  by right; right; right; right; left.
- unfold sumbool_not in *.
  break_match => //.
  by right; right; right; left.
- right; right; right; right; right; right; left.
  by exists client_id.
- right; right; right; right; right; right; left.
  by exists client_id.
- right; right; right; right; right; right; right; right; right; right; left.
  split => //.
  by exists client_id.
- right; right; right; right; right; right; right; right; right; right; right.
  split => //.
  by exists client_id.
- find_rewrite_lem send_level_adjacent_eq.
  find_injection.
  right; right; right; right; right; right; right; left.
  by rewrite app_nil_l -2!app_nil_end.
- by right; right; right; right; right; right; right; right; right; left.
- find_rewrite_lem send_level_adjacent_eq.
  find_injection.
  right; right; right; right; right; right; right; right; left.
  by rewrite app_nil_l -2!app_nil_end.
- by right; right; right; right; right; right; right; right; right; left.
Qed.

Ltac net_handler_cases :=
  find_apply_lem_hyp NetHandler_cases;
  intuition idtac; try break_exists;
  intuition idtac; subst;
  repeat find_rewrite.

Ltac io_handler_cases :=
  find_apply_lem_hyp IOHandler_cases;
  intuition idtac; try break_exists;
  intuition idtac; subst;
  repeat find_rewrite.

End TreeAggregation.


From Cyclic_PA.Maths Require Import naturals.
From Cyclic_PA.Maths Require Import ordinals.

From Cyclic_PA.Logic Require Import definitions.
From Cyclic_PA.Logic Require Import fol.
From Cyclic_PA.Logic Require Import PA_omega.
From Cyclic_PA.Logic Require Import proof_trees.
From Cyclic_PA.Logic Require Import substitute.

Fixpoint formula_sub_ptree_fit
  (P : ptree) (E F : formula) (S : subst_ind) : ptree :=
match P, S with
| deg_up d P', _ => deg_up d (formula_sub_ptree_fit P' E F S)

| ord_up alpha P', _ => ord_up alpha (formula_sub_ptree_fit P' E F S)

| node A, _ => node (formula_sub_ind_fit A E F S)

| exchange_ab A B d alpha P', lor_ind S_B S_A =>
    exchange_ab
      (formula_sub_ind_fit A E F S_A)
      (formula_sub_ind_fit B E F S_B)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind S_A S_B))

| exchange_cab C A B d alpha P', lor_ind (lor_ind S_C S_B) S_A =>
    exchange_cab
      (formula_sub_ind_fit C E F S_C)
      (formula_sub_ind_fit A E F S_A)
      (formula_sub_ind_fit B E F S_B)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind (lor_ind S_C S_A) S_B))

| exchange_abd A B D d alpha P', lor_ind (lor_ind S_B S_A) S_D =>
    exchange_abd
      (formula_sub_ind_fit A E F S_A)
      (formula_sub_ind_fit B E F S_B)
      (formula_sub_ind_fit D E F S_D)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind (lor_ind S_A S_B) S_D))

| exchange_cabd C A B D d alpha P', lor_ind (lor_ind (lor_ind S_C S_B) S_A) S_D =>
    exchange_cabd
      (formula_sub_ind_fit C E F S_C)
      (formula_sub_ind_fit A E F S_A)
      (formula_sub_ind_fit B E F S_B)
      (formula_sub_ind_fit D E F S_D)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind (lor_ind (lor_ind S_C S_A) S_B) S_D))

| contraction_a A d alpha P', _ =>
    contraction_a
      (formula_sub_ind_fit A E F S)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind S S))

| contraction_ad A D d alpha P', lor_ind S_A S_D =>
    contraction_ad
      (formula_sub_ind_fit A E F S_A)
      (formula_sub_ind_fit D E F S_D)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind (lor_ind S_A S_A) S_D))

| weakening_ad A D d alpha P', lor_ind S_A S_D =>
    weakening_ad
      (formula_sub_ind_fit A E F S_A)
      (formula_sub_ind_fit D E F S_D)
      d alpha
      (formula_sub_ptree_fit P' E F S_D)

| demorgan_ab A B d1 d2 alpha1 alpha2 P1 P2, _ => P

| demorgan_abd A B D d1 d2 alpha1 alpha2 P1 P2, lor_ind S_AB S_D =>
    demorgan_abd
      A B
      (formula_sub_ind_fit D E F S_D)
      d1 d2 alpha1 alpha2
      (formula_sub_ptree_fit P1 E F (lor_ind (0) S_D))
      (formula_sub_ptree_fit P2 E F (lor_ind (0) S_D))

| negation_a A d alpha P', _ => P

| negation_ad A D d alpha P', lor_ind S_A S_D =>
    negation_ad
      A
      (formula_sub_ind_fit D E F S_D)
      d alpha
      (formula_sub_ptree_fit P' E F (lor_ind (non_target A) S_D))

| quantification_a A n t d alpha P', _ => P

| quantification_ad A D n t d alpha P', lor_ind S_A S_D =>
    quantification_ad
      A
      (formula_sub_ind_fit D E F S_D)
      n t d alpha
      (formula_sub_ptree_fit P' E F (lor_ind (0) S_D))

| w_rule_a A n d alpha g, _ => P

| w_rule_ad A D n d alpha g, lor_ind S_A S_D =>
    w_rule_ad
      A
      (formula_sub_ind_fit D E F S_D)
      n d alpha
      (fun (t : c_term) =>
          formula_sub_ptree_fit (g t) E F (lor_ind (non_target A) S_D))

| cut_ca C A d1 d2 alpha1 alpha2 P1 P2, _ =>
    cut_ca
      (formula_sub_ind_fit C E F S)
      A d1 d2 alpha1 alpha2
      (formula_sub_ptree_fit P1 E F (lor_ind S (non_target A)))
      P2

| cut_ad A D d1 d2 alpha1 alpha2 P1 P2, _ =>
    cut_ad
      A
      (formula_sub_ind_fit D E F S)
      d1 d2 alpha1 alpha2
      P1
      (formula_sub_ptree_fit P2 E F (lor_ind (0) S))

| cut_cad C A D d1 d2 alpha1 alpha2 P1 P2, lor_ind S_C S_D =>
    cut_cad
      (formula_sub_ind_fit C E F S_C)
      A
      (formula_sub_ind_fit D E F S_D)
      d1 d2 alpha1 alpha2
      (formula_sub_ptree_fit P1 E F (lor_ind S_C (non_target A)))
      (formula_sub_ptree_fit P2 E F (lor_ind (0) S_D))

| _, _ => P
end.

Definition formula_sub_ptree
  (P : ptree) (E F : formula) (S : subst_ind) : ptree :=
match subst_ind_fit (ptree_formula P) S with
| false => P
| true => formula_sub_ptree_fit P E F S
end.

Lemma formula_sub_ptree_formula_true :
    forall (P : ptree) (E F : formula) (S : subst_ind),
        subst_ind_fit (ptree_formula P) S = true ->
            formula_sub_ptree_fit P E F S = formula_sub_ptree P E F S.
Proof.
intros P E F S FS.
unfold formula_sub_ptree.
rewrite FS.
reflexivity.
Qed.

Lemma formula_sub_ptree_fit_false :
    forall (P : ptree) (E F : formula) (S : subst_ind),
        subst_ind_fit (ptree_formula P) S = false ->
            formula_sub_ptree P E F S = P.
Proof.
intros P E F S FS.
unfold formula_sub_ptree.
rewrite FS.
reflexivity.
Qed.

Lemma formula_sub_ind_fit_closed :
    forall (A B C : formula),
        closed A = true ->
            (closed B = true -> closed C = true) ->
                forall (S : subst_ind),
                    subst_ind_fit A S = true ->
                        closed (formula_sub_ind_fit A B C S) = true.
Proof.
intros A B C CA CBC S FS.
destruct (closed (formula_sub_ind A B C S)) eqn:CFC.
- rewrite (sub_fit_true _ _ _ _ FS) in CFC.
  apply CFC.
- rewrite (formula_sub_ind_closed _ _ _ CA CBC) in CFC.
  inversion CFC.
Qed.

Lemma formula_sub_ptree_formula_atom' :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                subst_ind_fit (ptree_formula P) S = true ->
                    ptree_formula (formula_sub_ptree P (atom a) F S) =
                        formula_sub_ind (ptree_formula P) (atom a) F S.
Proof.
intros P a F PV.
induction P; intros S FS;
unfold formula_sub_ptree;
rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit.

1 : destruct PV as [ID PV].

2 : destruct PV as [[IO PV] NO].

1,2 : unfold ptree_formula in *; fold ptree_formula in *;
      rewrite (formula_sub_ptree_formula_true _ _ _ _ FS);
      apply (IHP PV _ FS).

1 : { unfold ptree_formula, formula_sub_ind in *.
      rewrite FS.
      reflexivity. }

1-4,6-15,18 : destruct S; inversion FS as [FS'];
              try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

2-4,9,12,15,18 : destruct S1; inversion FS' as [FS''];
      try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

4 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

all : unfold ptree_formula in *; fold ptree_formula in *;
      unfold formula_sub_ind, subst_ind_fit; fold subst_ind_fit;
      try rewrite FS;
      try rewrite FS1;
      try rewrite FS2;
      try rewrite FS1_1;
      try rewrite FS1_2;
      try rewrite FS1_1_1;
      try rewrite FS1_1_2;
      try reflexivity.
Qed.

Lemma formula_sub_ptree_formula_atom :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                ptree_formula (formula_sub_ptree P (atom a) F S) =
                    formula_sub_ind (ptree_formula P) (atom a) F S.
Proof.
intros P a F PV S.
destruct (subst_ind_fit (ptree_formula P) S) eqn:FS.
- apply (formula_sub_ptree_formula_atom' _ _ _ PV _ FS).
- unfold formula_sub_ptree, formula_sub_ind.
  rewrite FS.
  reflexivity.
Qed.

Lemma formula_sub_ptree_deg_atom :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                ptree_deg (formula_sub_ptree P (atom a) F S) = ptree_deg P.
Proof.
intros P a F PV.
pose (ptree_formula P) as A.
induction P; intros S;
case (subst_ind_fit A S) eqn:FS;
unfold formula_sub_ptree;
unfold A in FS;
try rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit;
try reflexivity.

1 : destruct PV as [[IO PV] NO].
    unfold ptree_formula in *; fold ptree_formula in *.
    rewrite (formula_sub_ptree_formula_true _ _ _ _ FS).
    apply (IHP PV).

all : destruct S; inversion FS as [FS'];
      try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

2-4,7-11 : destruct S1; inversion FS' as [FS''];
      try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

4 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

all : unfold ptree_deg; fold ptree_deg;
      reflexivity.
Qed.

Lemma formula_sub_ptree_ord_atom :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                ptree_ord (formula_sub_ptree P (atom a) F S) = ptree_ord P.
Proof.
intros P a F PV.
pose (ptree_formula P) as A.
induction P; intros S;
case (subst_ind_fit A S) eqn:FS;
unfold formula_sub_ptree;
unfold A in FS;
try rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit;
try reflexivity.

1 : destruct PV as [ID PV].
    unfold ptree_formula in *; fold ptree_formula in *.
    rewrite (formula_sub_ptree_formula_true _ _ _ _ FS).
    apply (IHP PV).

all : destruct S; inversion FS as [FS'];
      try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

2-4,7-11 : destruct S1; inversion FS' as [FS''];
      try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

4 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

all : unfold ptree_deg; fold ptree_deg;
      reflexivity.
Qed.

Lemma formula_sub_valid_atom :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            correct_a a = false ->
                closed F =true ->
                    forall (S : subst_ind),
                        subst_ind_fit (ptree_formula P) S = true ->
                            valid (formula_sub_ptree P (atom a) F S).
Proof.
intros P a F PV Aa CF.
induction P; intros S FS;
unfold formula_sub_ptree;
rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit.

1 : destruct PV as [ID PV].
2 : destruct PV as [[IO PV] NO].
4-9 : destruct PV as [[[PF PV] PD] PO].
10 : destruct PV as [[[[PF FC] PV] PD] PO].
11,12 : destruct PV as [[[[[[[P1F P1V] P2F] P2V] P1D] P2D] P1O] P2O].
13-16 : destruct PV as [[[PF PV] PD] PO].
19-21 : destruct PV as [[[[[[[P1F P1V] P2F] P2V] P1D] P2D] P1O] P2O].

4-7,9,10,12,14,16,18,21 : destruct S; inversion FS as [FS'];
                          try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

5-7,9-13 :  destruct S1; inversion FS' as [FS''];
            try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

7 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

3 : { unfold valid in *.
      destruct f.
      3,4 : inversion PV.
      2 : unfold formula_sub_ind_fit, form_eqb;
          apply PV.
      1 : destruct S;
          unfold formula_sub_ind_fit, form_eqb;
          fold formula_sub_ind_fit;
          case (atom_eqb a0 a) eqn:EQ;
          try apply PV.
          apply atom_beq_eq in EQ.
          destruct EQ.
          unfold PA_omega_axiom in PV.
          rewrite PV in Aa.
          inversion Aa. }

16,17 : repeat split;
        destruct (PV t) as [[[PF PcV] PD] PO];
        rewrite formula_sub_ptree_formula_true;
        try rewrite PF;
        try rewrite formula_sub_ptree_deg_atom;
        try apply PD;
        try rewrite formula_sub_ptree_ord_atom;
        try apply PO;
        try apply X;
        try rewrite formula_sub_ptree_formula_atom;
        try rewrite PF;
        try apply PV;
        try rewrite formula_sub_ind_lor;
        try rewrite (non_target_term_sub _ n (projT1 t));
        try rewrite non_target_sub;
        unfold formula_sub_ind;
        unfold subst_ind_fit; fold subst_ind_fit;
        try rewrite FS;
        try rewrite FS1;
        try rewrite FS2;
        try rewrite non_target_fit;
        reflexivity.

all : repeat rewrite formula_sub_ptree_formula_true;
      unfold ptree_formula in FS; fold ptree_formula in FS;
      try rewrite PF;
      try rewrite P1F;
      try rewrite P2F;
      repeat split;
      try rewrite formula_sub_ptree_deg_atom;
      try apply PD;
      try apply P1D;
      try apply P2D;
      try rewrite formula_sub_ptree_ord_atom;
      try apply PO;
      try apply P1O;
      try apply P2O;
      try apply IHP;
      try apply IHP1;
      try apply IHP2;
      try rewrite formula_sub_ptree_formula_atom;
      try apply PV;
      try apply P1V;
      try apply P2V;
      try rewrite PF;
      try rewrite P1F;
      try rewrite P2F;
      try rewrite formula_sub_ind_lor;
      try rewrite non_target_sub;
      unfold formula_sub_ind;
      unfold subst_ind_fit; fold subst_ind_fit;
      try rewrite FS;
      try rewrite FS1;
      try rewrite FS2;
      try rewrite FS1_1;
      try rewrite FS1_2;
      try rewrite FS1_1_1;
      try rewrite FS1_1_2;
      try rewrite non_target_fit;
      unfold "&&";
      try reflexivity;
      try apply ID;
      try apply IO;
      try apply NO.

all : refine (formula_sub_ind_fit_closed _ _ _ FC _ _ FS1);
      intros Caa;
      apply CF.
Qed.

Lemma formula_sub_ptree_formula_neg' :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                subst_ind_fit (ptree_formula P) S = true ->
                    ptree_formula (formula_sub_ptree P (neg (atom a)) F S) =
                        formula_sub_ind (ptree_formula P) (neg (atom a)) F S.
Proof.
intros P a F PV.
induction P; intros S FS;
unfold formula_sub_ptree;
rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit.

1 : destruct PV as [ID PV].

2 : destruct PV as [[IO PV] NO].

1,2 : unfold ptree_formula in *; fold ptree_formula in *;
      rewrite (formula_sub_ptree_formula_true _ _ _ _ FS);
      apply (IHP PV _ FS).

1 : { unfold ptree_formula, formula_sub_ind in *.
      rewrite FS.
      reflexivity. }

1-4,6-15,18 : destruct S; inversion FS as [FS'];
              try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

2-4,9,12,15,18 : destruct S1; inversion FS' as [FS''];
      try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

4 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

all : unfold ptree_formula in *; fold ptree_formula in *;
      unfold formula_sub_ind, subst_ind_fit; fold subst_ind_fit;
      try rewrite FS;
      try rewrite FS1;
      try rewrite FS2;
      try rewrite FS1_1;
      try rewrite FS1_2;
      try rewrite FS1_1_1;
      try rewrite FS1_1_2;
      try reflexivity.
Qed.

Lemma formula_sub_ptree_formula_neg :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                ptree_formula (formula_sub_ptree P (neg (atom a)) F S) =
                    formula_sub_ind (ptree_formula P) (neg (atom a)) F S.
Proof.
intros P a F PV S.
destruct (subst_ind_fit (ptree_formula P) S) eqn:FS.
- apply (formula_sub_ptree_formula_neg' _ _ _ PV _ FS).
- unfold formula_sub_ptree,formula_sub_ind.
  rewrite FS.
  reflexivity.
Qed.

Lemma formula_sub_ptree_deg_neg :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                ptree_deg (formula_sub_ptree P (neg (atom a)) F S) = ptree_deg P.
Proof.
intros P a F PV.
pose (ptree_formula P) as A.
induction P; intros S;
case (subst_ind_fit A S) eqn:FS;
unfold formula_sub_ptree;
unfold A in FS;
try rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit;
try reflexivity.

1 : destruct PV as [[IO PV] NO].
    unfold ptree_formula in *; fold ptree_formula in *.
    rewrite (formula_sub_ptree_formula_true _ _ _ _ FS).
    apply (IHP PV).

all : destruct S; inversion FS as [FS'];
      try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

2-4,7-11 : destruct S1; inversion FS' as [FS''];
      try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

4 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

all : unfold ptree_deg; fold ptree_deg;
      reflexivity.
Qed.

Lemma formula_sub_ptree_ord_neg :
    forall (P : ptree) (a : atomic_formula) (F : formula),
        valid P ->
            forall (S : subst_ind),
                ptree_ord (formula_sub_ptree P (neg (atom a)) F S) = ptree_ord P.
Proof.
intros P a F PV.
pose (ptree_formula P) as A.
induction P; intros S;
case (subst_ind_fit A S) eqn:FS;
unfold formula_sub_ptree;
unfold A in FS;
try rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit;
try reflexivity.

1 : destruct PV as [ID PV].
    unfold ptree_formula in *; fold ptree_formula in *.
    rewrite (formula_sub_ptree_formula_true _ _ _ _ FS).
    apply (IHP PV).

all : destruct S; inversion FS as [FS'];
      try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

2-4,7-11 : destruct S1; inversion FS' as [FS''];
      try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

4 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

all : unfold ptree_deg; fold ptree_deg;
      reflexivity.
Qed.

Lemma formula_sub_valid_neg :
    forall (P : ptree) (a : atomic_formula) (F : formula),
      valid P ->
          correct_a a = true ->
              closed F = true ->
                  forall (S : subst_ind),
                      subst_ind_fit (ptree_formula P) S = true ->
                          valid (formula_sub_ptree P (neg (atom a)) F S).
Proof.
intros P a F PV Aa CF.
induction P; intros S FS;
unfold formula_sub_ptree;
rewrite FS;
unfold formula_sub_ptree_fit;
fold formula_sub_ptree_fit.

1 : destruct PV as [ID PV].
2 : destruct PV as [[IO PV] NO].
4-9 : destruct PV as [[[PF PV] PD] PO].
10 : destruct PV as [[[[PF FC] PV] PD] PO].
11,12 : destruct PV as [[[[[[[P1F P1V] P2F] P2V] P1D] P2D] P1O] P2O].
13-16 : destruct PV as [[[PF PV] PD] PO].
19-21 : destruct PV as [[[[[[[P1F P1V] P2F] P2V] P1D] P2D] P1O] P2O].

4-7,9,10,12,14,16,18,21 : destruct S; inversion FS as [FS'];
                          try destruct (and_bool_prop _ _ FS') as [FS1 FS2].

5-7,9-13 :  destruct S1; inversion FS' as [FS''];
            try destruct (and_bool_prop _ _ FS1) as [FS1_1 FS1_2].

7 : destruct S1_1; inversion FS'' as [FS'''];
    try destruct (and_bool_prop _ _ FS1_1) as [FS1_1_1 FS1_1_2].

3 : { unfold valid in *.
      destruct f.
      3,4 : inversion PV.
      1 : unfold formula_sub_ind_fit, form_eqb;
          apply PV.
      1 : destruct S;
          unfold formula_sub_ind_fit, form_eqb;
          fold formula_sub_ind_fit form_eqb;
          case (form_eqb f (atom a)) eqn:EQ;
          try apply PV.
          apply form_eqb_eq in EQ.
          rewrite EQ in *.
          unfold PA_omega_axiom in PV.
          pose proof (correct_correctness _ Aa) as Ra.
          rewrite (incorrect_correctness _ PV) in Ra.
          inversion Ra. }

16,17 : repeat split;
        destruct (PV t) as [[[PF PcV] PD] PO];
        rewrite formula_sub_ptree_formula_true;
        try rewrite PF;
        try rewrite formula_sub_ptree_deg_neg;
        try apply PD;
        try rewrite formula_sub_ptree_ord_neg;
        try apply PO;
        try apply X;
        try rewrite formula_sub_ptree_formula_neg;
        try rewrite PF;
        try apply PV;
        try rewrite formula_sub_ind_lor;
        try rewrite (non_target_term_sub _ n (projT1 t));
        try rewrite non_target_sub;
        unfold formula_sub_ind;
        unfold subst_ind_fit; fold subst_ind_fit;
        try rewrite FS;
        try rewrite FS1;
        try rewrite FS2;
        try rewrite non_target_fit;
        try reflexivity.

all : repeat rewrite formula_sub_ptree_formula_true;
      unfold ptree_formula in FS; fold ptree_formula in FS;
      try rewrite PF;
      try rewrite P1F;
      try rewrite P2F;
      repeat split;
      try rewrite formula_sub_ptree_deg_neg;
      try apply PD;
      try apply P1D;
      try apply P2D;
      try rewrite formula_sub_ptree_ord_neg;
      try apply PO;
      try apply P1O;
      try apply P2O;
      try apply IHP;
      try apply IHP1;
      try apply IHP2;
      try rewrite formula_sub_ptree_formula_neg;
      try apply PV;
      try apply P1V;
      try apply P2V;
      try rewrite PF;
      try rewrite P1F;
      try rewrite P2F;
      try rewrite formula_sub_ind_lor;
      try rewrite non_target_sub;
      unfold formula_sub_ind;
      unfold subst_ind_fit; fold subst_ind_fit;
      try rewrite FS;
      try rewrite FS1;
      try rewrite FS2;
      try rewrite FS1_1;
      try rewrite FS1_2;
      try rewrite FS1_1_1;
      try rewrite FS1_1_2;
      try rewrite non_target_fit;
      unfold "&&";
      try reflexivity;
      try apply ID;
      try apply IO;
      try apply NO.

all : refine (formula_sub_ind_fit_closed _ _ _ FC _ _ FS1);
      intros Caa;
      apply CF.
Qed.


(* Copyright (c) 2011. Greg Morrisett, Gang Tan, Joseph Tassarotti,
   Jean-Baptiste Tristan, and Edward Gan.

   This file is part of RockSalt.

   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 2 of
   the License, or (at your option) any later version.
*)

Require Import List.
Require Import ZArith.
Require Import String.

Require Import Shared.Maps.
Require Import Shared.Bits.
Require Import Shared.RTL.
Require Import Shared.Monad.

Require Import MIPSModel.Parser.
Require Import MIPSModel.MIPSSyntax.

Set Implicit Arguments.
Unset Automatic Introduction.

Module MIPS_MACHINE.
  Local Open Scope Z_scope.
  Local Open Scope string_scope.

  Definition size_addr := size32.

  Inductive loc : nat -> Set :=
  | reg_loc : register -> loc size32
  | hi_loc : loc size32
  | lo_loc : loc size32
  | pc_loc : loc size32
    (*Whether or not the next instruction is executed as a branch delay*)
  | bdelay_active_loc : loc size1
    (*The pc to go to after the branch delay*)
  | bdelay_pc_loc : loc size32
.

  Definition location := loc.

  Definition fmap (A B:Type) := A -> B.
  Definition upd A (eq_dec:forall (x y:A),{x=y}+{x<>y}) B (f:fmap A B) (x:A) (v:B) :
    fmap A B := fun y => if eq_dec x y then v else f y.
  Definition look A B (f:fmap A B) (x:A) : B := f x.

  Record mach := {
    gp_regs : fmap register int32 ;
    hi_reg : int32;
    lo_reg : int32;
    pc_reg : int32;
    bdelay_active_f : int1;
    bdelay_pc_reg : int32
  }.
  Definition mach_state := mach.

  Definition get_location s (l:loc s) (m:mach_state) : int s :=
    match l in loc s' return int s' with
      | reg_loc r =>
        match r with
          | Reg i =>
            if Word.eq i Word.zero then Word.zero
            else look (gp_regs m) r
        end
      | hi_loc => hi_reg m
      | lo_loc => lo_reg m
      | pc_loc => pc_reg m
      | bdelay_active_loc => bdelay_active_f m
      | bdelay_pc_loc => bdelay_pc_reg m
    end.

  Definition set_gp_regs r v m :=
    {| gp_regs :=
       match r with
         | Reg i => if Word.eq i Word.zero then gp_regs m
                    else upd register_eq_dec (gp_regs m) r v
       end;
       hi_reg := hi_reg m;
       lo_reg := lo_reg m;
       pc_reg := pc_reg m;
       bdelay_active_f := bdelay_active_f m;
       bdelay_pc_reg := bdelay_pc_reg m
    |}.

  Definition set_pc v m :=
    {| gp_regs := gp_regs m ;
       hi_reg := hi_reg m;
       lo_reg := lo_reg m;
       pc_reg := v;
       bdelay_active_f := bdelay_active_f m;
       bdelay_pc_reg := bdelay_pc_reg m
    |}.

  Definition set_hi v m :=
    {| gp_regs := gp_regs m ;
       hi_reg := v;
       lo_reg := lo_reg m;
       pc_reg := pc_reg m;
       bdelay_active_f := bdelay_active_f m;
       bdelay_pc_reg := bdelay_pc_reg m
    |}.

  Definition set_lo v m :=
    {| gp_regs := gp_regs m ;
       hi_reg := hi_reg m;
       lo_reg := v;
       pc_reg := pc_reg m;
       bdelay_active_f := bdelay_active_f m;
       bdelay_pc_reg := bdelay_pc_reg m
    |}.

  Definition set_bdelay v m :=
    {| gp_regs := gp_regs m ;
       hi_reg := hi_reg m;
       lo_reg := lo_reg m;
       pc_reg := pc_reg m;
       bdelay_active_f := bdelay_active_f m;
       bdelay_pc_reg := v
    |}.

  Definition set_bdelay_f v m :=
    {| gp_regs := gp_regs m ;
       hi_reg := hi_reg m;
       lo_reg := lo_reg m;
       pc_reg := pc_reg m;
       bdelay_active_f := v;
       bdelay_pc_reg := bdelay_pc_reg m
    |}.
  Definition set_location s (l:loc s) (v:int s) m :=
    match l in loc s' return int s' -> mach_state with
      | reg_loc r => fun v => set_gp_regs r v m
      | hi_loc => fun v => set_hi v m
      | lo_loc => fun v => set_lo v m
      | pc_loc => fun v => set_pc v m
      | bdelay_active_loc => fun v => set_bdelay_f v m
      | bdelay_pc_loc => fun v => set_bdelay v m
    end v.

  (* The array component not needed for the MIPS machine *)
  Inductive arr : nat -> nat -> Set :=.
  Definition array := arr.

  Definition array_sub l s (a:array l s) :=
    match a in arr l' s' return int l' -> mach_state -> int s' with
    end.

  Definition array_upd l s (a:array l s) (i:int l) (v:int s) (m:mach_state) :=
    match a in arr l' s' return int l' -> int s' -> mach_state with
    end i v.

End MIPS_MACHINE.

Module MIPS_RTL := RTL.RTL(MIPS_MACHINE).

Module MIPS_Compile.
  Import MIPS_MACHINE.
  Import MIPS_RTL.

  Local Open Scope monad_scope.
  Record conv_state := { c_rev_i : list rtl_instr ; c_next : Z }.
  Definition Conv(T:Type) := conv_state -> T * conv_state.
  Instance Conv_monad : Monad Conv := {
    Return := fun A (x:A) (s:conv_state) => (x,s) ;
    Bind := fun A B (c:Conv A) (f:A -> Conv B) (s:conv_state) =>
      let (v,s') := c s in f v s'
  }.
  intros ; apply Coqlib.extensionality ; auto.
  intros ; apply Coqlib.extensionality ; intros. destruct (c x). auto.
  intros ; apply Coqlib.extensionality ; intros. destruct (f x) ; auto.
  Defined.
  Definition runConv (c:Conv unit) : (list rtl_instr) :=
    match c {|c_rev_i := nil ; c_next:=0|} with
      | (_, c') => (List.rev (c_rev_i c'))
    end.
  Definition EMIT(i:rtl_instr) : Conv unit :=
    fun s => (tt,{|c_rev_i := i::(c_rev_i s) ; c_next := c_next s|}).
  Notation "'emit' i" := (EMIT i) (at level 75) : monad_scope.

  Definition fresh s (almost_i : pseudo_reg s -> rtl_instr) : Conv (pseudo_reg s) :=
    fun ts => let r := c_next ts in
              let ts' := {|c_rev_i := (almost_i (ps_reg s r))::c_rev_i ts ;
                           c_next := r + 1|} in
                (ps_reg s r, ts').

  (* Begin: a set of basic conversion constructs *)
  Definition raise_error := emit error_rtl.
  Definition raise_trap := emit trap_rtl.
  Definition no_op := ret tt.
  (* Definition ret_exp s (e:rtl_exp s) := ret e. *)
  Definition load_int s (i:int s) := ret (imm_rtl_exp i).
  Definition arith s b (e1 e2:rtl_exp s) := ret (arith_rtl_exp b e1 e2).
  Definition test s t (e1 e2:rtl_exp s) := ret (test_rtl_exp t e1 e2).
  Definition cast_u s1 s2 (e:rtl_exp s1) := ret (@cast_u_rtl_exp s1 s2 e).
  Definition cast_s s1 s2 (e:rtl_exp s1) := ret (@cast_s_rtl_exp s1 s2 e).
  Definition read_loc s (l:loc s) := ret (get_loc_rtl_exp l).
  Definition write_loc s (e:rtl_exp s) (l:loc s)  := emit set_loc_rtl e l.

  (* store the value of e into the current pseudo reg and advance the
     index of the pseudo reg; it returns an rtl_exp that retrives the
     value from the storage; *)
  Definition write_ps_and_fresh s (e: rtl_exp s) : Conv (rtl_exp s) :=
    fun ts =>
      let r := c_next ts in
      let ts' := {|c_rev_i := (set_ps_reg_rtl e (ps_reg s r))::c_rev_i ts;
                   c_next := r + 1|} in
      (get_ps_reg_rtl_exp (ps_reg s r), ts').

  Definition read_byte (a:rtl_exp size32) := ret (get_byte_rtl_exp a).
  Definition write_byte (v:rtl_exp size8) (a:rtl_exp size32) :=
    emit set_byte_rtl v a.
  Definition if_exp s g (e1 e2:rtl_exp s) : Conv (rtl_exp s) :=
    ret (if_rtl_exp g e1 e2).
  Definition if_trap g : Conv unit := emit (if_rtl g trap_rtl).
  Definition if_set_loc cond s (e:rtl_exp s) (l:location s) :=
    emit (if_rtl cond (set_loc_rtl e l)).
  Definition choose s : Conv (rtl_exp s) :=
    emit (advance_oracle_rtl);;
    ret (@get_random_rtl_exp s).

  Definition load_Z s (z:Z) := load_int (@Word.repr s z).

  Definition load_reg (r:register) := read_loc (reg_loc r).
  Definition set_reg (p:rtl_exp size32) (r:register) :=
    write_loc p (reg_loc r).

  Definition get_pc := read_loc pc_loc.
  Definition set_pc v := write_loc v pc_loc.

  Definition get_hi := read_loc hi_loc.
  Definition get_lo := read_loc lo_loc.
  Definition set_hi v := write_loc v hi_loc.
  Definition set_lo v := write_loc v lo_loc.

  Definition get_bdelayf := read_loc bdelay_active_loc.
  Definition get_bdelaypc := read_loc bdelay_pc_loc.

  (*Helper functions for setting the pc in a branch delayed way with / without a
     conditional flag*)
  Definition set_bdelay_cond f v :=
    ttrue <- load_int Word.one;
    if_set_loc f ttrue bdelay_active_loc;;
    if_set_loc f v bdelay_pc_loc.

  Definition set_bdelay v:=
    ttrue <- load_int Word.one;
    write_loc ttrue bdelay_active_loc;;
    write_loc v bdelay_pc_loc.

  Definition clear_bdelay :=
    tzero <- load_int Word.zero;
    tfalse <- load_int Word.zero;
    write_loc tfalse bdelay_active_loc;;
    write_loc tzero bdelay_pc_loc.

  (*If a delay is present, set the pc. Always clear the delay*)
  (*This check should be done before every non-cflow instruction*)
  Definition bd_prologue :=
    bdf <- get_bdelayf;
    bdpc <- get_bdelaypc;
    if_set_loc bdf bdpc pc_loc;;
    clear_bdelay.

  (*If a delay is present for a control flow instruction, undefined behavior*)
  Definition bd_prologue_cf :=
    bdf <- get_bdelayf;
    emit (if_rtl bdf (error_rtl)).

  Definition not {s} (p: rtl_exp s) : Conv (rtl_exp s) :=
    mask <- load_Z s (Word.max_unsigned s);
    arith xor_op p mask.

  Definition bnot (p: rtl_exp size1) : Conv (rtl_exp size1) :=
    pone <- load_Z size1 1%Z;
    arith xor_op pone p.

  Definition neg {s} (p: rtl_exp s) : Conv (rtl_exp s) :=
    pzero <- load_Z s 0%Z;
    arith sub_op pzero p.

  Definition undef s := choose s.

  (* Copy the contents of rs to a new pseudo register *)
  (* Definition copy_ps s (rs:pseudo_reg s) := fresh (@cast_u_rtl s s rs). *)

  Definition onemaskl s n : Conv (rtl_exp s) :=
    numshift <- load_Z s ((Z_of_nat s)+1-n);
    mask <- load_Z s (Word.max_unsigned s);
    maskr <- arith shr_op mask numshift;
    arith shl_op maskr numshift.

  (* compute an effective address *)
  Definition compute_addr (r:register) (i:int16) : Conv (rtl_exp size32) :=
    b <- load_reg r;
    disp16 <- load_int i;
    disp32 <- cast_s size32 disp16;
    arith add_op b disp32.

  (* load a byte from memory, taking into account the specified segment *)
  Definition lmem (a:rtl_exp size32) : Conv (rtl_exp size8):=
    read_byte a.

  (* store a byte to memory, taking into account the specified segment *)
  Definition smem (v:rtl_exp size8) (a:rtl_exp size32) :
    Conv unit :=
    write_byte v a.

  (* load an n-byte vector from memory -- takes into account the segment *)
  Program Fixpoint load_mem_n (addr:rtl_exp size32)
    (nbytes_minus_one:nat) : Conv (rtl_exp ((nbytes_minus_one+1) * 8 -1)%nat) :=
    match nbytes_minus_one with
      | 0 => lmem addr
      | S n =>
        rec <- load_mem_n addr n ;
        count <- load_Z size32 (Z_of_nat (S n)) ;
        p3 <- arith add_op addr count ;
        nb <- lmem p3 ;
        p5 <- cast_u ((nbytes_minus_one + 1)*8-1)%nat rec ;
        p6 <- cast_u ((nbytes_minus_one + 1)*8-1)%nat nb ;
        p7 <- load_Z _ (Z_of_nat (S n) * 8) ;
        p8 <- arith shl_op p6 p7 ;
        arith or_op p5 p8
    end.

  Definition load_mem32 (addr:rtl_exp size32) :=
    load_mem_n addr 3.
  Definition load_mem16 (addr:rtl_exp size32) :=
    load_mem_n addr 1.
  Definition load_mem8 (addr:rtl_exp size32) :=
    load_mem_n addr 0.
  (* set memory with an n-byte vector *)
  Program Fixpoint set_mem_n {t}
    (v: rtl_exp (8*(t+1)-1)%nat) (addr : rtl_exp size32) : Conv unit :=
    match t with
      | 0 => smem v addr
      | S u =>
        p1 <- cast_u (8*(u+1)-1)%nat v ;
        set_mem_n p1 addr ;;
        p2 <- load_Z (8*(t+1)-1)%nat (Z_of_nat  ((S u) * 8)) ;
        p3 <- arith shru_op v p2 ;
        p4 <- cast_u size8 p3 ;
        p5 <- load_Z size32 (Z_of_nat (S u)) ;
        p6 <- arith add_op p5 addr ;
        smem p4 p6
    end.

  Definition set_mem32  (v a:rtl_exp size32) : Conv unit :=
    @set_mem_n 3 v a.
  Definition set_mem16  (v:rtl_exp size16)
    (a:rtl_exp size32) : Conv unit :=
      @set_mem_n 1 v a.
  Definition set_mem8  (v: rtl_exp size8)
    (a:rtl_exp size32) : Conv unit :=
      @set_mem_n 0 v a.

  (**********************************************)
  (*   Conversion functions for instructions    *)
  (**********************************************)

  (************************)
  (* Arith ops            *)
  (************************)

  Definition OF_check_add (a1:rtl_exp size32) a2 sum : Conv unit :=
        zero <- load_Z _ 0;
        sign1 <- test lt_op zero a1;
        sign2 <- test lt_op zero a2;
        signsum <- test lt_op zero sum;
        sign_op_eq <- test eq_op sign1 sign2;
        sign_res_neq <- @arith size1 xor_op sign1 signsum;
        of_flag <- @arith size1 and_op sign_op_eq sign_res_neq;
        (*Throw an exception if there was an overflow*)
        if_trap of_flag.

  Definition OF_check_sub (a1:rtl_exp size32) a2 diff : Conv unit :=
        zero <- load_Z _ 0;
        sign1 <- test lt_op zero a1;
        sign2 <- test lt_op zero a2;
        signdiff <- test lt_op zero diff;
        sign_op_neq <- arith xor_op sign1 sign2;
        sign_res_neq <- @arith size1 xor_op sign1 signdiff;
        of_flag <- @arith size1 and_op sign_op_neq sign_res_neq;
        (*Throw an exception if there was an overflow*)
        if_trap of_flag.

  Definition conv_ADD (op: reg3_operand) : Conv unit :=
    match op with
      | Reg3_op rs rt rd =>
        a1 <- load_reg rs;
        a2 <- load_reg rt;
        sum <- arith add_op a1 a2;
        OF_check_add a1 a2 sum;;
        set_reg sum rd
    end.

  Definition conv_ADDU (op: reg3_operand) : Conv unit :=
    match op with
      | Reg3_op rs rt rd =>
        a1 <- load_reg rs;
        a2 <- load_reg rt;
        sum <- arith add_op a1 a2;
        set_reg sum rd
    end.

  Definition conv_SUB (op:reg3_operand) : Conv unit :=
    match op with | Reg3_op rs rt rd =>
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      diff <- arith sub_op a1 a2;
      OF_check_sub a1 a2 diff;;
      set_reg diff rd
    end.

  Definition conv_SUBU (op:reg3_operand) : Conv unit :=
    match op with | Reg3_op rs rt rd =>
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      diff <- arith sub_op a1 a2;
      set_reg diff rd
    end.

  Definition conv_ADDIs (sign:bool) (iop: ioperand) : Conv unit :=
    match iop with
      | Iop rs rt i =>
        a1 <- load_reg rs;
        pi <- load_int i;
        a2 <- cast_s size32 pi;
        sum <- arith add_op a1 a2;
        match sign with
          | true =>
            OF_check_add a1 a2 sum;;
            set_reg sum rt
          | false =>
            set_reg sum rt
        end
    end.
  Definition conv_ADDI := conv_ADDIs true.
  Definition conv_ADDIU := conv_ADDIs false.

  Definition conv_MUL (op: reg3_operand) : Conv unit :=
    match op with | Reg3_op rs rt rd  =>
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      prod <- arith mul_op a1 a2;
      set_reg prod rd;;
      rand1 <- undef size32;
      set_hi rand1;;
      rand2 <- undef size32;
      set_lo rand2
    end.

  Definition conv_MULTs caster (op: reg2_operand) : Conv unit :=
    match op with | Reg2_op rs rt =>
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      a1e <- caster a1;
      a2e <- caster a2;
      prod <- arith mul_op a1e a2e;
      prodlo <- cast_u size32 prod;
      p32 <- load_Z 63 32;
      prods <- arith shru_op prod p32;
      prodhi <- cast_u size32 prods;
      set_hi prodhi;;
      set_lo prodlo
    end.
  Definition conv_MULT (op:reg2_operand) : Conv unit := conv_MULTs (cast_s 63) op.
  Definition conv_MULTU (op:reg2_operand) : Conv unit := conv_MULTs (cast_u 63) op.

  (*How is the signed modulus taken?*)
  Definition conv_DIVs (divop:bit_vector_op) (modop:bit_vector_op) (op: reg2_operand) : Conv unit :=
    match op with | Reg2_op rs rt =>
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      pzero <- load_Z size32 0;
      divbyzero <- test eq_op a2 pzero;
      emit (if_rtl divbyzero error_rtl);;
      quo <- arith divop a1 a2;
      rem <- arith modop a1 a2;
      set_lo quo;;
      set_hi rem
    end.
  Definition conv_DIV := conv_DIVs divs_op mods_op.
  Definition conv_DIVU := conv_DIVs divu_op modu_op.

  Definition conv_MFHI (rd: register) : Conv unit :=
    hi <- get_hi;
    set_reg hi rd.

  Definition conv_MFLO (rd: register) : Conv unit :=
    lo <- get_lo;
    set_reg lo rd.

  Definition conv_LUI (iop:ioperand) : Conv unit :=
    match iop with | Iop _ rt i =>
      pi <- load_int i;
      a1 <- cast_u size32 pi;
      p16 <- load_Z size32 16;
      a1s <- arith shl_op a1 p16;
      set_reg a1s rt
    end.
  (*-----Logical Ops-------*)
  Definition conv_log (op1: bit_vector_op) (op: reg3_operand) : Conv unit :=
    match op with
      | Reg3_op rs rt rd =>
        a1 <- load_reg rs;
        a2 <- load_reg rt;
        a3 <- arith op1 a1 a2;
        set_reg a3 rd
    end.

  Definition conv_AND := conv_log and_op.
  Definition conv_OR := conv_log or_op.
  Definition conv_XOR := conv_log xor_op.

  Definition conv_NOR (op: reg3_operand) : Conv unit :=
    match op with | Reg3_op rs rt rd =>
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      v_or <- arith or_op a1 a2;
      v_nor <- not v_or;
      set_reg v_nor rd
    end.

  Definition conv_logi (op1: bit_vector_op) (iop:ioperand): Conv unit :=
    match iop with
      | Iop rs rt i =>
        a1 <- load_reg rs;
        pi <- load_int i;
        a2 <- cast_u size32 pi;
        a3 <- arith op1 a1 a2;
        set_reg a3 rt
    end.
  Definition conv_ANDI := conv_logi and_op.
  Definition conv_ORI := conv_logi or_op.
  Definition conv_XORI := conv_logi xor_op.

  Definition conv_sh (op1: bit_vector_op) (op: reg2sh_operand) : Conv unit :=
    match op with | Reg2sh_op rt rd sa =>
      a1 <- load_reg rt;
      si <- load_int sa;
      a2 <- cast_u size32 si;
      a1s <- arith op1 a1 a2;
      set_reg a1s rd
    end.

  Definition conv_SLL := conv_sh shl_op.
  Definition conv_SRA := conv_sh shr_op.
  Definition conv_SRL := conv_sh shru_op.

  Definition conv_shv (op1: bit_vector_op) (op: reg3_operand) : Conv unit :=
    match op with | Reg3_op rs rt rd =>
      a1 <- load_reg rt;
      a2 <- load_reg rs;
      shamt5 <- cast_u 4 a2;
      shamt <- cast_u size32 shamt5;
      a3 <- arith op1 a1 shamt;
      set_reg a3 rd
    end.

  Definition conv_SLLV := conv_shv shl_op.
  Definition conv_SRAV := conv_shv shru_op.
  Definition conv_SRLV := conv_shv shr_op.

  Definition conv_set (op1: test_op) (op: reg3_operand) : Conv unit :=
    match op with | Reg3_op rs rt rd =>
      pone <- load_Z size32 1;
      pzero <- load_Z size32 0;
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      cond <- test op1 a1 a2;
      u <- cast_u _ pone;
      v_rd <- if_exp cond u pzero;
      set_reg v_rd rd
    end.

  Definition conv_SLT := conv_set lt_op.
  Definition conv_SLTU := conv_set ltu_op.

  Definition conv_seti caster (op1:test_op) (iop: ioperand) : Conv unit :=
    match iop with | Iop rs rt i =>
      pone <- load_Z size32 1;
      pzero <- load_Z size32 0;
      a1 <- load_reg rs;
      a2 <- load_int i;
      a2c <- caster a2;
      cond <- test op1 a1 a2c;
      u <- cast_u _ pone;
      v_rd <- if_exp cond u pzero;
      set_reg v_rd rt
    end.

  Definition conv_SLTI := conv_seti (cast_s size32) (lt_op).
  Definition conv_SLTIU := conv_seti (cast_u size32) (ltu_op).

  Definition conv_SEB (op:reg2_operand) : Conv unit :=
    match op with | Reg2_op rt rd =>
      a1 <- load_reg rt;
      a1b <- cast_u size8 a1;
      a1bc <- cast_s size32 a1b;
      set_reg a1bc rd
    end.

  Definition conv_SEH (op:reg2_operand) : Conv unit :=
    match op with | Reg2_op rt rd =>
      a1 <- load_reg rt;
      a1h <- cast_u size16 a1;
      a1hc <- cast_s size32 a1h;
      set_reg a1hc rd
    end.

  (*--------------------*)
  (* Memory Ops         *)
  (*--------------------*)

  Definition testparity (p:Z) (v:rtl_exp size32) : Conv (rtl_exp size1) :=
    pval <- load_Z size32 p;
    rem <- arith modu_op v pval;
    pzero <- load_Z size32 0;
    test eq_op rem pzero.

  Definition conv_L_tlc s
    (tester:rtl_exp size32 -> Conv (rtl_exp size1))
    (loader:rtl_exp size32 -> Conv (rtl_exp s))
    (caster:rtl_exp s -> Conv (rtl_exp size32))
      (iop: ioperand) : Conv unit :=
    match iop with
      | Iop rs rt i =>
        addr <- compute_addr rs i;
        addrgood <- tester addr;
        addrerrorf <- bnot addrgood;
        emit (if_rtl addrerrorf error_rtl);;
        membyte <- loader addr;
        m_extend <- caster membyte;
        set_reg m_extend rt
    end.

  Definition conv_LB (iop:ioperand) :=
    conv_L_tlc (testparity 1) (load_mem8) (cast_s size32) iop.
  Definition conv_LBU (iop:ioperand) :=
    conv_L_tlc (testparity 1) (load_mem8) (cast_u size32) iop.

  Definition conv_LH (iop:ioperand) :=
    conv_L_tlc (testparity 2) (load_mem16) (cast_s size32) iop.
  Definition conv_LHU (iop:ioperand) :=
    conv_L_tlc (testparity 2) (load_mem16) (cast_u size32) iop.

  Definition conv_LW (iop:ioperand) :=
    conv_L_tlc (testparity 4) (load_mem32) (cast_u size32) iop.

  Definition conv_S_tlc s
    (tester:rtl_exp size32 -> Conv (rtl_exp size1))
    (saver:rtl_exp s -> rtl_exp size32 -> Conv unit)
    (caster:rtl_exp size32 -> Conv (rtl_exp s))
      (iop: ioperand) : Conv unit :=
    match iop with
      | Iop rs rt i =>
        addr <- compute_addr rs i;
        addrgood <- tester addr;
        addrerrorf <- bnot addrgood;
        emit (if_rtl addrerrorf error_rtl);;

        rval <- load_reg rt;
        rvalc <- caster rval;
        saver rvalc addr
    end.

  Definition conv_SB (iop:ioperand) :=
    conv_S_tlc (testparity 1) (@set_mem8) (cast_u size8) iop.
  Definition conv_SH (iop:ioperand) :=
    conv_S_tlc (testparity 2) (@set_mem16) (cast_u size16) iop.
  Definition conv_SW (iop:ioperand) :=
    conv_S_tlc (testparity 4) (@set_mem32) (cast_u size32) iop.

  (*--------------------*)
  (* Control Flow Ops   *)
  (*--------------------*)

  (*Shift and cast ints into the actual memory offset*)
  Definition int2moffset_u s (i: int s) : Conv (rtl_exp size32) :=
    ibits <- load_int i;
    ioffset <- cast_u size32 ibits;
    ptwo <- load_Z size32 2;
    ishifted <- arith shl_op ioffset ptwo;
    ret ishifted.

  Definition int2moffset_s s (i: int s) : Conv (rtl_exp size32) :=
    ibits <- load_int i;
    ioffset <- cast_s size32 ibits;
    ptwo <- load_Z size32 2;
    ishifted <- arith shl_op ioffset ptwo;
    ret ishifted.

  (*pc is incremented before execution*)
  (*pc isn't actually set until next step, so we set the branch delay state of the machine here*)
  Definition conv_Jl (linkflag: bool) (jop: joperand) : Conv unit :=
    match jop with
      | Jop i =>
        curpc <- get_pc;
        pcmask <- onemaskl size32 4;
        pcmasked <- arith and_op curpc pcmask;
        ioffset <- int2moffset_u i;
        newpc <- arith or_op pcmasked ioffset;
        match (linkflag) with
          | false =>
            set_bdelay newpc
          | true =>
            pfour <- load_Z size32 4;
            pc4 <- arith add_op curpc pfour;
            set_reg pc4 (Reg (Word.repr 31));;
            set_bdelay newpc
        end
    end.

  Definition conv_J (jop: joperand) : Conv unit := conv_Jl false jop.
  Definition conv_JAL (jop: joperand) : Conv unit := conv_Jl true jop.

  Definition conv_JALR (op:reg2sh_operand) : Conv unit :=
    match op with | Reg2sh_op rs rd _ =>
      if (register_eq_dec rs rd) then (emit error_rtl) else (
        curpc <- get_pc;
        newpc <- load_reg rs;
        pfour <- load_Z size32 4;
        retpc <- arith add_op curpc pfour;
        set_bdelay newpc;;
        set_reg retpc rd
      )
    end.

  Definition conv_JR (rs: register) : Conv unit :=
    curpc <- get_pc;
    newpc <- load_reg rs;
    set_bdelay newpc.

  Definition b_getnewpc s (i: int s) : Conv (rtl_exp size32) :=
    curpc <- get_pc;
    ioffset <- int2moffset_s i;
    arith add_op curpc ioffset.

  Inductive condition2 : Set :=
  | Eq_cond : condition2
  | Ne_cond : condition2.

  Inductive condition1 : Set :=
  | Gez_cond : condition1
  | Gtz_cond : condition1
  | Lez_cond : condition1
  | Ltz_cond : condition1.

  Definition conv_cond1 s (c:condition1) (a:rtl_exp s) : Conv (rtl_exp size1) :=
    pzero <- load_Z s 0%Z;
    match c with
      | Gez_cond =>
        condf1 <- test lt_op pzero a;
        condf2 <- test eq_op pzero a;
        arith or_op condf1 condf2
      | Gtz_cond =>
        test lt_op pzero a
      | Lez_cond =>
        condf1 <- test lt_op a pzero;
        condf2 <- test eq_op a pzero;
        arith or_op condf1 condf2
      | Ltz_cond =>
        test lt_op a pzero
    end.
  Definition conv_cond2 s (c:condition2) (a1:rtl_exp s) (a2:rtl_exp s) :
    Conv (rtl_exp size1) :=
    match c with
      | Eq_cond => test eq_op a1 a2
      | Ne_cond => condf1 <- test eq_op a1 a2; bnot condf1
    end.

  (*Two argument comparison branches*)
  Definition conv_Bcond2 (c:condition2) (iop:ioperand) : Conv unit :=
    match iop with | Iop rs rt i =>
      newpc <- b_getnewpc i;
      a1 <- load_reg rs;
      a2 <- load_reg rt;
      condf <- conv_cond2 c a1 a2;
      set_bdelay_cond condf newpc
    end.
  Definition conv_BEQ (iop:ioperand) : Conv unit := conv_Bcond2 Eq_cond iop.
  Definition conv_BNE (iop:ioperand) : Conv unit := conv_Bcond2 Ne_cond iop.

  (*Single argument comparison branches*)
  Definition conv_Bcond1 (c:condition1) (bzop:bz_operand) : Conv unit :=
    match bzop with | BZ_op rs i =>
      newpc <- b_getnewpc i;
      a1 <- load_reg rs;
      condf <- conv_cond1 c a1;
      set_bdelay_cond condf newpc
    end.

  Definition conv_Bcondl1 (c:condition1) (bzop:bz_operand) : Conv unit :=
    match bzop with | BZ_op rs i =>
      curpc <- get_pc;
      pfour <- load_Z size32 4;
      retpc <- arith add_op curpc pfour;
      set_reg retpc (Reg (Word.repr 31));;
      newpc <- b_getnewpc i;
      a1 <- load_reg rs;
      condf <- conv_cond1 c a1;
      set_bdelay_cond condf newpc
    end.

  Definition conv_BGEZ (bzop:bz_operand) : Conv unit := conv_Bcond1 Gez_cond bzop.
  Definition conv_BGEZAL := conv_Bcondl1 Gez_cond.
  Definition conv_BGTZ (bzop:bz_operand) : Conv unit := conv_Bcond1 Gtz_cond bzop.
  Definition conv_BLEZ (bzop:bz_operand) : Conv unit := conv_Bcond1 Lez_cond bzop.
  Definition conv_BLTZ (bzop:bz_operand) : Conv unit := conv_Bcond1 Ltz_cond bzop.
  Definition conv_BLTZAL := conv_Bcondl1 Ltz_cond.
  (*---------End specific instruction implementations--------------*)

  (*Predicate to distinguish between control flow and noncflow instructions*)
  Definition iscflow (i:instr ) : bool:=
    match i with
      | BEQ _ | BGEZ _ | BGTZ _ | BLEZ _ | BLTZ _ | BNE _  => true
      | J _ | JAL _ | JALR _ | JR _ => true
      | _  => false
    end.

  Definition incr_pc : Conv unit :=
    pwsize <- load_Z size32 (4%Z);
    curpc <- get_pc;
    newpc <- arith add_op curpc pwsize;
    set_pc newpc
    .

  (*Main translation function*)
  Definition instr_to_rtl (i: instr) :=
    runConv (
      (*Increment the PC before executing any instruction.*)
      incr_pc;;
      (*Different preprocessing behaviors are needed to deal with branch delayed instructions*)
      match (iscflow i) with
        | true => bd_prologue_cf
        | false => bd_prologue
      end;;
      match i with
        | ADD rop => conv_ADD rop
        | ADDI iop => conv_ADDI iop
        | ADDIU iop => conv_ADDIU iop
        | ADDU rop => conv_ADDU rop
        | AND rop => conv_AND rop
        | ANDI iop => conv_ANDI iop
        | BEQ iop => conv_BEQ iop
        | BGEZ iop => conv_BGEZ iop
        | BGEZAL iop => conv_BGEZAL iop
        | BGTZ iop => conv_BGTZ iop
        | BLEZ iop => conv_BLEZ iop
        | BLTZ iop => conv_BLTZ iop
        | BLTZAL iop => conv_BLTZAL iop
        | BNE iop => conv_BNE iop
        | DIV rop => conv_DIV rop
        | DIVU rop => conv_DIVU rop
        | J jop => conv_J jop
        | JAL jop => conv_J jop
        | JALR rop => conv_JALR rop
        | JR rop => conv_JR rop
        | LB iop => conv_LB iop
        | LBU iop => conv_LBU iop
        | LH iop => conv_LH iop
        | LHU iop => conv_LHU iop
        | LUI iop => conv_LUI iop
        | LW iop => conv_LW iop
        | MFHI rop => conv_MFHI rop
        | MFLO rop => conv_MFLO rop
        | MUL rop => conv_MUL rop
        | MULT rop => conv_MULT rop
        | MULTU rop => conv_MULTU rop
        | NOR rop => conv_NOR rop
        | OR rop => conv_OR rop
        | ORI iop => conv_ORI iop
        | SB iop => conv_SB iop
        | SEB rop => conv_SEB rop
        | SEH rop => conv_SEH rop
        | SH iop => conv_SH iop
        | SLL rop => conv_SLL rop
        | SLLV rop => conv_SLLV rop
        | SLT rop => conv_SLT rop
        | SLTI iop => conv_SLTI iop
        | SLTIU iop => conv_SLTIU iop
        | SLTU rop => conv_SLTU rop
        | SRA rop => conv_SRA rop
        | SRAV rop => conv_SRAV rop
        | SRL rop => conv_SRL rop
        | SRLV rop => conv_SRLV rop
        | SUB rop => conv_SUB rop
        | SUBU rop => conv_SUBU rop
        | SW iop => conv_SW iop
        | XOR rop => conv_XOR rop
        | XORI iop => conv_XORI iop
      end
    ).

End MIPS_Compile.

Local Open Scope Z_scope.
Local Open Scope monad_scope.
Import MIPS_Compile.
Import MIPS_RTL.
Import MIPS_MACHINE.

Fixpoint RTL_step_list l :=
  match l with
    | nil => ret tt
    | i::l' => interp_rtl i;; RTL_step_list l'
  end.

Definition step_ins inst: RTL unit :=
  RTL_step_list (MIPS_Compile.instr_to_rtl inst).


Require Import FiatFormal.Language.Step.
Require Import FiatFormal.Language.SubstExpExp.
Require Import FiatFormal.Language.SubstTypeExp.
Require Import FiatFormal.Language.SubstTypeType.
Require Import FiatFormal.Language.TyJudge.
Require Import FiatFormal.Language.Progress.

(* A well typed program is either a value, can take a step, or has a choice. *)
Theorem progressP
  : forall pb ds pbOK p,
    (exists t, TYPEPROG ds nil nil p t)
    -> valueP p \/ (exists p', STEPP pb ds pbOK p p') \/ hasChoiceP p.
Proof.
  intros. gen pb ds pbOK. induction p; intros.
  Case "PLet".
  right. left.
  destruct H as [tp].
  inverts H.
  invert_adt_type.
  eexists.
  eapply SP_Prog.
  Case "PExp".
  inverts H.
  invert_prog_type.
  rename x into t.
  pose proof (progress pb ds pbOK e).
  destruct H as [| []];

  eauto.
  SCase "e value".
  left; repeat constructor.
  inversion H; auto.
  unfold closedP. simpl.
  inversion H. unfold closedX in *. assumption.

  SCase "e steps".
  right. left. dest H. exists (PEXP x'). auto.
Qed.


Require Import Rupicola.Lib.Api.
Require Import Rupicola.Examples.KVStore.KVStore.
Require Import Rupicola.Examples.KVStore.Properties.
Require Import Rupicola.Examples.KVStore.Tactics.

Require Import bedrock2.NotationsCustomEntry.

Local Open Scope nat_scope.

Section examples.
  Context {width: Z} {BW: Bitwidth width} {word: word.word width} {mem: map.map word Byte.byte}.
  Context {locals: map.map String.string word}.
  Context {env: map.map String.string (list String.string * list String.string * Syntax.cmd)}.
  Context {ext_spec: bedrock2.Semantics.ExtSpec}.
  Context {word_ok : word.ok word} {mem_ok : map.ok mem}.
  Context {locals_ok : map.ok locals}.
  Context {env_ok : map.ok env}.
  Context {ext_spec_ok : Semantics.ext_spec.ok ext_spec}.

  Section put_sum.
    Context {add : func}.
    Definition Int (addr : word) (x : Z) : mem -> Prop :=
      sep (emp (0 <= x < 2^width)%Z)
          (truncated_scalar access_size.word addr x).

    Context {ops} {key : Type}
            {kvp : kv_parameters}
            {ok : @kv_parameters_ok  _ BW _ mem ops key Z Int kvp}.

    Existing Instances map_ok annotated_map_ok key_eq_dec.
    Local Hint Extern 1 (spec_of "get") => unshelve simple refine (@spec_of_map_get _ _ _ _ _ _ _ _ _ _ _) : typeclass_instances.
    Local Hint Extern 1 (spec_of "put") => unshelve simple refine (@spec_of_map_put _ _ _ _ _ _ _ _ _ _ _) : typeclass_instances.

    Instance spec_of_add : spec_of "add" :=
      fun functions =>
        forall px x py y pout old_out R tr mem,
          (Int px x * Int py y * Int pout old_out * R)%sep mem ->
          WeakestPrecondition.call
            functions "add" tr mem [px; py; pout]
            (fun tr' mem' rets =>
               tr = tr'
               /\ rets = []
               /\ let out := word.wrap (x + y) in
                  (Int px x * Int py y * Int pout out * R)%sep mem').

    (* look up k1 and k2, add their values and store in k3 *)
    Definition put_sum_gallina (m : map.rep (map:=map))
               (k1 k2 k3 : key) : map.rep (map:=map) :=
      match map.get m k1, map.get m k2 with
      | Some v1, Some v2 => map.put m k3 (word.wrap (v1 + v2)%Z)
      | _, _ => m
      end.

    (* like put, returns a boolean indicating whether the put was an
         overwrite, and stores the old value in v if the boolean is true *)
    Definition put_sum := func! (m, k1, k2, k3, v) ~> ret {
      unpack! err, v1 = get(m, k1) ;
      require !err ;
      unpack! err, v2 = get(m, k2) ;
      require !err ;
      add(v1,v2,v);
      unpack! ret = put(m, k3, v)
    }.

    Instance spec_of_put_sum : spec_of "put_sum" :=
      fun functions =>
        forall pm m pk1 k1 pk2 k2 pk3 k3 pv v R tr mem,
          map.get m k1 <> None ->
          map.get m k2 <> None ->
          k1 <> k2 -> (* required because add needs arguments separate *)
          (Map pm m * Key pk1 k1 * Key pk2 k2 * Key pk3 k3
           * Int pv v * R)%sep mem ->
          WeakestPrecondition.call
            functions "put_sum" tr mem [pm; pk1; pk2; pk3; pv]
            (fun tr' mem' rets =>
               tr = tr'
               /\ length rets = 1
               /\ hd (word.of_Z 0) rets =
                  match map.get m k3 with
                  | Some _ => word.of_Z 1
                  | None => word.of_Z 0
                  end
               /\ (Map pm (put_sum_gallina m k1 k2 k3)
                   * Key pk1 k1 * Key pk2 k2 *
                   (match map.get m k3 with
                    | Some v3 =>
                      Key pk3 k3 * Int pv v3 * R
                    | None => R
                    end))%sep mem').

    (* Entire chain of separation-logic reasoning for put_sum
         (omitting keys for readability):

         { pm -> m; pv -> v}
         // annotate_iff1
         { pm -> annotate m; pv -> v}
         // get k1
         { pm -> map.put (annotate m) k1 (Reserved pv1, v1); pv -> v}
         // reserved_borrowed_iff1
         { pm -> map.put (annotate m) k1 (Borrowed pv1, v1);
                 pv1 -> v1; pv -> v}
         // get k2
         { pm -> map.put (map.put (annotate m) k1 (Borrowed pv1, v1))
                         k2 (Reserved pv2, v2); pv1 -> v1; pv -> v}
         // reserved_borrowed_iff1
         { pm -> map.put (map.put (annotate m) k1 (Borrowed pv1, v1))
                         k2 (Borrowed pv2, v2);
                 pv2 -> v2; pv1 -> v1; pv -> v}
         // add
         { pm -> map.put (map.put (annotate m) k1 (Borrowed pv1, v1))
                         k2 (Borrowed pv2, v2);
                 pv2 -> v2; pv1 -> v1; pv -> (v1 + v2)}
         // reserved_borrowed_iff1
         { pm -> map.put (map.put (annotate m) k1 (Borrowed pv1, v1))
                         k2 (Reserved pv2, v2); pv1 -> v1; pv -> (v1 + v2)}
         // commutativity of put (when k1 <> k2)
         { pm -> map.put (map.put (annotate m) k2 (Reserved pv2, v2))
                         k1 (Borrowed pv1, v1); pv1 -> v1; pv -> (v1 + v2)}
         // reserved_borrowed_iff1
         { pm -> map.put (map.put (annotate m) k2 (Reserved pv2, v2))
                         k1 (Reserved pv1, v1); pv -> (v1 + v2)}
         // reserved_impl1
         { pm -> map.put (map.put (annotate m) k2 (Reserved pv2, v2))
                         k1 (Owned, v1); pv -> (v1 + v2)}
         // commutativity of put (when k1 <> k2)
         { pm -> map.put (map.put (annotate m) k1 (Owned, v1))
                         k2 (Reserved pv2, v2); pv -> (v1 + v2)}
         // reserved_impl1
         { pm -> map.put (map.put (annotate m) k1 (Owned, v1))
                         k2 (Owned, v2); pv -> (v1 + v2)}
         // puts are noops
         { pm -> annotate m; pv -> (v1 + v2)}
         // annotate_iff1
         { pm -> m; pv -> (v1 + v2)}
         // put k3 (v1 + v2)
         { pm -> map.put m k3 (v1 + v2); match map.get m k3 with
                                         | Some v3 => { pv -> v3 }
                                         | None => emp True
                                         end }
     *)
    Lemma put_sum_ok : program_logic_goal_for_function! put_sum.
    Proof.
      repeat straightline'. cbv [put_sum_gallina].
      add_map_annotations.

      repeat match goal with
             | _ => progress kv_hammer
             | |- WeakestPrecondition.call _ ?f _ ?m ?args _ =>
               (* call add -- need to borrow all args first *)
               unify f "add";
                 let in0 := (eval hnf in (hd (word.of_Z 0) args)) in
                 let in1 := (eval hnf in (hd (word.of_Z 0) (tl args))) in
                 let in2 :=
                     (eval hnf in (hd (word.of_Z 0) (tl (tl args)))) in
                 try borrow in0; try borrow in1; try borrow in2;
                   handle_call; autorewrite with mapsimpl in *
             end.

      all: remove_map_annotations.

      all: subst; ssplit; try reflexivity.
      all: ecancel_assumption.
    Qed.
  End put_sum.

  Section swap.
    Context {ops} {key value : Type} {Value}
            {kvp : kv_parameters}
            {ok : @kv_parameters_ok _ BW _ mem ops key value Value kvp}.

    Existing Instances map_ok annotated_map_ok key_eq_dec.
    Local Hint Extern 1 (spec_of "get") => unshelve simple refine (@spec_of_map_get _ _ _ _ _ _ _ _ _ _ _) : typeclass_instances.
    Local Hint Extern 1 (spec_of "put") => unshelve simple refine (@spec_of_map_put _ _ _ _ _ _ _ _ _ _ _) : typeclass_instances.

    (* look up k1 and k2, add their values and store in k3 *)
    Definition swap_gallina (m : map.rep (map:=map))
               (k1 k2 : key) : map.rep (map:=map) :=
      match map.get m k1, map.get m k2 with
      | Some v1, Some v2 =>
        map.put (map.put m k2 v1) k1 v2
      | _, _ => m
      end.

    Definition swap := func! (m, k1, k2) {
      unpack! err, v1 = get (m, k1) ;
      require !err ;
      unpack! err, v2 = get (m, k2) ;
      require !err ;
      unpack! err = put (m, k2, v1)
      (* now v2 is stored in v1 -- no need for second put *)
    }.

    Instance spec_of_swap : spec_of "swap" :=
      fun functions =>
        forall pm m pk1 k1 pk2 k2 R tr mem,
          map.get m k1 <> None ->
          map.get m k2 <> None ->
          k1 <> k2 -> (* TODO: try removing *)
          (Map pm m * Key pk1 k1 * Key pk2 k2 * R)%sep mem ->
          WeakestPrecondition.call
            functions "swap" tr mem [pm; pk1; pk2]
            (fun tr' mem' rets =>
               tr = tr'
               /\ rets = []
               /\ (Map pm (swap_gallina m k1 k2)
                   * Key pk1 k1 * Key pk2 k2 * R)%sep mem').

    (* Entire chain of separation-logic reasoning for swap:

         { pm -> m; pk1 -> k1; pk2 -> k2 }
         // annotate
         { pm -> annotate m; pk1 -> k1; pk2 -> k2 }
         // get k1
         { pm -> map.put (annotate m) k1 (Reserved pv1, v1);
              pk1 -> k1; pk2 -> k2 }
         // get k2
         { pm -> map.put (map.put (annotate m) k1 (Reserved pv1, v1))
                         k2 (Reserved pv2, v2); pk1 -> k1; pk2 -> k2 }
         // borrow pv1
         { pm -> map.put (map.put (annotate m) k2 (Reserved pv2, v2))
                         k1 (Borrowed pv1, v1); pv1 -> v1;
                         pk1 -> k1; pk2 -> k2 }
         // put k2
         { pm -> map.put
                   (map.put (map.put (annotate m) k2 (Reserved pv2, v2))
                         k1 (Borrowed pv1, v1)) k2 (Reserved pv2, v1);
             pv1 -> v2; pk1 -> k1; pk2 -> k2 }
         // put_put
         { pm -> map.put (map.put (annotate m) k1 (Borrowed pv1, v1))
                         k2 (Reserved pv2, v1);
             pv1 -> v2; pk1 -> k1; pk2 -> k2 }
        // unborrow pv1
         { pm -> map.put (map.put (annotate m) k2 (Reserved pv2, v1))
                         k1 (Reserved pv1, v2);
             pk1 -> k1; pk2 -> k2 }
     *)
    Lemma swap_ok : program_logic_goal_for_function! swap.
    Proof.
      repeat straightline. cbv [swap_gallina].
      add_map_annotations.

      kv_hammer.

      remove_map_annotations.
      rewrite map.put_put_diff_comm by congruence.
      subst; ssplit; try reflexivity.
      ecancel_assumption.
    Qed.
  End swap.
End examples.


(* begin hide *)
Require Import Verse.Ast.
Require Import Verse.TypeSystem.
Require Import Verse.Language.Types.
Require Import Verse.Scope.
Require Import Verse.Utils.hlist.
Require Import Verse.Abstract.Machine.
Require Import Verse.Monoid.

Require Import PeanoNat.

Require Import EqdepFacts.
Import EqNotations.

(* end hide *)

(**

Verification of verse code is carried out by annotating code with
verification conditions. To begin the process, we need to interpret
verse types as types in Gallina using a type denotation. Given such a
denotation, the PHOAS style verse code is interpreted as interpreted
as instructions to an abstract machine whose states are captured by
hlists as given by the module [Verse.Abstract.Machine].  The hlist
machine is purely defined by the scope of the code that we are
interpreting.  The memory cells then act as the variables in this
context and unlike the generic variable, they can be checked for
equality (hetrogeneous though) and the associated state transitions
are computable.

The "universal" nature of this machine means that this is a natural
semantics to give to our code and hence verification here is
meaningful.

 *)

Section Hlist.

  Context [ts : typeSystem] (sc : Scope.type ts).
  Variable tyD : typeDenote ts.

  (** The memory of the hlist machine is the state *)
  Definition state    := Machine.memory (tyD : Variables.U ts) sc.

  (** The associated variable is just addresses into that memory *)
  Definition variable : Variables.U ts := Machine.address sc.

  (** The function that specialises the generic variable in the scoped
      code with the above variable type *)

  Definition specialise [CODE : Variables.U ts -> Type]
             (genC : forall v, scoped v sc (CODE v)) : CODE variable :=
    uncurry (genC variable) (all_membership sc).

  Local Definition type := typeOf ts.

  (** Compute the value of a variable *)
  Definition val (m : state) ty (x : variable ty)
    := index x m.

  (** Update the value of a variable *)
  Definition update ty (x : variable ty) f (m : state)
    := update x m (f (index x m)).

  (** The state transformation *)
  Definition transformation := state -> state.

  (** An assertion is a predicate on state. Here we have a product of
      state because we would like to assert things based on the
      initial value of the state. The first component of the product
      is therefore the initial state.
   *)
  Definition assertion := state * state -> Prop.

  (** The transformations form a monoid and acts on assertions in a
      natural way by transforming the second component state
      pair. Recall that the first component is the initial state and
      hence transformation acts trivially on it *)

  Global Instance assertion_action_op : LActionOp transformation assertion :=
    fun tr asrt => fun oAn => asrt (fst oAn, tr (snd oAn)).

  Add Parametric Morphism : lact with signature
      SetoidClass.equiv (A:=transformation) ==> SetoidClass.equiv (A:=assertion) ==> SetoidClass.equiv (A:=assertion) as lact_morp.
    unfold lact.
    unfold assertion_action_op; simpl.
    intros P Q.
    intro  PEQ.
    intros tr1 tr2.
    intro trEq.
    intro_destruct.
    rewrite (trEq (s, P s0)).
    now rewrite PEQ.
  Qed.

  Global Program Instance assertion_action : LAction transformation assertion :=
    {| lact_unit := _  |}.

  Next Obligation.
    intro_destruct; unfold lact. unfold assertion_action_op. intuition.
  Qed.

  Next Obligation.
    intro_destruct. unfold lact. unfold assertion_action_op. intuition.
  Qed.

  Next Obligation.
    intro_destruct; unfold lact; unfold assertion_action_op. intuition.
  Qed.

  Next Obligation.
    intro_destruct; unfold lact; unfold assertion_action_op;
    unfold binop; unfold transition_binop; unfold Basics.compose; simpl. intuition.
  Qed.

End Hlist.

Arguments val [ts sc tyD] _ [ty].
Arguments update [ts sc tyD ty].

(** * Meanings of program statements as machine transformations *)
Module Internals.

  Section MachineType.
    Variable ts : typeSystem.
    Variable sc : Scope.type ts.
    Variable tyD : typeDenote ts.

    Definition expr  T := Ast.expr  (variable sc) T.
    Definition lexpr T := Ast.lexpr (variable sc) T.

    Definition evalE {T} (st : state sc tyD) (e : expr T) :  tyD _ (projT2 T)
      := Ast.Expr.eval (Expr.rename (val st) e).

    Definition assign {T} (l : lexpr T) (e : expr T)(st : state sc tyD)
      : state sc tyD :=
      match l in Ast.lexpr _ T0 return  tyD _ (projT2 T0) -> state _ _ with
      | Ast.var reg  => fun v => update reg (fun _ => v) st
      | Ast.deref a idx => fun v => let arr := val st a in
                                let arrp := Vector.replace_order
                                              (rew [fun T0 : Type@{abstract_type_system.u0} => T0]
                                                   arrayCompatibility tyD _ _ _
                                                in arr)
                                              (proj2_sig idx) v
                                in
                                update a (fun _ =>
                                            (rew <- [id] arrayCompatibility tyD _ _ _ in arrp)
                                         ) st
      end (evalE st e).

    Definition binopUpdate {T}
               (l : lexpr (existT _ _ T))
               (o : operator ts T 2)
               (e : expr (existT _ _ T))
      : state sc tyD -> state sc tyD
      := let rhs := Ast.binOp o (Ast.valueOf l) e in
         assign l rhs.

    Definition uniopUpdate {T}
               (l : lexpr (existT _ _ T))
               (o : operator ts T 1)
      : state sc tyD -> state sc tyD
      := let rhs := Ast.uniOp o (Ast.valueOf l) in
         assign l rhs.

    Program Definition move {T}(l : lexpr T)(r : variable sc T) :=
      match l in Ast.lexpr _ T0 return variable sc T0 -> _ with
      | var _ as l1     => fun r0 => assign l1 (Ast.valueOf (Ast.var r0))
      | deref _ _ as l1 => fun r0 => assign l1 (Ast.valueOf (Ast.var r0))
      end r.

    Definition denoteInst {T}(inst : Ast.instruction (variable sc) T) : transformation sc tyD
      := match inst with
         | Ast.assign l e => assign l e
         | Ast.moveTo l r => move   l r
         | Ast.binopUpdate l o e => binopUpdate l o e
         | Ast.uniopUpdate l o   => uniopUpdate l o
         | Ast.clobber _     => fun x => x
         end.

    Definition denoteStmt (stmt : Ast.statement (variable sc) ) : transformation sc tyD
      := denoteInst (projT2 stmt).

  End MachineType.

End Internals.

(** ** The verification monoid.

Verification of annotated code involves working with both
transformations as well as assertions. Transformations form a monoid
under composition and assertions form a monoid under the /\ operation.
It turns out that the verification process can also be seen as working
with a monoid namely the semi direct product got by the action of the
transformation on the assertions.

*)
Section Machine.

  Variable ts : typeSystem.
  Variable sc : Scope.type ts.

  Variable tyD   : typeDenote ts.

  Definition mline := transformation sc tyD ⋉ assertion sc tyD.

  Definition justInst
    : transformation sc tyD -> mline
    := fun i => semiR i (fun _ => True).

  Definition justAssert
    : assertion sc tyD -> mline
    := fun a => semiR id a.

End Machine.

Arguments mline [ts].
Arguments justInst [ts sc tyD] .
Arguments justAssert [ts sc tyD].


From Hammer Require Import Hammer.

Require Export ClassicalChoice.
Require Export ExtensionalFunctionRepresentative.

Require Import ChoiceFacts.
Require Import ClassicalFacts.
Require Import RelationClasses.

Theorem setoid_choice :
forall A B,
forall R : A -> A -> Prop,
forall T : A -> B -> Prop,
Equivalence R ->
(forall x x' y, R x x' -> T x y -> T x' y) ->
(forall x, exists y, T x y) ->
exists f : A -> B, forall x : A, T x (f x) /\ (forall x' : A, R x x' -> f x = f x').
Proof. hammer_hook "SetoidChoice" "SetoidChoice.setoid_choice".
apply setoid_functional_choice_first_characterization. split; [|split].
- exact choice.
- exact extensional_function_representative.
- exact classic.
Qed.

Theorem representative_choice :
forall A (R:A->A->Prop), (Equivalence R) ->
exists f : A->A, forall x : A, R x (f x) /\ forall x', R x x' -> f x = f x'.
Proof. hammer_hook "SetoidChoice" "SetoidChoice.representative_choice".
apply setoid_fun_choice_imp_repr_fun_choice.
exact setoid_choice.
Qed.


Require Import AutoSep.

(** * Always-0 *)

Definition always0S res : spec := SPEC reserving res
  PRE[_] Emp
  POST[R] [| R = 0 |].

Definition always0 := bmodule "always0" {{
  bfunction "main"("f") [always0S 2]
    "f" <-- Lambda() [always0S 0]
      Return 0
    end;;

    "f" <-- ICall "f"()
    [PRE[_, R] [| R = 0 |]
     POST[R'] [| R' = 0 |] ];;
    Return "f"
  end
}}.

Hint Extern 1 (@eq W _ _) => words.
Hint Extern 1 (@eq (word _) _ _) => words.

Theorem always0Ok : moduleOk always0.
  vcgen; abstract (post; try icall (@nil string); (sep_auto; auto)).
Qed.

(** * Add three numbers *)

Definition add3S : spec := SPEC("x", "y", "z") reserving 5
  PRE[V] Emp
  POST[R] [| R = V "x" ^+ V "y" ^+ V "z" |].

Definition addS : spec := SPEC("x", "y") reserving 0
  PRE[V] Emp
  POST[R] [| R = V "x" ^+ V "y" |].

Definition add3 := bmodule "add3" {{
  bfunction "add3"("x", "y", "z", "f", "tmp") [add3S]
    "f" <-- Lambda("x", "y") [addS]
      Return "x" + "y"
    end;;

    "tmp" <-- ICall "f"("x", "y")
    [PRE[V, R] Emp
     POST[R'] [| R' = R ^+ V "z" |]];;

    Return "tmp" + "z"
  end
}}.

Theorem add3Ok : moduleOk add3.
  vcgen; abstract (post; try icall ("x" :: "y" :: nil); (sep_auto; auto)).
Qed.