From sflib Require Import sflib.
From Paco Require Import paco.

Require Export Coq.Strings.String.
From Coq Require Import Program.

From Fairness Require Export ITreeLib EventCategory FairBeh NatStructs.
From Fairness Require Export Mod.

Set Implicit Arguments.

Lemma unfold_iter E A B (f: A -> itree E (A + B)) (x: A)
  :
  ITree.iter f x =
    lr <- f x;;
    match lr with
    | inl l => tau;; ITree.iter f l
    | inr r => Ret r
    end.
Proof.
  apply bisim_is_eq. eapply unfold_iter.
Qed.

Module Th := NatMap.
Notation thread _Id E R := (itree ((((@eventE _Id) +' cE) +' callE) +' E) R).
Notation threads _Id E R := (Th.t (@thread _Id E R)).

Definition fn2th (m: Mod.t) (fn: fname) (args: Any.t): @thread (Mod.ident m) (sE (Mod.state m)) Any.t :=
  match Mod.funs m fn with
  | Some ktr => ktr args
  | None => Vis (inl1 (inl1 (inl1 Undefined))) (Empty_set_rect _)
  end.

Fixpoint _numbering {E} (l: list E) (n: NatMap.key): list (NatMap.key * E) :=
  match l with
  | hd :: tl => (n, hd) :: (_numbering tl (S n))
  | [] => []
  end.

Definition numbering {E} (l: list E): list (NatMap.key * E) := _numbering l O.

Definition prog2ths (m: Mod.t) (p: program): @threads (Mod.ident m) (sE (Mod.state m)) Any.t :=
  let pre_threads := List.map (fun '(fn, args) => fn2th m fn args) p in
  NatMapP.of_list (numbering pre_threads).

Section STATE.

  Variable S: Type.
  Variable E: Type -> Type.

  Let Es := E +' (sE S).

  Definition interp_state_aux {R} :
    (S * (itree Es R)) -> itree E (S * R).
  Proof.
    eapply ITree.iter. intros [state itr]. destruct (observe itr).
    - exact (Ret (inr (state, r))).
    - exact (Ret (inl (state, t))).
    - destruct e.
      + exact (Vis e (fun x => Ret (inl (state, k x)))).
      + destruct s as [f].
        exact (Ret (inl (fst (f state), k (snd (f state))))).
  Defined.

  Definition interp_state {R}:
    (S * (itree Es R)) -> itree E R :=
    fun x => r <- interp_state_aux x;;
          Ret (snd r).

  Lemma interp_state_aux_ret R st (r : R) :
    interp_state_aux (st, Ret r) = Ret (st, r).
  Proof. unfold interp_state_aux. rewrite unfold_iter. grind. Qed.

  Lemma interp_state_aux_vis R st X (e : E X) (ktr : ktree Es X R) :
    interp_state_aux (st, Vis (inl1 e) ktr) = Vis e (fun x => tau;; interp_state_aux (st, ktr x)).
  Proof. unfold interp_state_aux. rewrite unfold_iter. grind.
         apply observe_eta. ss. f_equal. extensionality x. grind.
  Qed.

  Lemma interp_state_aux_rmw R st X rmw (ktr : ktree Es X R) :
    interp_state_aux (st, Vis (inr1 (Rmw rmw)) ktr) =
      tau;; interp_state_aux (fst (rmw st), ktr (snd (rmw st))).
  Proof. unfold interp_state_aux. rewrite unfold_iter. grind. Qed.

  Lemma interp_state_aux_tau R st (itr : itree Es R) :
    interp_state_aux (st, Tau itr) = Tau (interp_state_aux (st, itr)).
  Proof. unfold interp_state_aux. rewrite unfold_iter. grind. Qed.

  Lemma interp_state_aux_bind {R1} {R2} st (itr : itree Es R1) (ktr : ktree Es R1 R2) :
    interp_state_aux (st, itr >>= ktr) =
      x <- interp_state_aux (st, itr);;
      interp_state_aux (fst x, ktr (snd x)).
  Proof.
    eapply bisim_is_eq. revert st itr. pcofix CIH. i.
    destruct_itree itr.
    - rewrite interp_state_aux_ret. grind.
      eapply paco2_mon.
      + eapply eq_is_bisim. ss.
      + ss.
    - grind. rewrite 2 interp_state_aux_tau. grind.
      pfold. econs. right. eapply CIH.
    - grind. destruct e.
      + rewrite 2 interp_state_aux_vis. grind.
        pfold. econs. intros. grind. left.
        pfold. econs. right. eapply CIH.
      + destruct s. rewrite ! interp_state_aux_rmw. grind.
        pfold. econs. right. eapply CIH.
  Qed.

  Lemma unfold_interp_state R st (itr : itree Es R) :
    interp_state (st, itr) = x <- interp_state_aux (st, itr);; Ret (snd x).
  Proof. unfold interp_state. ss. Qed.

  Lemma interp_state_bind R1 R2 st (itr : itree Es R1) (ktr : ktree Es R1 R2) :
    interp_state (st, itr >>= ktr) = x <- interp_state_aux (st, itr);; interp_state (fst x, ktr (snd x)).
  Proof. unfold interp_state. rewrite interp_state_aux_bind. grind. Qed.

  Lemma interp_state_ret
        R (r: R) (state: S)
    :
    interp_state (state, Ret r) = Ret r.
  Proof.
    unfold interp_state, interp_state_aux. rewrite unfold_iter. ss. grind.
  Qed.

  Lemma interp_state_tau
        R (state: S) (itr: itree Es R)
    :
    interp_state (state, tau;; itr) = tau;; interp_state (state, itr).
  Proof.
    unfold interp_state, interp_state_aux. rewrite unfold_iter. ss. grind.
  Qed.

  Lemma interp_state_rmw_vis
    R st X rmw (ktr : ktree Es X R)
    : interp_state (st, Vis (|Rmw rmw)%sum ktr) =
        tau;; interp_state (fst (rmw st), ktr (snd (rmw st))).
  Proof. unfold interp_state. rewrite interp_state_aux_rmw. grind. Qed.

  Lemma interp_state_rmw
    R st X rmw (ktr : ktree Es X R)
    : interp_state (st, trigger (Rmw rmw) >>= ktr) =
        tau;; interp_state (fst (rmw st), ktr (snd (rmw st))).
  Proof. rewrite bind_trigger. eapply interp_state_rmw_vis. Qed.

  Lemma interp_state_get_vis
        R (state: S) (ktr: S -> itree Es R)
    :
    interp_state (state, Vis (inr1 (Get id)) ktr)
    =
      tau;; interp_state (state, ktr state).
  Proof.
    unfold interp_state, interp_state_aux. rewrite get_rmw, unfold_iter. ss. grind.
  Qed.

  Lemma interp_state_get
        R (state: S) (ktr: S -> itree Es R)
    :
    interp_state (state, trigger (inr1 (Get id)) >>= ktr)
    =
      tau;; interp_state (state, ktr state).
  Proof.
    rewrite bind_trigger. apply interp_state_get_vis.
  Qed.

  Lemma interp_state_vis
        R (state: S) X (e: E X) (ktr: X -> itree Es R)
    :
    interp_state (state, Vis (inl1 e) ktr)
    =
      Vis e (fun x => tau;; interp_state (state, ktr x)).
  Proof.
    unfold interp_state, interp_state_aux. rewrite unfold_iter. ss. rewrite 2 bind_vis.
    apply observe_eta. ss. f_equal. extensionality x.
    rewrite bind_ret_l. rewrite bind_tau. reflexivity.
  Qed.

  Lemma interp_state_trigger
        R (state: S) X (e: E X) (ktr: X -> itree Es R)
    :
    interp_state (state, trigger (inl1 e) >>= ktr)
    =
      x <- trigger e;; tau;; interp_state (state, ktr x).
  Proof.
    rewrite ! bind_trigger. apply interp_state_vis.
  Qed.

End STATE.
Global Opaque interp_state_aux interp_state.

Section STATE_PROP.

  Variable State: Type.

  Lemma interp_state_aux_map_event E1 E2 R (embed : forall X, E1 X -> E2 X) st (itr : itree (E1 +' sE State) R) :
    interp_state_aux (st, map_event (embed_left embed) itr) = map_event embed (interp_state_aux (st, itr)).
  Proof.
    eapply bisim_is_eq. revert st itr. pcofix CIH. i.
    destruct_itree itr.
    - rewrite map_event_ret.
      rewrite 2 interp_state_aux_ret.
      rewrite map_event_ret.
      pfold. econs. ss.
    - rewrite map_event_tau.
      rewrite 2 interp_state_aux_tau.
      rewrite map_event_tau.
      pfold. econs. right. eapply CIH.
    - rewrite map_event_vis.
      destruct e.
      + ss. rewrite 2 interp_state_aux_vis.
        rewrite map_event_vis.
        pfold. econs. intros. left.
        rewrite map_event_tau.
        pfold. econs. right. eapply CIH.
      + ss. destruct s.
        rewrite ! interp_state_aux_rmw.
        rewrite ! map_event_tau.
        pfold. econs. right. eapply CIH.
  Qed.

End STATE_PROP.

Section SCHEDULE.

  Variant schedulerE (RT : Type) : Type -> Type :=
  | Execute : thread_id -> schedulerE RT (option RT)
  .

  Let eventE0 := @eventE thread_id.

  Definition scheduler RT R := itree (schedulerE RT +' eventE0) R.

  Context {_Ident: ID}.
  Variable E: Type -> Type.

  Let eventE1 := @eventE _Ident.
  Let eventE2 := @eventE (sum_tid _Ident).
  Let Es0 := ((eventE1 +' cE) +' callE) +' E.
  Let thread R := thread _Ident E R.
  Let threads R := threads _Ident E R.

  Definition interp_thread_aux {R} :
    thread_id * thread R -> itree (eventE1 +' E) (thread R + R).
  Proof.
    eapply ITree.iter.
    intros [tid itr].
    apply observe in itr; destruct itr as [r | itr | X e k].
    - (* Ret *)
      exact (Ret (inr (inr r))).
    - (* Tau *)
      exact (Ret (inl (tid, itr))).
    - (* Vis *)
      destruct e as [[[]|]|].
      + (* eventE *)
        exact (Vis (inl1 e) (fun x => Ret (inl (tid, k x)))).
      + (* cE *)
        destruct c.
        * (* Yield *)
          exact (Ret (inr (inl (k tt)))).
        * (* GetTid *)
          exact (Ret (inl (tid, k tid))).
      + (* callE *)
        exact (Vis (inl1 Undefined) (Empty_set_rect _)).
      + (* E *)
        exact (Vis (inr1 e) (fun x => Ret (inl (tid, k x)))).
  Defined.

  Definition interp_thread {R} :
    thread_id * thread R -> itree (eventE2 +' E) (thread R + R) :=
    fun x => map_event (embed_left (map_prism inrp)) (interp_thread_aux x).

  Definition interp_sched RT R : threads RT * scheduler RT R -> itree (eventE2 +' E) R.
  Proof.
    eapply ITree.iter. intros [ts sch].
    destruct (observe sch) as [r | sch' | X [e|e] ktr].
    - exact (Ret (inr r)).
    - exact (Ret (inl (ts, sch'))).
    - destruct e.
      destruct (Th.find n ts) as [t|].
      * exact (r <- interp_thread (n, t);;
               match r with
               | inl t' => Ret (inl (Th.add n t' ts, ktr None))
               | inr r => Ret (inl (Th.remove n ts, ktr (Some r)))
               end).
      * exact (Vis (inl1 Undefined) (Empty_set_rect _)).
    - exact (Vis (inl1 (map_prism inlp e)) (fun x => Ret (inl (ts, ktr x)))).
  Defined.

  Lemma unfold_interp_thread {R} tid (itr : thread R) :
    interp_thread (tid, itr) = map_event (embed_left (map_prism inrp)) (interp_thread_aux (tid, itr)).
  Proof. ss. Qed.

  Lemma interp_thread_ret {R} tid (r : R) :
    interp_thread (tid, Ret r) = Ret (inr r).
  Proof. unfold interp_thread, interp_thread_aux. rewrite unfold_iter. grind. eapply map_event_ret. Qed.

  Lemma interp_thread_tau R tid (itr : thread R) :
    interp_thread (tid, tau;; itr) = tau;; interp_thread (tid, itr).
  Proof. unfold interp_thread at 1, interp_thread_aux. rewrite unfold_iter. grind. eapply map_event_tau. Qed.

  Lemma interp_thread_vis R tid X (e : E X) (ktr : ktree Es0 X R) :
    interp_thread (tid, Vis (inr1 e) ktr) =
      Vis (inr1 e) (fun x => tau;; interp_thread (tid, ktr x)).
  Proof.
    unfold interp_thread at 1, interp_thread_aux. rewrite unfold_iter. grind. rewrite map_event_vis.
    eapply (f_equal (fun x => Vis (inr1 e) x)). extensionality x. grind. rewrite map_event_tau. grind.
  Qed.

  Lemma interp_thread_trigger R tid X (e : E X) (ktr : ktree Es0 X R) :
    interp_thread (tid, x <- trigger (inr1 e);; ktr x) =
      x <- trigger (inr1 e);; tau;; interp_thread (tid, ktr x).
  Proof. rewrite ! bind_trigger. eapply interp_thread_vis. Qed.

  Lemma interp_thread_vis_eventE R tid X (e : eventE1 X) (ktr : ktree Es0 X R) :
    interp_thread (tid, Vis (inl1 (inl1 (inl1 e))) ktr) =
      Vis (inl1 (map_prism inrp e)) (fun x => tau;; interp_thread (tid, ktr x)).
  Proof.
    unfold interp_thread at 1, interp_thread_aux. rewrite unfold_iter. grind. rewrite map_event_vis.
    eapply (f_equal (fun x => Vis (inl1 (map_prism inrp e)) x)). extensionality x. grind. rewrite map_event_tau. grind.
  Qed.

  Lemma interp_thread_trigger_eventE R tid X (e : eventE1 X) (ktr : ktree Es0 X R) :
    interp_thread (tid, x <- trigger (inl1 (inl1 e));; ktr x) =
      x <- trigger (inl1 (map_prism inrp e));; tau;; interp_thread (tid, ktr x).
  Proof. rewrite ! bind_trigger. eapply interp_thread_vis_eventE. Qed.

  Lemma interp_thread_vis_gettid R tid (ktr : ktree Es0 thread_id R) :
    interp_thread (tid, Vis (inl1 (inl1 (inr1 GetTid))) ktr) =
      tau;; interp_thread (tid, ktr tid).
  Proof.
    unfold interp_thread at 1, interp_thread_aux. rewrite unfold_iter. grind. rewrite map_event_tau. grind.
  Qed.

  Lemma interp_thread_trigger_gettid R tid (ktr : ktree Es0 thread_id R) :
    interp_thread (tid, x <- trigger (inl1 (inr1 GetTid));; ktr x) =
      tau;; interp_thread (tid, ktr tid).
  Proof. rewrite bind_trigger. eapply interp_thread_vis_gettid. Qed.

  Lemma interp_thread_vis_yield R tid (ktr : ktree Es0 () R) :
    interp_thread (tid, Vis (inl1 (inl1 (inr1 Yield))) ktr) =
      Ret (inl (ktr tt)).
  Proof.
    unfold interp_thread at 1, interp_thread_aux. rewrite unfold_iter. grind. rewrite map_event_ret. ss.
  Qed.

  Lemma interp_thread_trigger_yield R tid (ktr : ktree Es0 () R) :
    interp_thread (tid, x <- trigger (inl1 (inr1 Yield));; ktr x) =
      Ret (inl (ktr tt)).
  Proof. rewrite bind_trigger. apply interp_thread_vis_yield. Qed.

  Lemma interp_thread_call R tid fn args (ktr : ktree Es0 Any.t R) :
    interp_thread (tid, trigger (Call fn args) >>= ktr) = Vis (inl1 Undefined) (Empty_set_rect _).
  Proof. unfold interp_thread at 1, interp_thread_aux. rewrite unfold_iter. grind.
         rewrite map_event_vis. eapply observe_eta. ss. f_equal. extensionalities x. ss.
  Qed.

  Lemma interp_sched_ret RT R (ths : threads RT) (r : R) :
    interp_sched (ths, Ret r) = Ret r.
  Proof. unfold interp_sched. rewrite unfold_iter. grind. Qed.

  Lemma interp_sched_tau RT R ths (itr : scheduler RT R) :
    interp_sched (ths, Tau itr) = Tau (interp_sched (ths, itr)).
  Proof. unfold interp_sched. rewrite unfold_iter. grind. Qed.

  Lemma interp_sched_execute_Some RT R ths tid t (ktr : option RT -> scheduler RT R)
    (SOME : Th.find tid ths = Some t)
    : interp_sched (ths, Vis (inl1 (Execute _ tid)) ktr) =
      r <- interp_thread (tid, t);;
      match r with
      | inl t' => tau;; interp_sched (Th.add tid t' ths, ktr None)
      | inr r => tau;; interp_sched (Th.remove tid ths, ktr (Some r))
      end.
  Proof. unfold interp_sched. rewrite unfold_iter. grind. Qed.

  Lemma interp_sched_execute_None RT R ths tid (ktr : option RT -> scheduler RT R)
    (NONE : Th.find tid ths = None)
    : interp_sched (ths, Vis (inl1 (Execute _ tid)) ktr) =
        Vis (inl1 Undefined) (Empty_set_rect _).
  Proof. unfold interp_sched. rewrite unfold_iter. grind.
         eapply observe_eta. ss. f_equal. extensionality x. ss.
  Qed.

  Lemma interp_sched_vis RT R ths X (e : eventE0 X) (ktr : X -> scheduler RT R) :
    interp_sched (ths, Vis (inr1 e) ktr) =
      Vis (inl1 (map_prism inlp e)) (fun x => tau;; interp_sched (ths, ktr x)).
  Proof. unfold interp_sched. rewrite unfold_iter. grind.
         eapply observe_eta. ss. f_equal. extensionality x. grind.
  Qed.

End SCHEDULE.
Global Opaque interp_thread_aux interp_thread interp_sched.

Section SCHEDULE_NONDET.

  Definition sched_nondet_body {R} q tid r : scheduler R (thread_id * TIdSet.t + R) :=
    match r with
    | None =>
        tid' <- ITree.trigger (inr1 (Choose thread_id));;
        match nm_pop tid' (NatSet.add tid q) with
        | None => Vis (inr1 (Choose void)) (Empty_set_rect _)
        | Some (_, q') =>
            ITree.trigger (inr1 (Fair (tids_fmap tid' q')));;;
            Ret (inl (tid', q'))
        end
    | Some r =>
        if NatMap.is_empty q
        then Ret (inr r)
        else
          tid' <- ITree.trigger (inr1 (Choose thread_id));;
          match nm_pop tid' q with
          | None => Vis (inr1 (Choose void)) (Empty_set_rect _)
          | Some (_, q') =>
              ITree.trigger (inr1 (Fair (tids_fmap tid' q')));;;
              Ret (inl (tid', q'))
          end
    end.

  Definition sched_nondet R0 : thread_id * TIdSet.t -> scheduler R0 R0 :=
    ITree.iter (fun '(tid, q) =>
                  r <- ITree.trigger (inl1 (Execute _ tid));;
                  sched_nondet_body q tid r).

  Lemma unfold_sched_nondet R0 tid q :
    sched_nondet R0 (tid, q) =
      r <- ITree.trigger (inl1 (Execute _ tid));;
      match r with
      | None =>
          tid' <- ITree.trigger (inr1 (Choose thread_id));;
          match nm_pop tid' (NatSet.add tid q) with
          | None => Vis (inr1 (Choose void)) (Empty_set_rect _)
          | Some (_, q') =>
              ITree.trigger (inr1 (Fair (tids_fmap tid' q')));;;
              tau;; sched_nondet _ (tid', q')
          end
      | Some r =>
          if NatMap.is_empty q
          then Ret r
          else
            tid' <- ITree.trigger (inr1 (Choose thread_id));;
            match nm_pop tid' q with
            | None => Vis (inr1 (Choose void)) (Empty_set_rect _)
            | Some (_, q') =>
                ITree.trigger (inr1 (Fair (tids_fmap tid' q')));;;
                tau;; sched_nondet _ (tid', q')
            end
      end.
  Proof.
    unfold sched_nondet at 1, sched_nondet_body.
    rewrite unfold_iter.
    grind.
    - eapply observe_eta. ss. f_equal. extensionality x0. ss.
    - eapply observe_eta. ss. f_equal. extensionality x0. ss.
  Qed.

  Context {_Ident : ID}.
  Variable E: Type -> Type.

  Let eventE1 := @eventE _Ident.
  Let eventE2 := @eventE (sum_tid _Ident).
  Let Es0 := (eventE1 +' cE) +' E.
  Let thread R := thread _Ident E R.
  Let threads R := threads _Ident E R.

  Lemma unfold_interp_sched_nondet_Some R tid t (ths : threads R) q :
    Th.find tid ths = Some t ->
    interp_sched (ths, sched_nondet R (tid, q)) =
      r <- interp_thread (tid, t);;
      match r with
      | inl t' => Tau (interp_sched (Th.add tid t' ths,
                                     tid' <- ITree.trigger (inr1 (Choose thread_id));;
                                     match nm_pop tid' (NatSet.add tid q) with
                                     | None => Vis (inr1 (Choose void)) (Empty_set_rect _)
                                     | Some (_, q') =>
                                         ITree.trigger (inr1 (Fair (tids_fmap tid' q')));;;
                                         tau;; sched_nondet _ (tid', q')
                                     end))
      | inr r => Tau (interp_sched (Th.remove tid ths,
                                    if NatMap.is_empty q
                                    then Ret r
                                    else
                                      tid' <- ITree.trigger (inr1 (Choose thread_id));;
                                      match nm_pop tid' q with
                                      | None => Vis (inr1 (Choose void)) (Empty_set_rect _)
                                      | Some (_, q') =>
                                          ITree.trigger (inr1 (Fair (tids_fmap tid' q')));;;
                                          tau;; sched_nondet _ (tid', q')
                                      end))
      end.
  Proof.
    rewrite unfold_sched_nondet at 1.
    rewrite bind_trigger.
    eapply interp_sched_execute_Some.
  Qed.

  Lemma unfold_interp_sched_nondet_None R tid (ths : threads R) q :
    Th.find tid ths = None ->
    interp_sched (ths, sched_nondet R (tid, q)) =
      Vis (inl1 Undefined) (Empty_set_rect _).
  Proof.
    rewrite unfold_sched_nondet at 1.
    rewrite bind_trigger.
    eapply interp_sched_execute_None.
  Qed.

End SCHEDULE_NONDET.
Global Opaque sched_nondet_body sched_nondet.

Section INTERP.

  Variable State: Type.
  Variable _Ident: ID.
  Variable R: Type.

  Definition interp_all
    st (ths: @threads _Ident (sE State) R) tid : itree (@eventE (sum_tid _Ident)) R :=
    interp_state (st, interp_sched (ths, sched_nondet _ (tid, NatSet.remove tid (key_set ths)))).

  Lemma interp_all_tau
        st (ths: @threads _Ident (sE State) R) tid
        itr
    :
    (interp_all st (Th.add tid (Tau itr) ths) tid) = (Tau (interp_all st (Th.add tid itr ths) tid)).
  Proof.
    unfold interp_all. erewrite ! unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    rewrite interp_thread_tau. rewrite bind_tau. rewrite interp_state_tau.
    do 5 f_equal. extensionality x. destruct x.
    - rewrite ! nm_add_add_eq. rewrite ! key_set_pull_add_eq. auto.
    - erewrite 1 nm_rm_add_eq. rewrite ! key_set_pull_add_eq. eauto.
  Qed.

  Lemma interp_all_vis
        st (ths: @threads _Ident (sE State) R) tid
        X (e: @eventE _Ident X) ktr
    :
    (interp_all st (Th.add tid (Vis (((e|)|)|)%sum ktr) ths) tid) =
      (Vis (map_prism inrp e) (fun x => tau;; tau;; interp_all st (Th.add tid (ktr x) ths) tid)).
  Proof.
    unfold interp_all. erewrite ! unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    rewrite interp_thread_vis_eventE. rewrite bind_vis. rewrite interp_state_vis.
    do 2 f_equal. extensionality x. rewrite bind_tau. rewrite interp_state_tau.
    erewrite 1 unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    do 7 f_equal. extensionality r.
    destruct r.
    - rewrite ! nm_add_add_eq. rewrite ! key_set_pull_add_eq. auto.
    - erewrite 1 nm_rm_add_eq. rewrite ! key_set_pull_add_eq. eauto.
  Qed.

  Lemma interp_all_rmw
    st (ths : @threads _Ident (sE State) R) tid
    X rmw ktr
    : interp_all st (Th.add tid (Vis (|Rmw rmw)%sum ktr) ths) tid =
        tau;; tau;; interp_all (fst (rmw st)) (Th.add tid (ktr (snd (rmw st) : X)) ths) tid .
  Proof.
    unfold interp_all. erewrite ! unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    rewrite interp_thread_vis. rewrite bind_vis. rewrite interp_state_rmw_vis.
    rewrite bind_tau. rewrite interp_state_tau.
    repeat f_equal. extensionalities x.
    destruct x.
    - rewrite ! nm_add_add_eq. rewrite ! key_set_pull_add_eq. auto.
    - erewrite 1 nm_rm_add_eq. rewrite ! key_set_pull_add_eq. eauto.
  Qed.

  Lemma interp_all_get
        st (ths: @threads _Ident (sE State) R) tid
        ktr
    :
    (interp_all st (Th.add tid (Vis (|Get id)%sum ktr) ths) tid) =
      (tau;; tau;; interp_all st (Th.add tid (ktr st) ths) tid).
  Proof.
    unfold interp_all. erewrite ! unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    rewrite interp_thread_vis. rewrite bind_vis. rewrite interp_state_get_vis. rewrite bind_tau. rewrite interp_state_tau.
    repeat f_equal. extensionality x.
    destruct x.
    - rewrite ! nm_add_add_eq. rewrite ! key_set_pull_add_eq. auto.
    - erewrite 1 nm_rm_add_eq. rewrite ! key_set_pull_add_eq. eauto.
  Qed.

  Lemma interp_all_tid
        st (ths: @threads _Ident (sE State) R) tid
        ktr
    :
    (interp_all st (Th.add tid (Vis (((|GetTid)|)|)%sum ktr) ths) tid) =
      (tau;; interp_all st (Th.add tid (ktr tid) ths) tid).
  Proof.
    unfold interp_all. erewrite ! unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    rewrite interp_thread_vis_gettid. rewrite bind_tau. rewrite interp_state_tau.
    repeat f_equal. extensionality x.
    destruct x.
    - rewrite ! nm_add_add_eq. rewrite ! key_set_pull_add_eq. auto.
    - erewrite 1 nm_rm_add_eq. rewrite ! key_set_pull_add_eq. eauto.
  Qed.

  Lemma interp_all_call
    st (ths: @threads _Ident (sE State) R) tid
    fn args ktr
    : interp_all st (Th.add tid (trigger (Call fn args) >>= ktr) ths) tid = trigger Undefined >>= Empty_set_rect _.
  Proof.
    unfold interp_all. erewrite ! unfold_interp_sched_nondet_Some; eauto using nm_find_add_eq.
    rewrite interp_thread_call. rewrite bind_vis. rewrite interp_state_vis. rewrite ! bind_trigger.
    eapply observe_eta. ss. f_equal. extensionalities s. ss.
  Qed.

End INTERP.

Section MOD.

  Variable mod : Mod.t.
  Let st := (Mod.st_init mod).
  Let Ident := (Mod.ident mod).
  Let main := ((Mod.funs mod) "main").

  Definition interp_mod
             (tid: thread_id)
             (ths: @threads (Mod.ident mod) (sE (Mod.state mod)) Val)
             (sched: forall R, (thread_id * TIdSet.t)%type -> (scheduler R R)):
    itree (@eventE (sum_tid Ident)) Val :=
    interp_state (st, interp_sched (ths, sched _ (tid, NatSet.remove tid (key_set ths)))).

End MOD.


Require Import List.
Import ListNotations.

Require Import FSets FSets.FMapAVL FSets.FMapFacts.

From Verbatim Require Import state.
From Verbatim Require Import memo.
From Verbatim Require Import ltac.
From Verbatim Require Import Orders.
From Verbatim Require Import hashtrie.

Module FMemo (STT : state.T) <: MEMO STT.

  Import STT.Ty.
  Import STT.Defs.
  Import STT.R.Defs.
  Import Trie.

  Module Pointer_as_UOT <: UsualOrderedType := UOT_from_UCT Pointer_as_UCT.
  Module FM := FMapAVL.Make Pointer_as_UOT.
  Module FMF := FMapFacts.Facts FM.

  Definition Memo : Type := FM.t (@Trie (option (String * String * index))).
  Definition emptyMemo : Memo := FM.empty (@Trie (option (String * String * index))).

  Definition get_Memo (M : Memo) (pnt : Pointer) (i : index)
    : option (option (String * String * index)) :=
    match FM.find pnt M with
    | None => None
    | Some T => get_Trie T (index2list i)
    end.

  Definition set_Memo (M : Memo) (pnt : Pointer) (i : index)
             (o : (option (String * String * index))) : Memo :=
    match FM.find pnt M with
    | None => FM.add pnt (set_Trie Leaf (index2list i) o) M
    | Some T => FM.add pnt (set_Trie T (index2list i) o) M
    end.

  Lemma correct_Memo : forall M ptr i o, get_Memo (set_Memo M ptr i o) ptr i = Some o.
  Proof.
    intros. unfold get_Memo. unfold set_Memo. repeat dm.
    - rewrite FMF.add_eq_o in E; auto. repeat inj_all. apply get_set.
    - rewrite FMF.add_eq_o in E; auto. repeat inj_all. apply get_set.
    - rewrite FMF.add_eq_o in E; auto. discriminate.
    - rewrite FMF.add_eq_o in E; auto. discriminate.
  Qed.

  Lemma correct_Memo_moot : forall M ptr ptr' i i' o,
      (ptr <> ptr' \/ i <> i')
      ->
      get_Memo (set_Memo M ptr' i' o) ptr i = get_Memo M ptr i.
  Proof.
    intros. unfold get_Memo. unfold set_Memo.
    destruct (Pointer_as_UOT.eq_dec ptr ptr') eqn:E;
      destruct (index_eq_dec i i'); destruct H;
      repeat dm;
      try(rewrite FMF.add_neq_o in *; auto; rewrite E in *; repeat inj_all; auto; discriminate);
      try(rewrite FMF.add_eq_o in *; auto;
          repeat inj_all; rewrite E1 in *; repeat inj_all; rewrite get_set_moot; auto;
          intros C; destruct H; apply f_equal with (f := list2index) in C;
          repeat rewrite list_inv in C; auto; discriminate);
      try(try subst i; rewrite FMF.add_neq_o in *; auto; rewrite E0 in *;
          repeat inj_all; auto; discriminate);
      try(subst ptr; rewrite E1 in *; discriminate).
    - rewrite FMF.add_eq_o in *; auto; discriminate.
  Qed.

  Lemma correct_emptyMemo : forall stt z, get_Memo emptyMemo stt z = None.
  Proof.
    intros. unfold get_Memo. unfold emptyMemo. repeat dm.
    rewrite FMF.empty_o in E. discriminate.
  Qed.

End FMemo.

Module memoTFn (STT' : state.T) <: memo.T.
  Module STT := STT'.
  Module MemTy <: MEMO STT := FMemo STT.
  Module Defs := memo.MemoDefsFn STT MemTy.
End memoTFn.


Require Import HahnBase.

Require Import VST.floyd.proofauto.
Require Import VST.floyd.library.
Require Import printer.printer_files.compiled_format.
Require Import printer.verified_printer.Format.
Require Import printer.verified_printer.FormatTrivial.
Require Import Coq.Strings.Ascii.
Require Import format_specs.

Definition t_flist := Tstruct _format_list noattr.

Fixpoint listrepf (sigma: list t) (p: val) : mpred :=
 match sigma with
 | G::hs =>
    EX x:val, EX y: val,
    mformat G y *
    malloc_token Ews t_flist p *
    data_at Ews t_flist ((y, x) : @reptype CompSpecs t_flist) p *
    listrepf hs x
 | nil =>
    !! (<< LIST_NULL_PTR : p = nullval >> ) && emp
 end.

Arguments listrepf sigma p : simpl never.

Lemma listrepf_local_facts sigma p :
   listrepf sigma p |--
   !! (<< LIST_PTR_FACT : is_pointer_or_null p /\ (p=nullval <-> sigma=nil) >>).
Proof.
  intros.
  revert p; induction sigma; intros p.
  { unfold listrepf. unnw. entailer!. split; auto. }
  unff listrepf.
  destruct a. entailer. unnw. entailer!.
  split; ins.
  subst. eapply field_compatible_nullval; eauto.
Qed.
#[export] Hint Resolve listrepf_local_facts : saturate_local.

Lemma listrepf_valid_pointer sigma p :
   listrepf sigma p |-- valid_pointer p.
Proof.
  intros.
  unfold listrepf. destruct sigma; simpl; unnw.
  { entailer!. }
  Intros x y. auto with valid_pointer.
Qed.
#[export] Hint Resolve listrepf_valid_pointer : valid_pointer.

Fixpoint lsegf (sigma: list t) (x z: val) : mpred :=
  match sigma with
  | nil => !! (<< LSEG_PTR_FACT : x = z >>) && emp
  | G::hs => EX h: val, EX y:val,
      mformat G y *
      malloc_token Ews t_flist x *
      data_at Ews t_flist ((y, h) : @reptype CompSpecs t_flist) x *
      lsegf hs h z
  end.

Arguments lsegf sigma x z : simpl never.

Definition good_format (G : t) (w h : Z) : Prop :=
   (total_width G <= Z.to_nat w)%nat /\ (G.(height) <= Z.to_nat h)%nat.

Definition good_format_list (sigma : list t) (w h : Z) : Prop :=
  Forall (fun G => good_format G w h) sigma.

Definition clear_to_text_spec : ident * funspec :=
DECLARE _clear_to_text
   WITH l : list (Z * list byte), p : val, gv : globals
   PRE [ tptr t_list ]
      PROP() PARAMS(p) GLOBALS(gv)
      SEP (listrep l p; mem_mgr gv)
   POST [ tvoid ]
      PROP() RETURN() SEP(mem_mgr gv).

Definition clear_format_list_spec : ident * funspec :=
DECLARE _clear_format_list
   WITH fs : list t, p : val, gv : globals
   PRE [ tptr t_flist ]
      PROP() PARAMS(p) GLOBALS(gv)
      SEP (listrepf fs p; mem_mgr gv)
   POST [ tvoid ]
      PROP() RETURN() SEP(mem_mgr gv).

Definition beside_doc_spec : ident * funspec :=
DECLARE _beside_doc
   WITH fs1 : list t, fs2 : list t, p1 : val, p2 : val, w : Z, h : Z, gv : globals
   PRE [ tuint, tuint, tptr t_flist, tptr t_flist ]
      PROP (good_format_list fs1 w h; good_format_list fs2 w h;
            0 <= 4 * w <= Int.max_unsigned;
            0 <= 4 * h <= Int.max_unsigned)
      PARAMS(Vint (Int.repr w); Vint (Int.repr h); p1; p2) GLOBALS(gv)
      SEP (listrepf fs1 p1; listrepf fs2 p2; mem_mgr gv)
   POST [ tptr t_flist ]
      EX p: val, EX sigma: list t,
      PROP (good_format_list sigma w h;
            sigma = filter (fun G => (G.(height) <=? (Z.to_nat h))%nat) (besideDoc (Z.to_nat w) fs1 fs2))
      RETURN(p)
      SEP (listrepf sigma p; mem_mgr gv).

Definition Gprog : funspecs :=
        ltac:(with_library prog [
                   max_spec; strlen_spec; strcpy_spec; strcat_spec;
                   list_copy_spec; less_components_spec; is_less_than_spec;
                   empty_spec; line_spec; sp_spec;
                   get_applied_length_spec; format_copy_spec; get_list_tail_spec;
                   mdw_add_above_spec; list_concat_spec; to_text_add_above_spec;
                   new_list_spec; add_above_spec;
                   flw_add_beside_spec; shift_list_spec; add_beside_spec; line_concats_spec;
                   mdw_add_beside_spec; to_text_add_beside_spec;
                   mdw_add_fill_spec; flw_add_fill_spec; to_text_add_fill_spec;
                   llw_add_fill_spec; add_fill_spec; beside_doc_spec; clear_format_list_spec; clear_to_text_spec
 ]).


From iris.algebra Require Import frac.
From iris.proofmode Require Import tactics.
Require Import Eqdep_dec List.
From cap_machine Require Import stack_macros_helpers.
From cap_machine Require Export iris_extra addr_reg_sample contiguous.
From cap_machine Require Import macros.
From cap_machine.binary_model.rules_binary Require Import rules_binary rules_binary_StoreU_derived.
From cap_machine.binary_model.examples_binary Require Import req_binary.

Ltac iPrologue_s prog :=
  (try iPrologue_pre);
  iDestruct prog as "[Hi Hprog]".

Ltac iEpilogue_s :=
   iMod (do_step_pure _ [] with "[$Hspec $Hj]") as "Hj";auto;
   iSimpl in "Hj".

Section macros.
  Context {Σ:gFunctors} {memg:memG Σ} {regg:regG Σ}
          {cfg : cfgSG Σ}
          `{MP: MachineParameters}.

  (* TODO: move this to the rules_binary_Lea.v file. small issue with the spec of failure: it does not actually
     require/leave a trace on dst! It would be good if req_regs of a failing get does not include dst (if possible) *)
  Lemma step_Lea_fail_U E K pc_p pc_g pc_b pc_e pc_a w r1 rv p g b e a z a' :
    decodeInstrW w = Lea r1 (inr rv) →
    isCorrectPC (inr ((pc_p,pc_g),pc_b,pc_e,pc_a)) →
    (a + z)%a = Some a' ->
     (match p with
      | URW | URWL | URWX | URWLX => (a < a')%a
      | _ => False
      end) ->
     nclose specN ⊆ E →

     spec_ctx ∗ ⤇ fill K (Instr Executable)
              ∗ ▷ PC ↣ᵣ inr ((pc_p,pc_g),pc_b,pc_e,pc_a)
              ∗ ▷ pc_a ↣ₐ w
              ∗ ▷ r1 ↣ᵣ inr ((p,g),b,e,a)
              ∗ ▷ rv ↣ᵣ inl z
     ={E}=∗
         ⤇ fill K (Instr Failed).
  Proof.
    iIntros (Hdecode Hvpc Hz Hp Hnclose) "(#Hspec & Hj & >HPC & >Hpc_a & >Hsrc & >Hdst)".
    iDestruct (rules_binary_base.map_of_regs_3 with "HPC Hsrc Hdst") as "[Hmap (%&%&%)]".
    iMod (step_lea with "[$Hmap Hpc_a $Hspec $Hj]") as (regs' retv) "(Hj & #Hspec' & Hpc_a & Hmap)"; eauto; simplify_map_eq; eauto.
      by rewrite !dom_insert; set_solver+.
    iDestruct "Hspec'" as %Hspec.
    destruct Hspec as [* Hsucc |].
    { (* Success (contradiction) *) simplify_map_eq. destruct p0; try done; revert Hp H5;clear;solve_addr. }
    { (* Failure, done *) iFrame. done. }
  Qed.

  Definition prepstackU_s r minsize paramsize a : iProp Σ :=
    ([∗ list] a_i;w_i ∈ a;(prepstackU_instrs r minsize paramsize), a_i ↣ₐ w_i)%I.

  Lemma prepstackU_s_spec E r minsize paramsize a w pc_p pc_g pc_b pc_e a_first a_last :
    isCorrectPC_range pc_p pc_g pc_b pc_e a_first a_last ->
    contiguous_between a a_first a_last ->
    nclose specN ⊆ E →

    spec_ctx ∗ ⤇ Seq (Instr Executable)
    ∗ ▷ prepstackU_s r minsize paramsize a
    ∗ ▷ PC ↣ᵣ inr (pc_p,pc_g,pc_b,pc_e,a_first)
    ∗ ▷ r ↣ᵣ w
    ∗ ▷ (∃ w, r_t1 ↣ᵣ w)
    ∗ ▷ (∃ w, r_t2 ↣ᵣ w)
    ={E}=∗ (if isPermWord w URWLX then
           ∃ l b e a', ⌜w = inr (URWLX,l,b,e,a')⌝ ∗
           if (minsize + paramsize <? e - b)%Z then
             if ((b + paramsize) <=? a')%Z then
               (∃ a_param, ⌜(b + paramsize)%a = Some a_param⌝ ∧
                   ⤇ Seq (Instr Executable) ∗
                   PC ↣ᵣ inr (pc_p,pc_g,pc_b,pc_e,a_last) ∗ prepstackU_s r minsize paramsize a ∗
                   r ↣ᵣ inr (URWLX,l,b,e,a_param) ∗ r_t1 ↣ᵣ inl 0%Z ∗ r_t2 ↣ᵣ inl 0%Z)
             else ⤇ Seq (Instr Failed)
           else ⤇ Seq (Instr Failed)
         else ⤇ Seq (Instr Failed)).
  Proof.
    iIntros (Hvpc Hcont Hnclose) "(#Hspec & Hj & >Hprog & >HPC & >Hr & >Hr_t1 & >Hr_t2)".
    iDestruct (big_sepL2_length with "Hprog") as %Hlength. simpl in *.
    (* reqperm *)
    iPrologue_multi "Hprog" Hcont Hvpc link.
    iMod (reqperm_spec with "[- $HPC $Hcode $Hr $Hr_t1 $Hr_t2 Hprog $Hspec $Hj]") as "Hcond"; [apply Hvpc_code|apply Hcont_code|auto..].
    destruct (isPermWord w URWLX); auto.
    iDestruct "Hcond" as (l b e a' Heq) "Hφ".
    subst. iExists l,b,e,a'. iSplitR; auto.
    iDestruct "Hφ" as  "(Hj & HPC & Hprog_done & Hr & Hr_t1 & Hr_t2)".
    (* reqsize *)
    iPrologue_multi "Hprog" Hcont Hvpc link0.
    iMod (reqsize_spec with "[$HPC $Hcode $Hr $Hr_t1 $Hr_t2 $Hspec $Hj]") as "Hφ";
      [apply Hvpc_code0|eauto|auto..].
    destruct (minsize + paramsize <? e - b)%Z eqn:Hsize; auto.
    iDestruct "Hφ" as (w1 w2) "(Hj & Hreqsize & HPC & Hr & Hr_t1 & Hr_t2)".
    (* getb r_t1 r *)
    iDestruct (big_sepL2_length with "Hprog") as %Hlength'.
    prep_addr_list_full l_rest0 Hcont.
    iPrologue_s "Hprog".
    iMod (step_Get_success _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr $Hr_t1]")
      as "(Hj & HPC & Hi & Hr & Hr_t1) /=";
      [apply decode_encode_instrW_inv|auto|iCorrectPC link0 a_last|iContiguous_next Hcont 0|auto..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    (* geta r_t2 r *)
    iPrologue_s "Hprog".
    iMod (step_Get_success _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr $Hr_t2]")
      as "(Hj & HPC & Hi & Hr & Hr_t2) /=";
      [apply decode_encode_instrW_inv|auto|iCorrectPC link0 a_last|iContiguous_next Hcont 1|auto..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    (* add r_t2 r_t2 paramsize *)
    iPrologue_s "Hprog".
    iMod (step_add_sub_lt_success_dst_z _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr_t2]")
      as "(Hj & HPC & Hi & Hr_t2) /=";
      [apply decode_encode_instrW_inv|auto|iContiguous_next Hcont 2|iCorrectPC link0 a_last|auto..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    (* sub r_t1 r_t1 r_t2 *)
    iPrologue_s "Hprog".
    iMod (step_add_sub_lt_success_dst_r _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr_t2 $Hr_t1]")
      as "(Hj & HPC & Hi & Hr_t2 & Hr_t1) /=";
      [apply decode_encode_instrW_inv|auto|iContiguous_next Hcont 3|iCorrectPC link0 a_last|auto..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    (* we need to distinguish between the case where the capability is stuck, or usable *)
    assert (∃ a_param, (b + paramsize)%a = Some a_param) as [a_param Ha_param].
    { destruct (b + paramsize)%a eqn:Hnone;eauto. exfalso. clear -Hnone Hsize. apply Z.ltb_lt in Hsize. solve_addr. }
    assert ((a' + (b - (a' - paramsize)))%a = Some a_param) as Hlea;[clear -Ha_param; solve_addr|].
    destruct (decide (a_param <= a')%a).
    2: { (* lea fail *)
      iPrologue_s "Hprog".
      iMod (step_Lea_fail_U _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr_t1 $Hr]")
        as "Hj";
        [apply decode_encode_instrW_inv|iCorrectPC link0 a_last|apply Hlea|auto..].
      { simpl. solve_addr. }
      assert (b + paramsize <=? a' = false)%Z as ->;[apply Z.leb_gt;solve_addr|].
      iFrame. done. }
    (* lea r r_t1 *)
    iPrologue_s "Hprog".
    iMod (step_lea_success_reg _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr_t1 $Hr]")
      as "(Hj & HPC & Hi & Hr_t1 & Hr)";
      [apply decode_encode_instrW_inv|iCorrectPC link0 a_last|iContiguous_next Hcont 4|apply Hlea|auto..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    (* move r_t1 0 *)
    iPrologue_s "Hprog".
    iMod (step_move_success_z _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr_t1]")
      as "(Hj & HPC & Hi & Hr_t1)";
      [apply decode_encode_instrW_inv|iCorrectPC link0 a_last|iContiguous_next Hcont 5|auto..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    (* move r_t2 0 *)
    iPrologue_s "Hprog".
    apply contiguous_between_last with (ai:=a5) in Hcont as Hlast;[|auto].
    iMod (step_move_success_z _ [SeqCtx] with "[$Hspec $Hj $HPC $Hi $Hr_t2]")
      as "(Hj & HPC & Hi & Hr_t2)";
      [apply decode_encode_instrW_inv|iCorrectPC link0 a_last|apply Hlast|auto|..].
    iEpilogue_s; iCombine "Hi" "Hprog_done" as "Hprog_done".
    assert (b + paramsize <=? a' = true)%Z as ->;[apply Zle_is_le_bool;auto;clear -Ha_param l0;solve_addr|].
    iExists a_param. iSplitR;auto. iFrame.
    repeat (iDestruct "Hprog_done" as "[Hi Hprog_done]"; iFrame "Hi").
    iFrame "Hprog_done". done.
  Qed.

End macros.