From iris.heap_lang Require Import proofmode notation assert.

Definition null : loc := Loc 0.

Definition miter : val :=
  rec: "miter" "f" "p" :=
    if: "p" = #null then
      #()
    else
      "f" !"p";;
      "miter" "f" !("p" +ₗ #1).

Definition mlength : val :=
  λ: "p",
    let: "r" := ref #0 in
    miter (λ: "x", "r" <- !"r" + #1) "p";;
    !"r".

Definition msum : val :=
  λ: "p",
    let: "r" := ref #0 in
    miter (λ: "x", "r" <- !"r" + "x") "p";;
    !"r".

Section examples. Context `{!heapGS Σ}.

Fixpoint Mlist (L : list Z) (p : loc) : iProp Σ :=
  match L with
  | nil => ⌜ p = null ⌝
  | x :: L' => ∃ p' : loc, ⌜ p ≠ null ⌝ ∗ p ↦ #x ∗ (p +ₗ 1) ↦ #p' ∗ Mlist L' p'
  end%I.

Definition iter_spec_future (it : val) := ∀ (f : val) l (I : list Z → iProp Σ),
  (∀ x s, {{{ I (x :: s) }}} f #x {{{ RET #(); I s }}}) →
  ∀ p, {{{ Mlist l p ∗ I l }}} it f #p {{{ RET #(); Mlist l p ∗ I [] }}}.

Lemma miter_correct : iter_spec_future miter.
Proof.
  iIntros (f l I Hf p φ) "(Hp & i) ret".
  iInduction l as [ | x l ] "IH" forall (p).
  - iDestruct "Hp" as %->. wp_rec. wp_pures. iApply "ret". auto.
  - simpl Mlist at 3. iDestruct "Hp" as "(%p' & %Hp' & Hp & Hp1 & Hp')".
    wp_rec. wp_pures. case_bool_decide. congruence. wp_pures. wp_load.
    wp_apply (Hf with "i") as "i". wp_pures. wp_load.
    wp_apply ("IH" with "Hp' i") as "(Hp' & i)". iApply "ret". by iFrame.
Qed.

Definition iter_spec_past (it : val) := ∀ (f : val) l (I : list Z → iProp Σ),
  (∀ x k, {{{ I k }}} f #x {{{ RET #(); I (k ++ [x]) }}}) →
  ∀ p, {{{ Mlist l p ∗ I [] }}} it f #p {{{ RET #(); Mlist l p ∗ I l}}}.

Lemma future_past it : iter_spec_future it → iter_spec_past it.
Proof.
  iIntros (F f l I Hf p φ) "(Hp & i) ret".
  wp_apply (F f l (λ k, ∃ prefix, I prefix ∗ ⌜l = prefix ++ k⌝)%I with "[$Hp $i //]")
    as "(Hp & %s & i & ->)".
  - clear φ. iIntros (x k φ) "(%s & i & ->) ret".
    wp_apply (Hf with "i") as "i".
    iApply "ret". iFrame. iPureIntro. rewrite -app_assoc //.
  - iApply "ret". rewrite app_nil_r. iFrame.
Qed.

Lemma miter_correct_past : iter_spec_past miter.
Proof. apply future_past, miter_correct. Qed.

Definition Zsum := foldl Z.add 0%Z.

Lemma msum_correct (p : loc) l :
  {{{ Mlist l p }}} msum #p {{{ RET #(Zsum l); Mlist l p }}}.
Proof.
  iIntros (φ) "Hp ret". wp_lam. wp_alloc r. wp_pures.
  wp_apply (miter_correct_past _ _ (λ k, r ↦ #(Zsum k))%I with "[$]") as "(Hp & Hr)".
  2:  wp_load; by iApply "ret".
  clear φ p l.
  iIntros (x k φ) "Hr ret". wp_pures. wp_load. wp_store. iApply "ret".
  rewrite /Zsum foldl_app /= //.
Qed.

Lemma mlength_correct (p : loc) l :
  {{{ Mlist l p }}} mlength #p {{{ RET #(length l); Mlist l p }}}.
Proof.
  iIntros (φ) "Hp ret". wp_lam. wp_alloc r. wp_pures.
  wp_apply (miter_correct_past _ _ (λ k, r ↦ #(length k))%I with "[$]") as "(Hp & Hr)".
  2:  wp_load; by iApply "ret".
  clear φ p l.
  iIntros (x k φ) "Hr ret". wp_pures. wp_load. wp_store. iApply "ret".
  rewrite /Zsum length_app /=.
  replace (length k + 1 : Z)%nat with (Z.of_nat (length k + 1)) by lia. auto.
Qed.

Lemma Zsum_cons x l : Zsum (x :: l) = (x + Zsum l)%Z.
Proof.
  rewrite /Zsum /=.
  generalize 0%Z; induction l; intros z; rewrite /=. lia.
  rewrite -IHl. f_equal. lia.
Qed.

Lemma msum_correct_future (p : loc) l :
  {{{ Mlist l p }}} msum #p {{{ RET #(Zsum l); Mlist l p }}}.
Proof.
  iIntros (φ) "Hp ret". wp_lam. wp_alloc r as "Hr". wp_pures.
  wp_apply (miter_correct _ _ (λ k, r ↦ #(Zsum l - Zsum k))%I with "[$Hp Hr]") as "(Hp & Hr)".
  2: { by rewrite Z.sub_diag. }
  2:  wp_load; rewrite Z.sub_0_r; by iApply "ret".
  clear φ p.
  iIntros (x k φ) "Hr ret". wp_pures. wp_load. wp_store. iApply "ret".
  rewrite Zsum_cons.
  replace (Zsum l - (x + Zsum k) + x)%Z with (Zsum l - Zsum k)%Z by lia.
  auto.
Qed.

Lemma mlength_correct_future (p : loc) l :
  {{{ Mlist l p }}} mlength #p {{{ RET #(length l); Mlist l p }}}.
Proof.
  iIntros (φ) "Hp ret". wp_lam. wp_alloc r as "Hr". wp_pures.
  wp_apply (miter_correct _ _ (λ k, r ↦ #(length l - length k))%I with "[$Hp Hr]") as "(Hp & Hr)".
  2: { by rewrite Z.sub_diag. }
  2:  wp_load; rewrite Z.sub_0_r; by iApply "ret".
  clear φ p.
  iIntros (x k φ) "Hr ret". wp_pures. wp_load. wp_store. iApply "ret".
  replace (length l - S (length k) + 1)%Z with (length l - length k)%Z by lia.
  auto.
Qed.

Lemma example_only_future (p : loc) :
  {{{ Mlist [1; 2; 3; 4; 5; 6]%Z p }}}
    miter (λ: "x", assert: "x" < #7)%V #p
  {{{ RET #(); Mlist [1; 2; 3; 4; 5; 6]%Z p }}}.
Proof.
  iIntros (φ) "Hp ret".
  wp_apply (miter_correct _ _ (λ k : list Z, ⌜Forall (Z.gt 7) k⌝)%I with "[$Hp]")
    as "(Hp & %)".
  2: iPureIntro; repeat (lia || constructor).
  2: by iApply "ret".
  clear φ p.
  iIntros (x k φ) "%Hr ret". wp_pures. wp_pures.
  inversion Hr; subst.
  iApply wp_assert. wp_pures.
  case_bool_decide.
  - iModIntro. iSplit; auto. iApply "ret"; auto.
  - lia.
Qed.

End examples.