Module AlarmMon

Require Import
  Utf8 String Coqlib Util AdomLib.

Set Implicit Arguments.

Definition alarm_mon (T:Type) : Type := (T*list (unit -> string))%type.

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

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

Instance alarm_monad : monad alarm_mon :=
  { bind A B x f :=
      let '(x, err) := x in
      match err with
      | nil => f x
      | _ =>
        let '(y, err') := f x in
        (y, err++err')
      end;
    ret A x := (x, nil)
  }.

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

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

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

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

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

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

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

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

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

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

Lemma am_rev_map_correct A B (f: B → alarm_mon 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 A B (f: B → alarm_mon 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 A B (f: B → am_map_t A) :
  ∀ l,
    match am_get (am_map' 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' ?f ?l ] =>
  let map_f := fresh "map_f" in
  remember f as map_f;
  generalize (am_map'_correct map_f l);
    destruct (am_map' map_f l) as (? & [|]);
    [subst map_f| easy ]
  end.