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
    end.

  Lemma iexpr_get_tag_correct:
    ∀ ρ e (te:tagged_iexpr e) n,
      ρ ∈ γ te ->
      eval_iexpr ρ e n ->
      n ∈ γ (iexpr_get_tag te).

  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 ρ)
          end
        in
        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
        in
        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
        in
        IEbinop_tag ab _ te1 te2
      end.
  Definition eval_expr_fuel := 5%nat.
  Opaque eval_expr_fuel.

  Lemma eval_expr_correct:
    ∀ fuel e ρ ab,
      ρ ∈ γ ab ->
      ρ ∈ γ (eval_expr fuel e ab).

  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
            end).

  Lemma assign_correct:
    ∀ x e ρ, ∀ n ab c,
      ρ ∈ γ ab -> (upd ρ x n) ∈ γ c ->
      n ∈ eval_iexpr ρ e ->
      (upd ρ x n) ∈ γ (assign x e ab c).

  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)
    end.

  Lemma var_refine_sound :
    ∀ ρ x abv ab, ρ x ∈ γ abv -> ρ ∈ γ ab -> ρ ∈ γ (var_refine x abv ab).

  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 =>
        backward_var
          (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 ρ)
      end
    | 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 ρ)
      end
    | 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 ρ)
        end;
      match ve2 with
      | Just ve2 => backward_eval' te2 (ρ', snd ρ) ve2
      | All => ret ρ'
      end
    end.
  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).

  Program Definition accumulate_assume_msgs (ρold ρnew:t): messages_chan :=
    let messages :=
      T.shcombine_diff
        (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
             end
           | Some vOld, Some vNew =>
             if vOld ⊑ vNew then None
             else
               match assume_msgs x (Just vOld) vNew return _ with
                 | nil => None
                 | x => Some x
               end
           end) _ ρold ρnew
    in
    T.fold (fun acc _ l => app l acc) messages nil.

  Lemma accumulate_assume_msgs_correct:
    forall ρold ρnew ρ,
      ρ ∈ γ ρold ->
      ρ ∈ γ ρnew ->
      ρ ∈ γ (accumulate_assume_msgs ρold ρnew).

  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
    end.

  Lemma assume'_correct:
    ∀ e ρ ab (b:bool), ρ ∈ γ ab ->
      (INz (if b then F1 else F0)) ∈ eval_iexpr ρ e ->
      ρ ∈ γ (assume' e b ab).

  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 ρ)
    end.

  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).

  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).

  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)
          end
        end
      in
      (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)
          end
        | _ =>
          ret (fst ab, nil)
        end;
      do (ab, c2) <-
        process_msg
          (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
  }.

End S.