Require Import base ars ord wp.states wp.gc wp.ic tower.tarski
        tower.direct_induction wp.prelim data.fintype wp.gcsemantics wp.icsemantics.

(* move to tower.tarski *)
Global Instance lfp_mono {X : clat} : monotone (@lfp X).
Proof.
  move=> f g le_f_g. rewrite/lfp/lfp_above. focus=> y [_ h].
  apply: allEc (y) _ le_refl. split=>//. by rewrite le_f_g.
Qed.

Global Instance gfp_mono {X : clat} : monotone (@gfp X).
Proof.
  move=> f g le_f_g. rewrite/gfp/gfp_below. focus=> y [_ h].
  apply: exIc (y) _ le_refl. split=>//. rewrite {1}h. exact: le_f_g.
Qed.

Opaque lfp gfp.

Definition simple {n} (s : tm n) := if s is Jump _ then true else false.

Definition LetDef {n} (s : tm n) (t : tm n.+1) : tm n :=
  if simple s then inst (s .: Jump) t else Def (rinst shift s) t.

Lemma lfp_const {X : clat} (f : X -> X) y :
  (forall x, f x = y) -> lfp f = y.
Proof.
  move=> h. apply: le_eq. apply: lfp_above_ind => //. by rewrite h.
  rewrite -(h (lfp f)). apply: lfp_above_fold. move=> a b _. by rewrite !h.
Qed.

Lemma gfp_const {X : clat} (f : X -> X) y :
  (forall x, f x = y) -> gfp f = y.
Proof.
  move=> h. apply: le_eq. rewrite -(h (gfp f)). apply: gfp_below_unfold.
  move=> a b _. by rewrite !h. apply: gfp_below_coind => //. by rewrite h.
Qed.

Lemma wpi_let n (s : tm n) t Q :
  wpi (LetDef s t) Q = wpi t (wpi s Q .: Q).
Proof.
  rewrite/LetDef. case: ifP => _. by rewrite wpi_beta.
  rewrite wpi_def. f_equal; f_equal; apply: lfp_const => P.
    by rewrite wpi_weaken.
Qed.

Lemma wlpi_let n (s : tm n) t Q :
  wlpi (LetDef s t) Q = wlpi t (wlpi s Q .: Q).
Proof.
  rewrite/LetDef. case: ifP => _. by rewrite wlpi_beta.
  rewrite wlpi_def. f_equal; f_equal; apply: gfp_const => P.
    by rewrite wlpi_weaken.
Qed.

Definition flatmap {n} (f : cmd -> tm n) (u : tm n) : gc -> tm n :=
  fix rec G := match G with
  | [::] => u
  | (b,s) :: G => If b (f s) (rec G)
  end.

Fixpoint comp {n} (u : tm n) (s : cmd) : tm n :=
  match s with
    | Skip => u
    | Assn a => Act a u
    | Seq s t => comp (comp u t) s
    | Case [::] => Ω (* dummy case *)
    | Case ((_,s) :: G) => LetDef u (flatmap (comp (Jump bound)) (comp (Jump bound) s) G)
    | Do G => Def (flatmap (comp (Jump bound)) (rinst shift u) G) (Jump bound)
  end.
Local Notation compile := (comp (Jump bound)).

Section CompilerCorrectness.

Lemma gtest_cons (G : gc) b s x :
  gtest ((b,s) :: G) x = b x || G x.
Proof. by []. Qed.

Lemma gtestP (G : gc) x :
  reflect (exists (b:guard) (s:cmd), (b,s) \in G /\ b x) (G x).
Proof.
  apply: (iffP hasP) => [[b/mapP[[b' s]/=mem->bx]]|[b[s[mem bx]]]].
  by exists b', s. exists b => //. apply/mapP. by exists (b,s).
Qed.

Lemma gtest_contra (G : gc) b s x :
  (b,s) \in G -> ~~G x -> ~~b x.
Proof.
  move=> mem. apply/contra => bx. apply/gtestP. by exists b, s.
Qed.

Section Flatten.
Context {n : nat} (Q : fin n -> Pred state) (f : cmd -> tm n) (u : tm n) (x : state).
Implicit Types (G : gc).

Lemma wpi_flatmap G :
  (forall (b : guard) (s : cmd), (b, s) \in G -> b x -> wpi (f s) Q x) ->
  (~~G x -> wpi u Q x) ->
  wpi (flatmap f u G) Q x.
Proof.
  elim: G => //[|[e s]G ih]h1 h2. exact: h2. rewrite/= wpi_if. 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 wpi_flatmapT G :
  (forall (b : guard) (s : cmd), (b, s) \in G -> b x -> wpi (f s) Q x) ->
  G x ->
  wpi (flatmap f u G) Q x.
Proof.
  move=> h1 h2. apply: wpi_flatmap h1 _. by rewrite h2.
Qed.

Lemma wpi_flatmapF G :
  ~~G x -> wpi u Q x -> wpi (flatmap f u G) Q x.
Proof.
  move=> gt h. apply: wpi_flatmap => // b s h1 h2. exfalso.
  move/gtestP: gt => []; eauto.
Qed.

Lemma wlpi_flatmap G :
  (forall (b : guard) (s : cmd), (b, s) \in G -> b x -> wlpi (f s) Q x) ->
  (~~G x -> wlpi u Q x) ->
  wlpi (flatmap f u G) Q x.
Proof.
  elim: G => //[|[e s]G ih]h1 h2. exact: h2. rewrite/= wlpi_if. 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 wlpi_flatmapT G :
  (forall (b : guard) (s : cmd), (b, s) \in G -> b x -> wlpi (f s) Q x) ->
  G x ->
  wlpi (flatmap f u G) Q x.
Proof.
  move=> h1 h2. apply: wlpi_flatmap h1 _. by rewrite h2.
Qed.

Lemma wlpi_flatmapF G :
  ~~G x -> wlpi u Q x -> wlpi (flatmap f u G) Q x.
Proof.
  move=> gt h. apply: wlpi_flatmap => // b s h1 h2. exfalso.
  move/gtestP: gt => []; eauto.
Qed.
End Flatten.

Lemma wpi_comp {n} (s : cmd) (u : tm n) Q :
  wp s (wpi u Q) ≤ wpi (comp u s) Q.
Proof.
  elim: s n Q u => [|a|s t ihs iht|G ih|G ih] n Q u.
  - by rewrite wp_skip /=.
  - by rewrite wp_act wpi_act.
  - rewrite wp_seq /=. by rewrite iht ihs.
  - rewrite wp_case; hnf=> x. rewrite prod_meetE prop_meetE. hnf; case=> gx h.
    cbn. destruct G as [|[e s]G] => //=. rewrite wpi_let.
    apply: wpi_flatmap => //[b t h1 h2|/negbTE h1].
    + apply: (@ih b) => //. by rewrite inE h1 orbT. rewrite wpi_jump /=.
      apply: h h2. by rewrite inE/=h1 orbT.
    + apply: (ih e). by rewrite inE eqxx. rewrite wpi_jump /=.
      move: gx. rewrite/Gtest gtest_cons h1 orbF. apply: h. by rewrite inE eqxx.
  - rewrite wp_do /= wpi_def wpi_jump /=. apply: lfp_mono => P. hnf=> x; hnf=> h.
    apply: wpi_flatmap.
    + move=> b s mem bx. apply: (ih _ _ mem). rewrite wpi_jump.
      move: h. rewrite/wp_loop ifT. apply; eassumption. apply/gtestP.
        by exists b, s.
    + move=> gx. rewrite wpi_weaken. move: h. by rewrite/wp_loop ifN.
Qed.

Lemma wlpi_comp {n} (s : cmd) (u : tm n) Q :
  wlp s (wlpi u Q) ≤ wlpi (comp u s) Q.
Proof.
  elim: s n Q u => [|a|s t ihs iht|G ih|G ih] n Q u.
  - by rewrite wlp_skip /=.
  - by rewrite wlp_act wlpi_act.
  - rewrite wlp_seq /=. by rewrite iht ihs.
  - rewrite wlp_case; hnf=> x. rewrite prod_meetE prop_meetE. hnf; case=> gx h.
    cbn. destruct G as [|[e s]G] => //=. rewrite wlpi_let.
    apply: wlpi_flatmap => //[b t h1 h2|/negbTE h1].
    + apply: (@ih b) => //. by rewrite inE h1 orbT. rewrite wlpi_jump /=.
      apply: h h2. by rewrite inE/=h1 orbT.
    + apply: (ih e). by rewrite inE eqxx. rewrite wlpi_jump /=.
      move: gx. rewrite/Gtest gtest_cons h1 orbF. apply: h. by rewrite inE eqxx.
  - rewrite wlp_do /= wlpi_def wlpi_jump /=. apply: gfp_mono => P. hnf=> x; hnf=> h.
    apply: wlpi_flatmap.
    + move=> b s mem bx. apply: (ih _ _ mem). rewrite wlpi_jump.
      move: h. rewrite/wlp_loop ifT. apply; eassumption. apply/gtestP.
        by exists b, s.
    + move=> gx. rewrite wlpi_weaken. move: h. by rewrite/wlp_loop ifN.
Qed.

End CompilerCorrectness.

Theorem wpi_compile_correct n (Q : Pred state) (s : cmd) :
  wp s Q ≤ wpi (compile s : tm n.+1) (Q .: ⊥).
Proof. by rewrite -wpi_comp wpi_jump. Qed.

Theorem wlpi_compile_correct n (Q : Pred state) (s : cmd) :
  wlp s Q ≤ wlpi (compile s : tm n.+1) (Q .: ⊥).
Proof. by rewrite -wlpi_comp wlpi_jump. Qed.