Require Import Facts States GCSemantics ICSemantics.
Set Implicit Arguments.

Module Compiler (Sigma : State).
Module GCSem := GCSemantics.GCSemantics Sigma.
Module ICSem := ICSemantics.ICSemantics Sigma.
Export ICSem.
Export GCSem.

Definition simple s := if s is Jump _ then true else false.

Definition LetDef (s : term) (t : {bind term}) : term :=
  if simple s then t.[s/] else Def s.[ren (+1)] t.

Definition flatmap (f : cmd → term) (u : term) : gc → term :=
  fix rec G := match G with
  | [::] ⇒ u
  | (b,s) :: G ⇒ If b (f s) (rec G)
  end.

Definition abort : term := Def (Jump 0) (Jump 0).

Fixpoint comp (u : term) (s : cmd) : term :=
  match s with
    | Skip ⇒ u
    | Assn a ⇒ Act a u
    | Seq s t ⇒ comp (comp u t) s
    | Case [::] ⇒ abort
    | Case ((_,s) :: G) ⇒ LetDef u (flatmap (comp (Jump 0)) (comp (Jump 0) s) G)
    | Do G ⇒ Def (flatmap (comp (Jump 0)) u.[ren (+1)] G) (Jump 0)
  end.
Local Notation compile := (comp (Jump 0)).

Lemma fix_weaken Q s x :
  Fix (fun P ⇒ wp (P .: Q) s.[ren (+1)]) x ↔ wp Q s x.

Lemma wp_let Q s t x :
  wp Q (LetDef s t) x ↔ wp (wp Q s .: Q) t x.

Section Flatten.

Variable (Q : NPred state) (f : cmd → term) (u : term) (x : state).
Implicit Types (G : gc).

Lemma flatmap_sound G :
  (∀ (b : guard) (s : cmd), (b, s) \in G → b x → wp Q (f s) x) →
  (~~G x → wp Q u x) →
  wp Q (flatmap f u G) x.

Lemma flatmap_soundT G :
  (∀ (b : guard) (s : cmd), (b, s) \in G → b x → wp Q (f s) x) →
  G x →
  wp Q (flatmap f u G) x.

Lemma flatmap_soundF G :
  ~~G x → wp Q u x → wp Q (flatmap f u G) x.

End Flatten.

Lemma compile_correct´ (P : Pred state) (x : state) (s : cmd) :
  wps P s x →
  ∀ (Q : NPred state) (u : term),
    P <<= wp Q u →
    wp Q (comp u s) x.

Theorem compile_correct (Q : Pred state) (s : cmd) :
  wps Q s <<= wp (Q .: @Bot _) (compile s).

End Compiler.