Module IdealBoxDomain

Require Import Coqlib Utf8 FastArith.
Require Import Maps ShareTree Util.
Require Import AdomLib.
Require Import IdealExpr.
Require Import AbIdealEnv.
Require Import AbIdealNonrel.
Require Import IdealIntervals.
Require Import FactMsgHelpers.

Module Dom := WPFun T.

Section S.

  Context {Val: Type} {Dval: ab_ideal_nonrel Val}.
  Context (G:gamma_op Val ideal_num) (Dval_cor: ab_ideal_nonrel_correct Dval G).

  Definition t := T.t Val.

  Instance leb_tree : leb_op t := Dom.leb_tree _ leb_refl.
  Instance join_tree : join_op t t := Dom.join_tree _ join_idem.
  Definition widen_tree : t -> t -> t := Dom.widen_tree _ widen_idem.

  Inductive tagged_iexpr : iexpr -> Type :=
  | IEvar_tag : Val+⊤+⊥ -> ∀ ty x, tagged_iexpr (IEvar ty x)
  | IEconst_tag : Val+⊤+⊥ -> ∀ n, tagged_iexpr (IEconst n)
  | IEZitv_tag : Val+⊤+⊥ -> ∀ i, tagged_iexpr (IEZitv i)
  | IEunop_tag : Val+⊤+⊥ -> ∀ op {e}, tagged_iexpr e -> tagged_iexpr (IEunop op e)
  | IEbinop_tag : Val+⊤+⊥ -> ∀ op {e1 e2}, tagged_iexpr e1 -> tagged_iexpr e2 ->
                                       tagged_iexpr (IEbinop op e1 e2).

  Inductive tagged_iexpr_gamma:
    ∀ {e}, gamma_op (tagged_iexpr e) (var -> ideal_num) :=
  | IEvar_tag_ok: ∀ ab x ty ρ,
    ρ x ∈ γ ab ->
    tagged_iexpr_gamma (IEvar_tag ab ty x) ρ
  | IEZconst_tag_ok: ∀ ab n ρ,
    n ∈ γ ab ->
    tagged_iexpr_gamma (IEconst_tag ab n) ρ
  | IEZitv_tag_ok: ∀ ab itv ρ,
    (∀ (x:Z), x ∈ γ itv -> INz x ∈ γ ab) ->
    tagged_iexpr_gamma (IEZitv_tag ab itv) ρ
  | IEunop_tag_ok: ∀ ab op e (te:tagged_iexpr e) ρ,
    (∀ r, eval_iexpr ρ (IEunop op e) r -> r ∈ γ ab) ->
    tagged_iexpr_gamma te ρ ->
    tagged_iexpr_gamma (IEunop_tag ab op te) ρ
  | IEbinop_tag_ok: ∀ ab op e1 (te1:tagged_iexpr e1) e2 (te2:tagged_iexpr e2) ρ,
    (∀ r, eval_iexpr ρ (IEbinop op e1 e2) r -> r ∈ γ ab) ->
    tagged_iexpr_gamma te1 ρ ->
    tagged_iexpr_gamma te2 ρ ->
    tagged_iexpr_gamma (IEbinop_tag ab op te1 te2) ρ.

  Existing Instance tagged_iexpr_gamma.

  Definition iexpr_get_tag {e:iexpr} (te:tagged_iexpr e) : Val+⊤+⊥ :=
    match te with
    | IEvar_tag ab _ _ => ab
    | IEconst_tag ab _ => ab
    | IEZitv_tag ab _ => ab
    | IEunop_tag ab _ _ _ => ab
    | IEbinop_tag ab _ _ _ _ _ => ab

  Lemma iexpr_get_tag_correct:
    ∀ ρ e (te:tagged_iexpr e) n,
      ρ ∈ γ te ->
      eval_iexpr ρ e n ->
      n ∈ γ (iexpr_get_tag te).
    destruct 1; intros EVAL; inv EVAL; eauto.

  Fixpoint eval_expr (fuel:nat) :
    ∀ e:iexpr, t * query_chan -> tagged_iexpr e :=
    fix eval_expr' (e:iexpr) (ρ:t * query_chan) : tagged_iexpr e :=
      match e with
      | IEvar _ x =>
        let ab := Dom.get x (fst ρ) in
        let ab :=
          match fuel with
            | O => NotBot ab
            | S fuel =>
              forward_var (fun e => iexpr_get_tag (eval_expr fuel e ρ))
                          x ab (snd ρ)
        IEvar_tag ab _ _
      | IEconst n => IEconst_tag (NotBot (const n)) _
      | IEZitv itv => IEZitv_tag (NotBot (z_itv itv)) _
      | IEunop op e =>
        let te := eval_expr' e ρ in
        let ab :=
          do ab <- iexpr_get_tag te;
          forward_unop op ab
        IEunop_tag ab _ te
      | IEbinop op e1 e2 =>
        let te1 := eval_expr' e1 ρ in
        let te2 := eval_expr' e2 ρ in
        let ab :=
          do ab1 <- iexpr_get_tag te1;
          do ab2 <- iexpr_get_tag te2;
          forward_binop op ab1 ab2
        IEbinop_tag ab _ te1 te2
  Definition eval_expr_fuel := 5%nat.
  Opaque eval_expr_fuel.

  Lemma eval_expr_correct:
    ∀ fuel e ρ ab,
      ρ ∈ γ ab ->
      ρ ∈ γ (eval_expr fuel e ab).
    fix 1.
    intros fuel e ρ ab Hρ. induction e.
    - destruct fuel; constructor. apply Hρ.
      eapply forward_var_sound; try apply Hρ; auto.
      repeat intro. eapply iexpr_get_tag_correct; eauto.
    - destruct fuel; constructor; try apply const_correct; auto.
    - destruct fuel; constructor; try constructor; apply z_itv_correct.
    - destruct fuel; constructor; eauto; intros; inv H; eapply botbind_spec;
      eauto using iexpr_get_tag_correct, forward_unop_sound.
    - destruct fuel; constructor; eauto; intros; inv H; (eapply botbind_spec; [intros; eapply botbind_spec|]);
      eauto using iexpr_get_tag_correct, forward_binop_sound.

  Definition assign (x:var) (e:iexpr) (ρ:t * query_chan) (chan:query_chan) : (t*messages_chan)+⊥ :=
    do abv <- iexpr_get_tag (eval_expr eval_expr_fuel e ρ);
    ret (Dom.set (fst ρ) x abv,
            match abv with
            | Just abv => assign_msgs x abv
            | All => nil

  Lemma assign_correct:
    ∀ x e ρ, ∀ n ab c,
      ρ ∈ γ ab -> (upd ρ x n) ∈ γ c ->
      neval_iexpr ρ e ->
      (upd ρ x n) ∈ γ (assign x e ab c).
    intros x e ρ n ab c Hρ Hc Hn.
    eapply botbind_spec, iexpr_get_tag_correct; eauto. 2:apply eval_expr_correct; auto.
    - intro. simpl. rewrite Dom.gsspec.
      unfold upd. unfold eq_dec, var_eq. destruct T.elt_eq. auto. apply Hρ.
    - destruct a. constructor. apply assign_msgs_sound.
      rewrite upd_s. auto.

  Definition var_refine x : var_refiner t := fun abx abenv =>
    match T.get_set x abenv with
    | (None, setter) => ret (setter abx)
    | (Some abx', setter) =>
      do abx <- meet abx abx';
      ret (setter abx)

  Lemma var_refine_sound :
    ∀ ρ x abv ab, ρ x ∈ γ abv -> ρ ∈ γ ab -> ρ ∈ γ (var_refine x abv ab).
    intros. unfold var_refine.
    pose proof T.get_set_spec x ab. destruct T.get_set.
    destruct H1. simpl in H1. subst.
    change (∀ v, t0 v = Dom.set ab x (Just v)) in H2.
    pose proof (H0 x). unfold Dom.get in H1. destruct T.get.
    - eapply botbind_spec withx). intros. rewrite H2. eapply Dom.gamma_set; eauto.
      apply meet_correct, conj; auto.
    - rewrite H2. eapply Dom.gamma_set; eauto.

  Fixpoint backward_eval (fuel:nat) :
    ∀ {e:iexpr}, tagged_iexpr e -> t * query_chan -> Val -> t+⊥ :=
    fix backward_eval' {e:iexpr} (te:tagged_iexpr e) (ρ:t * query_chan) (v:Val) : t+⊥ :=
    do abres <- meet (NotBot (Just v)) (iexpr_get_tag te);
    match te with
    | IEvar_tag abx _ x =>
      match abres, fuel with
      | Just abres, S fuel =>
          (fun e => iexpr_get_tag (eval_expr fuel e ρ))
          (fun e abe abenv => backward_eval fuel (eval_expr fuel e ρ) abenv abe)
          (var_refine x)
          x abres ρ
      | Just abres, _ => var_refine x abres (fst ρ)
      | All, _ => ret (fst ρ)
    | IEconst_tag _ _ | IEZitv_tag _ _ =>
      ret (fst ρ)
    | IEunop_tag _ op _ te' =>
      do ve' <- iexpr_get_tag te';
      do ve' <- backward_unop op abres ve';
      match ve' with
      | Just ve' => backward_eval' te' ρ ve'
      | All => ret (fst ρ)
    | IEbinop_tag _ op _ _ te1 te2 =>
      do ve1 <- iexpr_get_tag te1;
      do ve2 <- iexpr_get_tag te2;
      do (ve1, ve2) <- backward_binop op abres ve1 ve2;
      do ρ' <-
        match ve1 with
        | Just ve1 => backward_eval' te1 ρ ve1
        | All => ret (fst ρ)
      match ve2 with
      | Just ve2 => backward_eval' te2 (ρ', snd ρ) ve2
      | All => ret ρ'
  Definition backward_eval_fuel := 5%nat.

  Lemma backward_eval_correct:
    ∀ fuel e (te:tagged_iexpr e) res res_ab,
      res ∈ γ res_ab ->
    ∀ ρ ρ_ab,
      ρ ∈ γ ρ_ab ->
      ρ ∈ γ te ->
      eval_iexpr ρ e res ->
      ρ ∈ γ (backward_eval fuel te ρ_ab res_ab).
    fix 1.
    induction te; intros res res_ab Hres ρ ρ_ab Hρ Hte EVAL; inversion EVAL; subst; clear EVAL.
    - assert (Hρ':=Hρ). destruct Hρ' as [Hρ0 Hρ1].
      destruct fuel; (eapply botbind_spec; [|apply meet_correct; eauto using iexpr_get_tag_correct]);
      intros [|abx] Habx; auto using var_refine_sound.
      apply backward_var_sound; eauto.
      + repeat intro. eapply iexpr_get_tag_correct, H. apply eval_expr_correct; eauto.
      + repeat intro. eapply backward_eval_correct; eauto. apply eval_expr_correct; eauto.
      + repeat intro. apply var_refine_sound; auto.
    - assert (Hρ':=Hρ). destruct Hρ' as [Hρ0 Hρ1].
      destruct fuel; auto; (eapply botbind_spec; [auto|]);
      apply meet_correct; eauto using iexpr_get_tag_correct, const_correct.
    - assert (Hρ':=Hρ). destruct Hρ' as [Hρ0 Hρ1].
      destruct fuel; auto; eapply botbind_spec; auto;
      apply meet_correct; eauto using iexpr_get_tag_correct.
    - pose proof Hte. inv Hte. apply Classical_Prop.EqdepTheory.inj_pair2 in H6. subst te0.
      destruct fuel;
      (repeat (eapply botbind_spec; [intros|]);
        [destruct a1; [apply Hρ|eapply IHte; eauto]|
         eapply backward_unop_sound; eauto |
         eapply iexpr_get_tag_correct; eauto |
         eapply meet_correct; eauto |
         eapply iexpr_get_tag_correct; eauto ]).
    - pose proof Hte. inv Hte.
      apply Classical_Prop.EqdepTheory.inj_pair2 in H8. subst te0.
      apply Classical_Prop.EqdepTheory.inj_pair2 in H9. subst te3.
      destruct fuel;
      (repeat (eapply botbind_spec; [intros; try destruct a2|]);
      [| |eapply backward_binop_sound; eauto|
       eapply iexpr_get_tag_correct; eauto|
       eapply iexpr_get_tag_correct; eauto|
       eapply meet_correct; eauto|
       eapply iexpr_get_tag_correct; eauto];
      destruct H6; [destruct t1; [|eapply IHte2]|destruct t0; [|eapply IHte1]];
      try split; auto; apply Hρ.

  Program Definition accumulate_assume_msgsold ρnew:t): messages_chan :=
    let messages :=
        (fun x vOld vNew =>
           match vOld, vNew return _ with
           | _, None => None
           | None, Some vNew =>
             match assume_msgs x All vNew return _ with
             | nil => None
             | x => Some x
           | Some vOld, Some vNew =>
             if vOldvNew then None
               match assume_msgs x (Just vOld) vNew return _ with
                 | nil => None
                 | x => Some x
           end) _ ρold ρnew
    T.fold (fun acc _ l => app l acc) messages nil.
Next Obligation.
destruct v; auto. rewrite leb_refl; auto. Qed.

  Lemma accumulate_assume_msgs_correct:
    forall ρold ρnew ρ,
      ρ ∈ γ ρold ->
      ρ ∈ γ ρnew ->
      ρ ∈ γ (accumulate_assume_msgs ρold ρnew).
    unfold accumulate_assume_msgs.
    apply Dom.P.fold_rec. auto. constructor. intros.
    assertv ρ).
    { rewrite T.gshcombine_diff in H2.
      specialize (H k). specialize (H0 k). unfold Dom.get in H, H0.
      destruct T.get, T.get; try discriminate.
      - destruct (v0v1). discriminate.
        pose proof assume_msgs_sound k (Just v0) v1 ρ.
        destruct (assume_msgs k (Just v0) v1); inv H2. auto.
      - pose proof assume_msgs_sound k All v0 ρ.
        destruct (assume_msgs k All v0); inv H2. auto. }
    clear H2. induction v.
    auto. inv H4. constructor. auto. apply IHv. auto.

  Definition assume' (e:iexpr) (b:bool) (ρ:t * query_chan): t+⊥ :=
    let n := INz (if b then F1 else F0) in
    match const n with
    | All => NotBot (fst ρ)
    | Just ab =>
      backward_eval backward_eval_fuel (eval_expr eval_expr_fuel e ρ)
                    ρ ab

  Lemma assume'_correct:
    ∀ e ρ ab (b:bool), ρ ∈ γ ab ->
      (INz (if b then F1 else F0)) ∈ eval_iexpr ρ e ->
      ρ ∈ γ (assume' e b ab).
    intros. unfold assume'.
    destruct (const (INz (if b then F1 else F0))) eqn:EQ.
    apply H.
    eapply backward_eval_correct; eauto.
    change (γ (Just v) (INz (if b then F1 else F0))). rewrite <- EQ.
    apply const_correct. apply eval_expr_correct; auto.

  Fixpoint assume_rec (fuel: nat) (e: iexpr) (b: bool) (ρ: t * query_chan): t+⊥ :=
    let ρ' := assume' e b ρ in
    match fuel with
    | O => ρ'
    | S fuel' =>
      do ρ' <- ρ';
      if fst ρ ⊑ ρ'
      then NotBot (fst ρ)
      else assume_rec fuel' e b (ρ', snd ρ)

  Parameter assume_fuel: unit -> nat.

  Definition assume := assume_rec (assume_fuel tt).

  Lemma assume_correct:
    ∀ e ρ ab b, ρ ∈ γ ab ->
                (INz (if b:bool then 1 else 0)) ∈ eval_iexpr ρ e ->
                ρ ∈ γ (assume e b ab).
    unfold assume. elim (assume_fuel tt).
    apply assume'_correct.
    intros n H e ρ ab b H0 H1. simpl.
    eapply botbind_spec. 2: eauto using assume'_correct.
    intros a H2.
    bif. exact (proj1 H0).
    apply H; eauto.
    split. exact H2. exact (proj2 H0).

  Definition forget (x:var) (ρ:t * query_chan) (c:query_chan): t+⊥ :=
    NotBot (Dom.set (fst ρ) x All).

  Lemma forget_correct:
    ∀ x ρ, ∀ n ab c,
      ρ ∈ γ ab -> (upd ρ x n) ∈ γ c ->
      (upd ρ x n) ∈ γ (forget x ab c).
    intros x ρ n ab c Hρ Hc y. unfold upd.
    rewrite Dom.gsspec. unfold eq_dec, var_eq. destruct T.elt_eq; auto. exact I.
    destruct Hρ. auto.

  Instance ideal_box_domain : ab_ideal_env t (t+⊥) := {
    id_top c := NotBot (top, nil);
    id_leb x y := leb (fst x) y;
    id_join x y c := NotBot (join (fst x) (fst y), nil);
    id_init_iter := id;
    id_widen x y c :=
      let w :=
        match x with
        | Bot => fmap fst y
        | NotBot x =>
          match y with
          | Bot => ret x
          | NotBot (y, _) => ret (widen_tree x y)
      (w, do w <- w; NotBot (w, nil));
    assign x e ab c := assign x e ab c;
    forget x ab c := fmap (fun x => (x, nil)) (forget x ab c);
    process_msg m ab :=
      do (ab', c1) <-
        match m return (t*messages_chan)+⊥ with
        | Fact_msg e b =>
          do res <- assume_unfold (fun e b ab => assume e b ab) e b ab;
          let c := accumulate_assume_msgs (fst ab) res in
          ret (res, c)
        | Known_value_msg x v =>
          match const v with
          | Just v =>
            do ab' <- var_refine x v (fst ab);
            ret (ab', nil)
          | All => ret (fst ab, nil)
        | _ =>
          ret (fst ab, nil)
      do (ab, c2) <-
          (fun e => iexpr_get_tag (eval_expr 0 e (ab', snd ab)))
          (fun e ab' abenv => backward_eval 0 (eval_expr 0 e abenv) abenv ab')
          m (ab', snd ab);
      NotBot (ab, c1++c2);
    enrich_query_chan ab c :=
      enrich_query_chan (fun e => iexpr_get_tag (eval_expr eval_expr_fuel e (ab, c))) c
    - split. intro. unfold Dom.get. rewrite T.gempty. constructor. constructor.
    - intros. intro. destruct H0. eapply leb_correct in H0. 2:apply H.
      specialize (H0 e). unfold flat_top_leb in H. auto.
    - intros. split. apply join_correct. destruct H as [[]|[]]; auto. constructor.
    - auto.
    - intros. destruct ab1. contradiction.
      destruct ab2. split. auto. eapply botbind_spec with ρ.
      intros. split. auto. constructor. auto.
      destruct x0 as [y ?].
      unfold widen_tree.
      pose proof Dom.widen_tree_incr widen widen_idem AbIdealNonrel.widen_incr.
      specialize (H1 x y _ H). split. apply H1. split. apply H1. constructor.
    - eapply assign_correct.
    - intros. eapply botbind_spec, forget_correct; eauto. split. auto. constructor.
    - intros. repeat (eapply botbind_spec with ρ; [intros [] ?|]).
      + split. apply H2.
        destruct H1, H2. simpl in *. induction l. auto. inv H3. constructor.
        auto. apply IHl, H8.
      + destruct H1, H. apply process_msg_sound; repeat intro; auto.
        eapply iexpr_get_tag_correct; eauto. apply eval_expr_correct; split; auto.
        eapply backward_eval_correct; eauto. apply eval_expr_correct. auto.
        split; auto.
      + destruct m; try (eapply botbind_spec; intros); repeat constructor; try apply H.
        apply accumulate_assume_msgs_correct. apply H. auto.
        apply assume_unfold_correct; auto. intros. apply join_correct; auto.
        intros. apply assume_correct; auto.
        inv H0. pose proof const_correctv). destruct const. split. apply H. constructor.
        eapply botbind_spec, var_refine_sound, H. split. auto. constructor. auto.
    - intros. apply enrich_query_chan_sound; auto.
      intros e ve He. eapply iexpr_get_tag_correct; eauto. eapply eval_expr_correct; split; eauto.

End S.