Swierstra's Hoare state monad

hstate.v

From Coq Require Import Program.
From mathcomp Require Import ssreflect ssrfun ssrbool eqtype seq ssrnat.
Export Unset Program Cases.
#[local] Obligation Tactic := auto.

Section HoareState.
  Variable S : Type.

  Definition Pre : Type := S -> Prop.
  Definition Post (A : Type) : Type := S -> A -> S -> Prop.

  Definition top : Pre := fun s => True.

  Definition HoareState (pre : Pre) (A : Type) (post : Post A) : Type :=
    forall i : S, pre i -> { (x,f) : A * S | post i x f}.

  Definition ret {A} :
    forall x, HoareState top A (fun i y f => i = f /\ y = x) :=
    fun x s _ => exist _ (x,s) (conj (erefl s) (erefl x)).

  Definition bind {A B} P1 P2 Q1 Q2 :
    HoareState P1 A Q1 ->
    (forall (x : A), HoareState (P2 x) B (Q2 x)) ->
    HoareState (fun s1 => P1 s1 /\ forall x s2, Q1 s1 x s2 -> P2 x s2)
                B
               (fun s1 y s3 => exists x s2, Q1 s1 x s2 /\ Q2 x s2 y s3).
  Proof.
  move=>c1 c2 s1 [H1 H2].
  case: (c1 s1 H1)=>[[x s2] H3].
  case: (c2 x s2 (H2 _ _ H3))=>[[y s3] H4].
  by exists (y, s3), x, s2.
  Defined.

  Definition get : HoareState top S (fun i x f => i = f /\ x = i) :=
    fun s _ => exist _ (s, s) (conj (erefl s) (erefl s)).

  Definition put (x : S) : HoareState top unit (fun _ _ f => f = x) :=
    fun _ _ => exist _ (tt, x) (erefl x).

  Definition do {A} (P1 P2 : Pre) (Q1 Q2 : Post A) :
     (forall i, P2 i -> P1 i) ->
     (forall i x f, Q1 i x f -> Q2 i x f) ->
     HoareState P1 A Q1 -> HoareState P2 A Q2.
  Proof.
  move=>str wkn c s Hs.
  case: (c s (str _ Hs))=>[[a s1 H1]].
  by exists (a, s1); apply: wkn.
  Defined.

End HoareState.

Arguments top [S] _.
Arguments ret [S A] _ _.
Arguments bind [S A B] {P1 P2 Q1 Q2} _ _ _.

Notation "x '<--' c1 ';' c2" := (bind c1 (fun x => c2))
  (at level 81, right associativity).
Notation "c1 ';;' c2" := (bind c1 (fun _ => c2))
  (at level 81, right associativity).
Arguments get [S] _.
Arguments put [S] _ _.
Arguments do [S A P1 P2 Q1 Q2] _ _ _ _.
Notation "'Do' e" := (do _ _ e) (at level 80).

(* tree + functions *)

Inductive tree (A : Type) : Type :=
  Leaf of A | Node of tree A & tree A.

Arguments Leaf [A] _.
Arguments Node [A] _ _.

Fixpoint size {A} (t : tree A) : nat :=
  if t is Node l r then size l + size r else 1.

Arguments size [A] _.

Fixpoint flatten {A} (t : tree A) : seq A :=
  match t with
  | Leaf x => [::x]
  | Node l r => flatten l ++ flatten r
  end.

Arguments flatten [A] _.

Fixpoint sameShape {A B} (ta : tree A) (tb : tree B) : bool :=
  match ta, tb with
  | Leaf _, Leaf _ => true
  | Node la ra, Node lb rb => sameShape la lb && sameShape ra rb
  | _, _ => false
  end.

Arguments sameShape [A B] _ _.

Program Definition fresh :
  HoareState nat (@top nat) nat
                 (fun i n f => f = i.+1 /\ n = i) :=
  Do (n <-- @get nat;
      put n.+1;;
      ret n).
Next Obligation. by []. Qed.
Next Obligation.
by move=>i x f /= [_][_][[_ ->]][_][_][->][<- ->].
Qed.

Opaque fresh.

(* stateful relabel function *)

Program Fixpoint relabel {A} (t : tree A) :
  HoareState nat
             (@top nat)
             (tree nat)
             (fun i t' f => [/\ f = i + size t',
                                flatten t' = iota i (size t') &
                                sameShape t t'])
  := match t with
     | Leaf x => Do (n <-- fresh;
                     ret (Leaf n))
     | Node l r => Do (l' <-- relabel l;
                       r' <-- relabel r;
                       ret (Node l' r'))
     end.
Next Obligation. by []. Qed.
Next Obligation.
move=>_ _ _ i _ _ /= [_][_][[->->]][<- ->] /=; split=>//.
by rewrite addn1.
Qed.
Next Obligation. by []. Qed.
Next Obligation.
move=>/= A t l r i _ _ [tm][_][[-> Etm Esm][tn][_]][[-> Etn Esn][<- ->]] /=.
split; last by apply/andP.
- by rewrite addnA.
by rewrite iotaD Etm Etn.
Qed.

Program Fixpoint final {A} (t : tree A) :
   HoareState nat
              (@top nat)
              (tree nat)
              (fun _ t' _ => uniq (flatten t'))
   := do _ _ (relabel t).
Next Obligation.
move=>A t i x f /= [_ -> _].
by apply: iota_uniq.
Qed.