Module AlarmMon

Require Import
  Utf8 String Coqlib Util AdomLib.

Set Implicit Arguments.

Definition alarm_mon (L T:Type) : Type := (T * (list (unit -> string) * list L))%type.
Definition am_lift {L T T'} (f: T → T') (m: alarm_mon L T) : alarm_mon L T' :=
  (f(fst m), snd m).

Instance block_mon_gamma L T T' `(gamma_op T T') : gamma_op (alarm_mon L T) (option T') :=
  λ t, match fst (snd t) with
       | nil => fun x => match x with Some x => x ∈ γ(fst t) | None => False end
       | _ => λ _, True
       end.

Definition am_get {L T} (x: alarm_mon L T) : option T :=
  match x with (y, (nil, _)) => Some y | _ => None end.

Instance alarm_monad L : monad (alarm_mon L) :=
  { bind A B x f :=
  let '(x, (err, log)) := x in
  let '(y, (err', log')) := f x in
  (y, (err ++ err', log ++ log'));
    ret A x := (x, (nil, nil))
  }.

Notation alarm_am s := (tt, ((fun _ => s)::nil, nil)).

Lemma bind_am_assoc {L A B C} (m: alarm_mon L A) (f: A → alarm_mon L B) (g: B → alarm_mon L C) :
  bind (bind m f) g = bind m (λ a, bind (f a) g).

Lemma am_bind_case {L T T'} (f: alarm_mon L T) (b: T → alarm_mon L T') :
  ∀ P : alarm_mon L T' → Prop,
    (∀ x y l l', P (x, (y, l)) → P (x, (y, l'))) →
    (∀ x y z l, P (x, (y :: z, l))) →
    (∀ x l, f = (x, (nil, l)) → P (b x)) →
    P (do a <- f; b a).

Lemma am_get_bind {L A B} (m: alarm_mon L A) (f: A → alarm_mon L B) bs :
  am_get (do a <- m; f a) = Some bs →
  ∃ x, am_get m = Some x ∧ am_get (f x) = Some bs.

Notation am_assert b err :=
  (if b then ret tt else alarm_am err).

Lemma am_assert_spec L b err B (k: alarm_mon L B) :
  ∀ P : _ → Prop,
    (∀ x y l l', P (x, (y, l)) → P (x, (y, l'))) →
    (∀ x y z l, P (x, (y :: z, l))) →
    (b = true → P k) →
    P (do _ <- am_assert b err; k).

Lemma am_assert_inv {L} (b: bool) err :
  am_assert b err = (tt, (nil, @nil L)) →
  b = true.

Lemma am_get_ret L A (a: A) : @am_get L A (ret a) = Some a.

Definition am_fold L A B (f: A → B → alarm_mon L A) (l: list B) (a: A) : alarm_mon L A :=
  fold_left (λ acc b, do acc' <- acc; f acc' b) l (ret a).

Lemma am_fold_app L A B f l l' a :
  @am_fold L A B f (l ++ l') a = do a' <- am_fold f l a; am_fold f l' a'.

Definition am_rev_map L A B (f: B → alarm_mon L A) (l: list B) : alarm_mon L (list A) :=
  am_fold (λ acc b, do a <- f b; ret (a :: acc)) l nil.

Inductive am_map_t (A: Type) : Type :=
| Result (a: A)
| Errorc (a: option A) (msg: unit -> string).
Notation Error a msg := (Errorc a (fun _ => msg)).
Lemma Result_Error {A} (a: A) a' m : Result a = Errorc a' m → ∀ P : Prop, P.
Proof (λ E, match E in _ = t return match t return Prop with Result _ => True | _ => _ end with eq_refl => I end).
Lemma Result_eq_inv {A} (a a': A) : Result a = Result a' → a = a'.
Proof (λ E, match E in _ = t return match t with Result x => a = x | _ => False end with eq_refl => eq_refl end).
Ltac result_eq_inv :=
  repeat match goal with
  | H : Result _ = Errorc _ _ |- _ => exact (Result_Error H _)
  | H : Errorc _ _ = Result _ |- _ => exact (Result_Error (eq_sym H) _)
  | H : Result _ = Result _ |- _ => apply Result_eq_inv in H
  end.

Definition am_map' L A B (f: B → am_map_t A) (l: list B) : alarm_mon L (list A) :=
  am_fold (λ acc b, match f b with
                      | Result a => ret (a :: acc)
                      | Errorc a msg =>
                        do _ <- (tt, (msg::nil, nil));
                        ret match a with
                                 | None => acc
                                 | Some a => a :: acc
                               end
                    end) l nil.

Lemma am_rev_map_correct L A B (f: B → alarm_mon L A) :
  ∀ l,
    match am_get (am_rev_map f l) with
    | Some m =>
      (∀ y, In y m → ∃ x, am_get (f x) = Some y ∧ In x l)
    ∧ (∀ x, In x l → ∃ y, am_get (f x) = Some y ∧ In y m)
    | None => ∃ x, In x l ∧ am_get (f x) = None
    end.

Lemma am_rev_map_inv L A B (f: B → alarm_mon L A) :
  ∀ l,
    match am_get (am_rev_map f l) with
    | Some m =>
      m = rev (map (fst ∘ f) l)
    | None => True
    end.

Lemma am_map'_correct L A B (f: B → am_map_t A) :
  ∀ l,
    match am_get (am_map' L f l) with
    | Some m =>
      (∀ y, In y m → ∃ x, f x = Result y ∧ In x l)
    ∧ (∀ x, In x l → ∃ y, f x = Result y ∧ In y m)
    | None => True
    end.

Ltac with_am_map' :=
  match goal with
  |- appcontext[ am_map' ?L ?f ?l ] =>
  let map_f := fresh "map_f" in
  remember f as map_f;
  generalize (am_map'_correct L map_f l);
    destruct (am_map' L map_f l) as (? & ([|], ?));
    [subst map_f| easy ]
  end.