From iris.program_logic Require Import atomic.
From iris.heap_lang Require Import notation proofmode.
From iris.algebra.lib Require Import excl_auth.
From iris.base_logic.lib Require Import invariants.
From iris.base_logic.lib Require Import invariants mono_nat mono_list.
From iris_named_props Require Import custom_syntax.

From lmpmc Require fin_queue_spec lmpmc_queue_spec.
From lmpmc Require Import upstream util.

Section lmpmc_queue.
  Context {fq : fin_queue_spec.fin_queue_hl}.

  Definition new : val := λ: <>,
    let: "ℓ_q" := fq.(fin_queue_spec.new) #() in ("ℓ_q", "ℓ_q").

  Definition enqueue : val := (* passing arguments [ℓ_enq data] *)
    fq.(fin_queue_spec.enqueue).
  
  Definition try_dequeue : val := rec: "go" "ℓ_q" :=
    match: fq.(fin_queue_spec.try_dequeue) "ℓ_q" with
      InjL "v" =>
        SOME "v"
    | InjR "mℓ_nextq" =>
        match: "mℓ_nextq" with
          NONE => NONE
        | SOME "ℓ_nextq" => "go" "ℓ_nextq"
        end
    end.
  
  Definition link : val := (* passing arguments [ℓ_enq1 ℓ_deq2] *)
    fq.(fin_queue_spec.finalize).

  Section proofs.
    Context `{!heapGS Σ} {fqp : fin_queue_spec.fin_queue_iris Σ fq}.
  
    Record node := Node { ℓ_node : loc; vs_node : list val }.
    Add Printing Constructor node.
    Definition linked_single_queue_repr (n : node) (mℓ_nextq : option loc) : iProp Σ :=
      fqp.(fin_queue_spec.queue_repr fq) n.(ℓ_node) n.(vs_node) (loc_to_val <$> mℓ_nextq).
    
    Definition chain := list node.
    Definition queue_repr ℓ_deq ℓ_enq (cvs : list val) : iProp Σ :=
      ∃ (c : chain),
        ⌜length c > 0⌝ ∗ ⌜ℓ_node <$> head c = Some ℓ_deq⌝ ∗
        ⌜ℓ_node <$> last c = Some ℓ_enq⌝ ∗ ⌜cvs = c ≫= vs_node⌝ ∗
        [∗ list] i ↦ l ∈ c, linked_single_queue_repr l (ℓ_node <$> c !! S i).

    #[global] Instance queue_repr_timeless ℓ_deq ℓ_enq cvs : Timeless (queue_repr ℓ_deq ℓ_enq cvs) := _.
    
    Lemma new_spec :
      {{{ True }}}
        new #()
        {{{ (ℓ_deq ℓ_enq : loc), RET (#ℓ_deq, #ℓ_enq); queue_repr ℓ_deq ℓ_enq [] }}}.
    Proof.
      iIntros (Φ _) "HΦ". iUnfold new. wp_pures.
      wp_apply (fqp.(fin_queue_spec.new_spec fq) with "[//]").
      iIntros (ℓ) "H". wp_pures. iModIntro. iApply "HΦ".
      unfold queue_repr. iExists [Node ℓ []].
      iFrame. iSimplifyEq. by iSplit; first (iPureIntro; lia).
    Qed.
    
    Lemma enqueue_spec (ℓ_enq : loc) (data : val) :
      ⊢ <<{ ∀∀ ℓ_deq cvs, queue_repr ℓ_deq ℓ_enq cvs }>>
        enqueue #ℓ_enq data @ ∅
        <<{ queue_repr ℓ_deq ℓ_enq (cvs ++ [data]) | RET #() }>>.
    Proof.
      iIntros "%Φ AU". iUnfold enqueue.
      
      wp_bind (fin_queue_spec.enqueue _ _ _)%E.
      awp_apply fqp.(fin_queue_spec.enqueue_spec fq).
      rewrite /atomic_acc /=.
      iMod "AU" as "(%ℓ_deq & %cvs & (%c & %Hclen & %Hhead & %Hlast & %Hcvs & Hqs) & Hclose)".
      iModIntro.
      
      destruct (last c) as [n_last|] eqn:Hlastn; last done.
      iExists n_last.(vs_node).
      
      apply last_Some in Hlastn as [c' ->].
      iDestruct (big_sepL_insert_acc _ _ (length c') with "Hqs") as "[Hnode Hrepr]".
      { by rewrite lookup_app_r  // Nat.sub_diag. }
      iSimplifyEq. iUnfold linked_single_queue_repr in "Hnode".
      
      rewrite lookup_app_r; last lia.
      replace (S (length c') - length c') with 1 by lia.
      simplify_eq/=. iFrame.
      
      iSplit.
      - iDestruct "Hclose" as "[Hclose _]".
        iIntros "Hnode".
        iDestruct ("Hrepr" with "[Hnode //]") as "Hrepr".
        rewrite list_insert_id; last first.
        { by rewrite lookup_app_r // Nat.sub_diag. }
        iMod ("Hclose" with "[$Hrepr]") as "AU".
        { by rewrite last_snoc. }
        done.
      - iDestruct "Hclose" as "[_ Hclose]".
        iIntros "Hnode".
        
        set n_last' := Node n_last.(ℓ_node) (n_last.(vs_node) ++ [data]).
        set c := c' ++ [n_last'].
        
        iDestruct ("Hrepr" with "[Hnode]") as "Hrepr".
        { rewrite /linked_single_queue_repr /=.
          replace (ℓ_node n_last) with n_last'.(ℓ_node) by done.
          by replace (vs_node n_last ++ [data]) with n_last'.(vs_node). }
        
        iMod ("Hclose" with "[Hrepr]") as "HΦ".
        { unfold queue_repr. iExists c. repeat iSplit; try done.
          - iPureIntro. rewrite length_app /=. lia.
          - iPureIntro. rewrite /c. by destruct c'; simplify_eq/=.
          - iPureIntro. by rewrite /c last_snoc.
          - iPureIntro. by rewrite /c !bind_app /= !app_nil_r app_assoc.
          - rewrite insert_app_r_alt // Nat.sub_diag /=.
            iApply (big_sepL_mono with "Hrepr").
            { iIntros (i n Hi) "Hrepr".
              rewrite -list_lookup_fmap fmap_app.
              replace (ℓ_node <$> [n_last]) with (ℓ_node <$> [n_last']) by done.
              by rewrite -fmap_app list_lookup_fmap. } }
        
        iModIntro. iApply "HΦ".
    Qed.
    
    Lemma try_dequeue_spec (ℓ_deq : loc) :
      ⊢ <<{ ∀∀ ℓ_enq cvs, queue_repr ℓ_deq ℓ_enq cvs }>>
        try_dequeue #ℓ_deq @ ∅
        <<{ queue_repr ℓ_deq ℓ_enq (tail cvs) | RET (val_opt_hl (head cvs)) }>>.
    Proof.
      revert ℓ_deq. iLöb as "IH". iIntros (ℓ_deq Φ) "AU". wp_rec.
      
      awp_apply (fqp.(fin_queue_spec.try_dequeue_spec fq)).
      rewrite /atomic_acc /=.
      iMod "AU" as "(%ℓ_enq & %cvs & (%c & %Hclen & %Hhead & %Hlast & %Hcvs & Hqs) & Hclose)". iModIntro.
      
      destruct c as [|n_head c']; first done.
      iDestruct (big_sepL_cons with "Hqs") as "[Hnode Hqs]".
      simpl. destruct (c' !! 0) as [n_next|] eqn:Hnextq.
      - iExists n_head.(vs_node), (Some #n_next.(ℓ_node)). iSplitL "Hnode"; last iSplit.
        + simplify_eq/=. by rewrite /linked_single_queue_repr /=.
        + simplify_eq/=. iIntros "Hnode". iDestruct "Hclose" as "[Hclose _]".
          pose c := n_head :: c'.
          iAssert (linked_single_queue_repr n_head (ℓ_node <$> c !! S 0)) with "[Hnode]" as "Hnode".
          { by rewrite /linked_single_queue_repr /= Hnextq.  }
          iAssert ([∗ list] i↦n ∈ c', linked_single_queue_repr n (ℓ_node <$> c !! S (S i)))%I with "[Hqs //]" as "Hqs".
          iAssert ([∗ list] i↦n ∈ c, linked_single_queue_repr n (ℓ_node <$> c !! S i))%I with "[Hnode Hqs]" as "Hqs".
          { iApply big_sepL_cons. iFrame. }
          iAssert (queue_repr n_head.(ℓ_node) ℓ_enq (vs_node n_head ++ c' ≫= vs_node)) with "[$Hqs //]" as "Hrepr".
          by iMod ("Hclose" with "Hrepr") as "Hrepr".
        + simplify_eq/=. iIntros "Hnode".
          destruct n_head.(vs_node) as [|x vs_head'] eqn:Hvs_head.
          * simplify_eq/=. iDestruct "Hclose" as "[Hclose _]".
            
            iDestruct (fqp.(fin_queue_spec.queue_fin_obtain fq) with "Hnode") as "#Hfin".
            
            pose c := n_head :: c'.
            iAssert (linked_single_queue_repr n_head (ℓ_node <$> c !! S 0)) with "[Hnode]" as "Hnode".
            { by rewrite /linked_single_queue_repr /= Hnextq Hvs_head /=.  }
            iAssert ([∗ list] i↦n ∈ c', linked_single_queue_repr n (ℓ_node <$> c !! S (S i)))%I with "[Hqs //]" as "Hqs".
            iAssert ([∗ list] i↦n ∈ c, linked_single_queue_repr n (ℓ_node <$> c !! S i))%I with "[Hnode Hqs]" as "Hqs".
            { iApply big_sepL_cons. iFrame. }
            iAssert (queue_repr n_head.(ℓ_node) ℓ_enq (vs_node n_head ++ c' ≫= vs_node)) with "[$Hqs //]" as "Hrepr".
            rewrite Hvs_head /=. iMod ("Hclose" with "Hrepr") as "AU". iModIntro. wp_pures.
            clear c c' Hclen Hlast Hnextq Hvs_head.
            
            awp_apply "IH". rewrite /atomic_acc /=.
            iMod "AU" as "(%ℓ_enq' & %cvs' & (%c'' & _ & %Hhead' & %Hlast' & %Hcvs' & Hqs') & Hclose)".
            destruct c'' as [|n_head' c'']; first done.
            iDestruct (big_sepL_cons with "Hqs'") as "[Hq0 Hqs']". simplify_eq/=.
            iEval (rewrite /linked_single_queue_repr /=) in "Hq0".
            rewrite -Hhead'.
            iDestruct (fqp.(fin_queue_spec.queue_fin_agree fq) with "Hfin Hq0") as %[Hhead Hc''].
            destruct (c'' !! 0) as [n_next'|] eqn:Hc''head; last done. simplify_eq/=.
            rewrite -Hc''. clear Hhead' Hc'' n_next n_head.
            assert (Hc''len : length c'' > 0). { by apply lookup_lt_Some in Hc''head. }
            
            iAssert (queue_repr n_next'.(ℓ_node) ℓ_enq' (c'' ≫= vs_node)) with "[$Hqs']" as "Hrepr'".
            { iPureIntro. repeat split; try done.
              - by rewrite head_lookup Hc''head.
              - enough (last c'' = last (n_head' :: c'')) as ->; first done.
                destruct c''; first done. by rewrite last_cons_cons. }
            
            iModIntro. iFrame. iSplit.
            -- iIntros "(%d & %Hdlen & %Hdhead & %Hdlast & %Hdvs & Hrepr)".
               rewrite Hhead /=. simplify_eq/=.
               pose c := n_head' :: d.
               iAssert (linked_single_queue_repr n_head' (ℓ_node <$> c !! S 0)) with "[Hq0]" as "Hq0".
               { by rewrite /linked_single_queue_repr /= Hhead /= -head_lookup Hdhead /=. }
               iAssert ([∗ list] i↦n ∈ c, linked_single_queue_repr n (ℓ_node <$> c !! S i))%I with "[Hq0 Hrepr]" as "Hrepr".
               { iApply big_sepL_cons. iFrame. }
               iAssert (queue_repr (ℓ_node n_head') ℓ_enq' (c'' ≫= vs_node)) with "[$Hrepr]" as "Hrepr".
               { iPureIntro. repeat split; try done.
                 - unfold c. simpl. lia.
                 - unfold c. rewrite -Hdlast. by destruct d.
                 - unfold c. rewrite Hdvs. simplify_eq/=.
                   by rewrite Hhead. }
               by iMod ("Hclose" with "Hrepr") as "AU".
            -- iIntros "(%d & %Hdlen & %Hdhead & %Hdlast & %Hdvs & Hrepr)".
               rewrite Hhead /=. simplify_eq/=.
               pose c := n_head' :: d.
               iAssert (linked_single_queue_repr n_head' (ℓ_node <$> c !! S 0)) with "[Hq0]" as "Hq0".
               { by rewrite /linked_single_queue_repr /= Hhead /= -head_lookup Hdhead /=. }
               iAssert ([∗ list] i↦n ∈ c, linked_single_queue_repr n (ℓ_node <$> c !! S i))%I with "[Hq0 Hrepr]" as "Hrepr".
               { iApply big_sepL_cons. iFrame. }
               iAssert (queue_repr (ℓ_node n_head') ℓ_enq' (tail (c'' ≫= vs_node))) with "[$Hrepr]" as "Hrepr".
               { iPureIntro. repeat split; try done.
                 - unfold c. simpl. lia.
                 - unfold c. rewrite -Hdlast. by destruct d.
                 - unfold c. rewrite Hdvs. simplify_eq/=.
                   by rewrite Hhead. }
               by iMod ("Hclose" with "Hrepr") as "AU".
          * simplify_eq/=. iDestruct "Hclose" as "[_ Hclose]".
            
            pose n_head' := Node (ℓ_node n_head) vs_head'.
            pose c := n_head' :: c'.
            
            iAssert (linked_single_queue_repr n_head' (ℓ_node <$> c !! S 0)) with "[Hnode]" as "Hnode".
            { by rewrite /linked_single_queue_repr /= Hnextq /=.  }
            iAssert ([∗ list] i↦n ∈ c', linked_single_queue_repr n (ℓ_node <$> c !! S (S i)))%I with "[Hqs //]" as "Hqs".
            iAssert ([∗ list] i↦n ∈ c, linked_single_queue_repr n (ℓ_node <$> c !! S i))%I with "[Hnode Hqs]" as "Hqs".
            { iApply big_sepL_cons. iFrame. }
            iAssert (queue_repr n_head.(ℓ_node) ℓ_enq (vs_node n_head' ++ c' ≫= vs_node)) with "[$Hqs]" as "Hrepr".
            { iPureIntro. repeat split; try done. by destruct c'. }
            
            iMod ("Hclose" with "Hrepr") as "Hrepr".
            iModIntro. wp_pures. iApply "Hrepr".
      - destruct c'; last done. simplify_eq/=. iClear "Hqs".
        rewrite app_nil_r. clear Hclen.
        iExists n_head.(vs_node), None. iSplitL "Hnode".
        + by rewrite /linked_single_queue_repr /=.
        + iSplit.
          * iIntros "Hrepr". iDestruct "Hclose" as "[Hclose _]".
            iMod ("Hclose" with "[Hrepr]") as "AU".
            { iExists [n_head]. iFrame.
              repeat iSplit; try done.
              - iPureIntro. simpl. lia.
              - simpl. by rewrite app_nil_r. }
            done.
          * iIntros "Hrepr". iDestruct "Hclose" as "[_ Hclose]".
            set n_head' := Node n_head.(ℓ_node) (tail n_head.(vs_node)).
            iMod ("Hclose" with "[Hrepr]") as "HΦ".
            { iExists [n_head']. iFrame.
              repeat iSplit; try done.
              - iPureIntro. simpl. lia.
              - simpl. by rewrite app_nil_r. }
            destruct n_head.(vs_node) as [|v vs'];
              iModIntro; wp_pures; iApply "HΦ".
    Qed.
    
    (** Linking two queues:
     *     _______________      _______________
     *    /      Q1       \    /       Q2      \
     *    ℓ_deq1 ... ℓ_enq1 <> ℓ_deq2 ... ℓ_enq2
     *    \          \ link op args /          /
     *     \          \____________/          /
     *      \            Q1 <> Q2            /
     *       \______________________________/
     *)
    Lemma link_spec (ℓ_enq1 ℓ_deq2 : loc) :
      ⊢ <<{ ∀∀ (b : bool) ℓ_deq1 ℓ_enq2 cvs1 cvs2, queue_repr ℓ_deq1 ℓ_enq1 cvs1 ∗
                          if b then queue_repr ℓ_deq2 ℓ_enq2 cvs2 else ⌜ℓ_deq1 = ℓ_deq2 ∧ ℓ_enq1 = ℓ_enq2⌝ }>>
        link #ℓ_enq1 #ℓ_deq2 @ ∅
        <<{ if b then queue_repr ℓ_deq1 ℓ_enq2 (cvs1 ++ cvs2) else True | RET #() }>>.
    Proof.
      iIntros "%Φ AU". iUnfold link.
      
      wp_bind (fin_queue_spec.finalize _ _ _)%E.
      awp_apply fqp.(fin_queue_spec.finalize_spec fq).
      rewrite /atomic_acc /=.
      iMod "AU" as (b ℓ_deq1 ℓ_enq2 cvs1 cvs2) "[[(%c1 & %Hclen1 & %Hhead1 & %Hlast1 & %Hcvs1 & Hqs1) Hqs2] Hclose]".
      iModIntro.
      
      destruct (last c1) as [n_last1|] eqn:Hlastn1; last done.
      iExists n_last1.(vs_node).
      
      apply last_Some in Hlastn1 as [c1' ->].
      iDestruct (big_sepL_app with "Hqs1") as "[Hqs1 Hnode]".
      iSimplifyEq. iDestruct  "Hnode" as "[Hnode _]".
      rewrite Nat.add_0_r.
      
      rewrite lookup_ge_None_2 /=; last first.
      { rewrite length_app /=. lia. }
      iUnfold linked_single_queue_repr in "Hnode".
      iSimplifyEq. iFrame.
      
      iSplit.
      - iDestruct "Hclose" as "[Hclose _]".
        iIntros "Hnode".
        
        iAssert (linked_single_queue_repr n_last1 (ℓ_node <$> (c1' ++ [n_last1]) !! S (length c1'))) with "[Hnode]" as "Hnode".
        { rewrite /linked_single_queue_repr !lookup_app_r; last lia.
          by assert (S (length c1') - length c1' = 1) as -> by lia. }
        
        iDestruct (big_sepL_snoc with "[$Hqs1 $Hnode]") as "Hqs1'".
        
        iPoseProof ("Hclose" with "[$Hqs1' $Hqs2]") as "Hqs".
        { iPureIntro. repeat split; try done. by rewrite last_app. }
        done.
        
      - iDestruct "Hclose" as "[_ Hclose]".
        iIntros "Hnode".
        
        destruct b; last by iApply "Hclose".
        iDestruct "Hqs2" as "(%c2 & %Hclen2 & %Hhead2 & %Hlast2 & %Hcvs2 & Hqs2)".
        
        pose c := (c1' ++ [n_last1]) ++ c2.
        
        iAssert (linked_single_queue_repr n_last1 (ℓ_node <$> c !! S (length c1'))) with "[Hnode]" as "Hnode".
        { rewrite /linked_single_queue_repr /c !lookup_app_r; try (rewrite length_app /=; lia).
          rewrite length_app /=. assert (S (length c1') - (length c1' + 1) = 0) as -> by lia.
          rewrite -head_lookup Hhead2 //=. }
        
        iDestruct (big_sepL_mono with "Hqs1") as "Hqs1".
        { iIntros (i l Hi) "Hrepr".
          iAssert (linked_single_queue_repr l (ℓ_node <$> c !! S i))%I with "[Hrepr]" as "Hrepr".
          { rewrite /c lookup_app_l // length_app /=. apply lookup_lt_Some in Hi. lia. }
          iExact "Hrepr". }
        
        iDestruct (big_sepL_snoc with "[$Hqs1 $Hnode]") as "Hqs1'".
        
        iDestruct (big_sepL_mono with "Hqs2") as "Hqs2".
        { iIntros (i l Hi) "Hrepr".
          iAssert (linked_single_queue_repr l (ℓ_node <$> c !! S (length (c1' ++ [n_last1]) + i)))%I with "[Hrepr]" as "Hrepr".
          { rewrite /c length_app /= lookup_app_r length_app /=; last lia.
            by assert (S i = S (length c1' + 1 + i) - (length c1' + 1)) as -> by lia. }
          iExact "Hrepr". }
        iSimplifyEq.
        
        iDestruct (big_sepL_app with "[$Hqs1' $Hqs2]") as "Hqs". fold c.
        rewrite -bind_app -/c.
        iMod ("Hclose" with "[Hqs]") as "HΦ".
        { iFrame. iPureIntro. repeat split; try done.
          - rewrite /c !length_app /=. lia.
          - rewrite /c head_app. case_match; last done. congruence.
          - rewrite /c last_app. case_match; last done. congruence. }
        iModIntro. iApply "HΦ".
    Qed.

  End proofs.

End lmpmc_queue.

Definition lmpmc_queue_hl (fq : fin_queue_spec.fin_queue_hl) : lmpmc_queue_spec.lmpmc_queue_hl.
Proof.
  by refine {| lmpmc_queue_spec.new := new;
               lmpmc_queue_spec.enqueue := enqueue;
               lmpmc_queue_spec.try_dequeue := try_dequeue;
               lmpmc_queue_spec.link := link |}.
Defined.

Definition lmpmc_queue_iris `{!heapGS Σ}
  (fq : fin_queue_spec.fin_queue_hl)
  (fqp : fin_queue_spec.fin_queue_iris Σ fq) : lmpmc_queue_spec.lmpmc_queue_iris Σ (lmpmc_queue_hl fq).
Proof.
  by refine (lmpmc_queue_spec.LmpmcQueueIris _ _
               (lmpmc_queue_hl fq) _ _
               new_spec
               enqueue_spec
               try_dequeue_spec
               link_spec).
Defined.

#[global] Typeclasses Opaque queue_repr.