Require Import Facts States GCSemantics ICSemantics.
Set Implicit Arguments.
Unset Strict Implicit.

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.
Proof.
  split. apply: wp_fix. by rewrite wp_ren. move⇒ wps.
  apply: fix_fold. by rewrite wp_ren.
Qed.

Lemma wp_let Q s t x :
  wp Q (LetDef s t) x ↔ wp (wp Q s .: Q) t x.
Proof.
  rewrite/LetDef. case: ifP ⇒ _/=. rewrite wp_subst. by apply: wp_equiv ⇒ -[].
  apply: wp_equiv ⇒ -[]//=y. by rewrite fix_weaken.
Qed.

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.
Proof.
  elim: G ⇒ //=[|[e s]G ih]h1 h2/=. exact: h2. case: ifPn ⇒ p.
  - apply: h1 p. exact: mem_head.
  - apply: ih ⇒ [b t|]h3. apply: h1. exact: mem_behead. apply: h2.
      by rewrite gtest_cons negb_or p.
Qed.

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.
Proof.
  move⇒ h1 h2. apply: flatmap_sound h1 _. by rewrite h2.
Qed.

Lemma flatmap_soundF G :
  ~~G x → wp Q u x → wp Q (flatmap f u G) x.
Proof.
  move⇒ gt h. apply: flatmap_sound ⇒ // b s h1 h2. exfalso.
  move/gtestP: gt ⇒ []; eauto.
Qed.

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.
Proof.
  elim⇒ {P s x}/=[|||P [//|[e s]G]x h _ ih Q k hk|
    P G x I gt _ ih1 _ ih2 Q k hk|P G x gt Qa Q k hk]; eauto.
  - apply/wp_let. apply: flatmap_sound ⇒ [b t h1/ih{ih}ih|/negbTE hG].
    + apply: ih ⇒ //. by rewrite inE h1 orbT.
    + apply: (ih e s) ⇒ //=. by rewrite inE eqxx.
      move: h. by rewrite gtest_cons hG orbF.
  - apply: fix_fold. apply: flatmap_soundT gt ⇒ b s h1 h2.
    apply: ih1 h1 h2 _ _ _ ⇒ beta ib /=. exact: ih2.
  - apply: fix_fold. apply: flatmap_soundF ⇒ //.
    rewrite wp_weaken. exact: hk.
Qed.

Theorem compile_correct (Q : Pred state) (s : cmd) :
  wps Q s <<= wp (Q .: @Bot _) (compile s).
Proof. move⇒ x/compile_correct´. exact. Qed.

End Compiler.