(* iris/program_logic/total_adequacy.v *)
From iris.algebra Require Import gmap auth agree gset coPset list
From iris.bi Require Import big_op fixpoint_mono
From iris.proofmode Require Import proofmode
From iris.program_logic Require Export total_weakestpre adequacy
From iris.prelude Require Import options

Definition twptp_pre(twptp : list (expr Λ) → iProp Σ)
    (t1 : list (expr Λ)) : iProp Σ :=
  ∀ t2 σ1 ns κ κs σ2 nt, ⌜step (t1,σ1) κ (t2,σ2)⌝ -∗
    state_interp σ1 ns κs nt ={⊤}=∗
              ∃ nt', ⌜κ = []⌝ ∗ state_interp σ2 (S ns) κs nt' ∗ twptp t2.

Lemma twptp_pre_mono(twptp1 twptp2 : list (expr Λ) → iProp Σ) :
  □ (∀ t, twptp1 t -∗ twptp2 t) -∗
  ∀ t, twptp_pre twptp1 t -∗ twptp_pre twptp2 t.
Proof.
  iIntros "#H"; iIntros (t) "Hwp". rewrite /twptp_pre.
  iIntros (t2 σ1 ns κ κs σ2 nt1) "Hstep Hσ".
  iMod ("Hwp" with "[$] [$]") as (n2) "($ & Hσ & ?)".
  iModIntro. iExists n2. iFrame "Hσ". by iApply "H".
Qed.

Local Instancetwptp_pre_mono' : BiMonoPred twptp_pre.
Proof.
  constructor; first (intros ????; apply twptp_pre_mono).
  intros wp Hwp n t1 t2 ?%(discrete_iff _ _ _)%leibniz_equiv; solve_proper.
Qed.

Definition twptp(t : list (expr Λ)) : iProp Σ :=
  bi_least_fixpoint twptp_pre t.

Lemma twptp_unfoldt : twptp t ⊣⊢ twptp_pre twptp t.
Proof. by rewrite /twptp least_fixpoint_unfold. Qed.

Lemma twptp_indΨ :
  ⊢ (□ ∀ t, twptp_pre (λ t, Ψ t ∧ twptp t) t -∗ Ψ t) → ∀ t, twptp t -∗ Ψ t.
Proof.
  iIntros "#IH" (t) "H".
  assert (NonExpansive Ψ).
  { by intros n ?? ->%(discrete_iff _ _ _)%leibniz_equiv. }
  iApply (least_fixpoint_ind _ Ψ with "[] H").
  iIntros "!>" (t') "H". by iApply "IH".
Qed.

Local Instancetwptp_Permutation : Proper ((≡ₚ) ==> (⊢)) twptp.
Proof.
  iIntros (t1 t1' Ht) "Ht1". iRevert (t1' Ht); iRevert (t1) "Ht1".
  iApply twptp_ind; iIntros "!>" (t1) "IH"; iIntros (t1' Ht).
  rewrite twptp_unfold /twptp_pre. iIntros (t2 σ1 ns κ κs σ2 nt Hstep) "Hσ".
  destruct (step_Permutation t1' t1 t2 κ σ1 σ2) as (t2'&?&?); [done..|].
  iMod ("IH" $! t2' with "[% //] Hσ") as (n2) "($ & Hσ & IH & _)".
  iModIntro. iExists n2. iFrame "Hσ". by iApply "IH".
Qed.

Lemma twptp_appt1 t2 : twptp t1 -∗ twptp t2 -∗ twptp (t1 ++ t2).
Proof.
  iIntros "H1". iRevert (t2). iRevert (t1) "H1".
  iApply twptp_ind; iIntros "!>" (t1) "IH1". iIntros (t2) "H2".
  iRevert (t1) "IH1"; iRevert (t2) "H2".
  iApply twptp_ind; iIntros "!>" (t2) "IH2". iIntros (t1) "IH1".
  rewrite twptp_unfold /twptp_pre. iIntros (t1'' σ1 ns κ κs σ2 nt Hstep) "Hσ1".
  destruct Hstep as [e1 σ1' e2 σ2' efs' t1' t2' [=Ht ?] ? Hstep]; simplify_eq/=.
  apply app_eq_inv in Ht as [(t&?&?)|(t&?&?)]; subst.

Lemma twp_twptps Φ e : WP e @ s; ⊤ [{ Φ }] -∗ twptp [e].
Proof.
  iIntros "He". remember (⊤ : coPset) as E eqn:HE.
  iRevert (HE). iRevert (e E Φ) "He". iApply twp_ind.
  iIntros "!>" (e E Φ); iIntros "IH" (->).
  rewrite twptp_unfold /twptp_pre /twp_pre.
  iIntros (t1' σ1' ns κ κs σ2' nt Hstep) "Hσ1".
  destruct Hstep as [e1 σ1 e2 σ2 efs [|? t1] t2 ?? Hstep];
    simplify_eq/=; try discriminate_list.
  destruct (to_val e1) as [v|] eqn:He1.
  { apply val_stuck in Hstep; naive_solver. }
  iMod ("IH" with "Hσ1") as "[_ IH]".
  iMod ("IH" with "[% //]") as "($ & Hσ & [IH _] & IHfork)".
  iModIntro. iExists (length efs + nt). iFrame "Hσ".
  iApply (twptp_app [_] with "(IH [//])").
  clear. iInduction efs as [|e efs IH]; simpl.
  { rewrite twptp_unfold /twptp_pre. iIntros (t2 σ1 ns κ κs σ2 nt1 Hstep).
    destruct Hstep; simplify_eq/=; discriminate_list. }
  iDestruct "IHfork" as "[[IH' _] IHfork]".
  iApply (twptp_app [_] with "(IH' [//])"). by iApply "IH".
Qed.

Lemma twptp_totalσ ns nt t :
  state_interp σ ns [] nt -∗ twptp t ={⊤}=∗ ⌜sn erased_step (t, σ)⌝.
Proof.
  iIntros "Hσ Ht". iRevert (σ ns nt) "Hσ". iRevert (t) "Ht".
  iApply twptp_ind; iIntros "!>" (t) "IH"; iIntros (σ ns nt) "Hσ".
  iApply (pure_mono _ _ (Acc_intro _)). iIntros ([t' σ'] [κ Hstep]).
  rewrite /twptp_pre.
  iMod ("IH" with "[% //] Hσ") as (n' ->) "[Hσ [H _]]".
  by iApply "H".
Qed.

Theorem twp_totalΣ Λ `{!invGpreS Σ} s e σ Φ n :
  (∀ `{Hinv : !invGS_gen HasNoLc Σ},
     ⊢ |={⊤}=> ∃
         (stateI : state Λ → nat → list (observation Λ) → nat → iProp Σ)