From iris.heap_lang Require Import proofmode notation.

Section mtree_examples.

(** Telling Iris what algebras to use in the following *)
Context `{!heapGS Σ}.

Definition null : loc := Loc 0.

(* [malloc] is just like normal allocation but returns a non-null location *)
Definition malloc : val :=
  λ: "n" "v",
    if: #0 < "n" then
      let: "l" := AllocN "n" "v" in
      if: "l" = #null then AllocN "n" "v" else "l"
    else #(Loc 91).

Lemma malloc_spec (n : Z) (v : val):
  {{{ True }}}
    malloc #n v
  {{{ (l : loc), RET #l; l ↦∗ replicate (Z.to_nat n) v ∗ ⌜l ≠ null⌝ }}}.
Proof.
  iIntros (Φ) "_ ret". wp_lam. wp_pures.
  case_bool_decide.
  - wp_alloc l as "Hl". lia. wp_pures.
    destruct (decide (l = null)) as [-> | nl].
    + wp_if. wp_alloc l' as "Hl'". lia. iApply "ret".
      assert (Z.to_nat n = S (Z.to_nat (n - 1))) as -> by lia.
      rewrite /= !array_cons.
      iModIntro; iSplit; iFrame.
      iDestruct "Hl" as "(? & _)".
      iDestruct "Hl'" as "(? & _)".
      iApply (pointsto_ne with "[$] [$]").
    + rewrite bool_decide_eq_false_2. 2: congruence.
      wp_if. iApply "ret"; auto.
  - wp_if. iApply "ret".
    assert (Z.to_nat n = 0) as -> by lia. rewrite /array /=. auto.
Qed.


(** Trees and logical assertions *)

Inductive tree := Leaf | Node (x : Z) (T1 T2 : tree).

Inductive depth : nat -> tree -> Prop :=
  | depth_Leaf : depth 0 Leaf
  | depth_Node n x T1 T2 : depth n T1 -> depth n T2 -> depth (S n) (Node x T1 T2).

Fixpoint Mtree (T : tree) (l : loc) : iProp Σ :=
  match T with
  | Leaf => ⌜ l = null ⌝
  | Node x T1 T2 =>
      ∃ l1 l2 : loc,
        ⌜ l ≠ null ⌝
        ∗ l ↦∗ [(#x); #l1; #l2]
        ∗ Mtree T1 l1 ∗ Mtree T2 l2
  end%I.


(** Building a new record, i.e. allocation + updating fields *)

Definition new_node : val :=
  λ: "v0" "v1" "v2",
    let: "p" := malloc #3 #() in
    ("p" +ₗ #0) <- "v0";;
    ("p" +ₗ #1) <- "v1";;
    ("p" +ₗ #2) <- "v2";;
    "p".

Lemma array_3_split l v0 v1 v2 :
  l ↦∗ [v0; v1; v2] ⊣⊢
    (l +ₗ 0) ↦ v0
  ∗ (l +ₗ 1) ↦ v1
  ∗ (l +ₗ 2) ↦ v2.
Proof.
  rewrite !array_cons array_nil /Loc.add -Z.add_assoc /= Z.add_0_r.
  destruct l.
  iSplit; iIntros "(? & (? & (? & ?)))"; iFrame.
Qed.


(* [copy] program from the lecture *)

Definition tree_copy : val :=
  rec: "copy" "l" :=
    if: "l" = #null then #null else
      new_node
        !("l" +ₗ #0)
        ("copy" !("l" +ₗ #1))
        ("copy" !("l" +ₗ #2)).

Lemma tree_copy_spec : forall T l,
  {{{ Mtree T l }}}
    tree_copy #l
  {{{ (l' : loc), RET #l'; Mtree T l ∗ Mtree T l' }}}.
Proof.
  iIntros (T l Φ) "T post".
  iInduction T as [ | x T1 ? T2 ] "IH" forall (l Φ).
  - wp_rec. simpl.
    iDestruct "T" as %->.
    wp_pures.
    iApply "post". auto.
  - wp_rec.
    simpl.
    iDestruct "T" as (l1 l2) "(%n & l & T1 & T2)".
    wp_pure.
    rewrite bool_decide_eq_false_2. 2: congruence.
    wp_if.
    unfold new_node.
    rewrite array_3_split.
    iDestruct "l" as "(l0 & l1 & l2)".
    wp_load.
    wp_apply ("IH1" with "T2") as (l2') "(T2 & T2')".
    wp_load.
    wp_apply ("IH" with "T1") as (l1') "(T1 & T1')".
    wp_pures.
    wp_load.
    wp_pures.
    wp_apply (malloc_spec 3) as (l') "(l' & %l'n)"; auto.
    wp_pures.
    rewrite array_3_split.
    iDestruct "l'" as "(l'0 & l'1 & l'2)".
    wp_store.
    wp_store.
    wp_store.
    iApply "post". iModIntro.
    iSplitL "l0 l1 l2 T1 T2".
    + iExists l1, l2. iSplit. auto. rewrite array_3_split. iFrame.
    + iExists l1', l2'. iSplit. auto. rewrite array_3_split. iFrame.
Qed.

Definition Mtreecomplete (T : tree) (l : loc) : iProp Σ :=
  ∃ n, Mtree T l ∗ ⌜depth n T⌝.

Lemma complete_tree_copy_spec : forall T l,
  {{{ Mtreecomplete T l }}}
    tree_copy #l
  {{{ (l' : loc), RET #l'; Mtreecomplete T l ∗ Mtreecomplete T l' }}}.
Proof.
  iIntros (T l Φ) "Hl HΦ".
  unfold Mtreecomplete at 1.
  (* Exists *)
  iDestruct "Hl" as (n) "Hl".
  (* Prop *)
  iDestruct "Hl" as "(Hl & %HT)".
  (* using the normal spec of [tree_copy]: *)
  wp_apply (tree_copy_spec with "Hl") as (l') "(Hl & Hl')".
  iApply "HΦ".
  iSplitL "Hl".
  - unfold Mtreecomplete. iExists n. iFrame. iPureIntro. auto.
  - iExists n. iSplit. iFrame. auto.
Qed.

End mtree_examples.