include "utilities/extralib.ma".
include "common/Errors.ma".

(* IO monad *)

inductive IO (output:Type[0]) (input:output → Type[0]) (T:Type[0]) : Type[0] ≝
| Interact : ∀o:output. (input o → IO output input T) → IO output input T
| Value : T → IO output input T
| Wrong : errmsg → IO output input T.

include "utilities/proper.ma".
(* a weak form of extensionality *)
axiom interact_proper :
  ∀O,I,T,o.∀f,g : I o → IO O I T.(∀i. f i = g i) → Interact … o f = Interact … o g.

let rec bindIO (O:Type[0]) (I:O → Type[0]) (T,T':Type[0]) (v:IO O I T) (f:T → IO O I T') on v : IO O I T' ≝
match v with
[ Interact out k ⇒ (Interact ??? out (λres. bindIO O I T T' (k res) f))
| Value v' ⇒ (f v')
| Wrong m ⇒ Wrong O I T' m
].

include "utilities/monad.ma".

definition IOMonad ≝ λO,I.
  MakeMonadProps
  (* the monad *)
  (IO O I)
  (* return *)
  (λX.Value … X)
  (* bind *)
  (bindIO O I)
  ????. / by /
[(* bind_ret *)
 #X#m elim m normalize // #o#f#Hi @interact_proper //
|(* bind_bind *)
 #X#Y#Z#m#f#g elim m normalize [2,3://]
 (* Interact *)
  #o#f #Hi @interact_proper //
|#X #Y #m #f #g #H elim m normalize
  [ #o #x @interact_proper ] //
]
qed.

definition bindIO2 ≝ λO,I. m_bind2 (IOMonad O I).

include "hints_declaration.ma".
alias symbol "hint_decl" (instance 1) = "hint_decl_Type1".
unification hint 0 ≔ O, I, T;
  N ≟ IOMonad O I, M ≟ max_def N
(*******************************************) ⊢
  IO O I T ≡ monad M T
.

let rec rel_io O I A B (Perr : relation errmsg) (P : A → B → Prop) (v1 : IO O I A)
  (v2 : IO O I B) on v1 : Prop ≝
  match v1 with
  [ Value x ⇒
    match v2 with
    [ Value y ⇒
      P x y
    | _ ⇒ False
    ]
  | Wrong e1 ⇒
    match v2 with
    [ Wrong e2 ⇒ Perr e1 e2
    | _ ⇒ False
    ]
  | Interact o1 f1 ⇒
    match v2 with
    [ Interact o2 f2 ⇒
      ∃prf:o1 = o2.∀i.rel_io … Perr P (f1 i) (f2 ?)
    | _ ⇒ False
    ]
  ]. <prf @i qed.

definition IORel ≝ λO,I,Perr.
  mk_MonadRel (IOMonad O I) (IOMonad O I)
    (λA,B.rel_io O I A B Perr)
    ???.
[//
|#X #Y #Z #W #relin #relout #m elim m
  [ #o #f #IH * [#o' #f' | #v | #e] * #EQ destruct(EQ) #G
    #f1 #f2 #G' whd %{(refl …)} #i
    @(IH … (G i) f1 f2 G')
  | #v * [#o #f * | #v' | #e *]
    #Pvv' #f #g #H normalize @H @Pvv'
  | #e1 * [#o #f * | #v' *| #e2] //
  ]
| #X #Y #P #Q #H #m elim m [#o #f #IH | #v | #e] *
  [1,4,7: #o' #f' [2,3: *]
    normalize * #prf destruct(prf) normalize #G
    % [%] normalize #i @IH @G
  |2,5,8: #v' [1,3: *]
    @H
  |*: #e' [1,2: *] //
  ]
]
qed.

lemma IORel_refl : ∀O,I,Perr.reflexive ? Perr → ∀X,rel.reflexive X rel →
  reflexive ? (m_rel ?? (IORel O I Perr) ?? rel).
#O #I #Perr #G #X #rel #H #m elim m
[ #o #f #IH whd %[%] #i normalize @IH
| #v @H
| #e @G
]
qed.

lemma IORel_transitive : ∀O,I,Perr1,Perr2,Perr3.
  (∀e1,e2,e3.Perr1 e1 e2 → Perr2 e2 e3 → Perr3 e1 e3) →
  ∀X,Y,Z,rel1,rel2,rel3.
  (∀x : X.∀y : Y.∀z : Z.rel1 x y → rel2 y z → rel3 x z) →
  ∀m,n,o.
  m_rel ?? (IORel O I Perr1) … rel1 m n →
  m_rel ?? (IORel O I Perr2) … rel2 n o →
  m_rel ?? (IORel O I Perr3) … rel3 m o.
#O #I #Perr1 #Perr2 #Perr3 #Herr #X #Y #Z #rel1 #rel2 #rel3 #H #m elim m
[ #o #f #IH * [#o' #f' * [#o'' #f'' | #v #_ * | #e #_ * ] | #v #x * | #e #x * ]
  normalize * #EQ #H1 * #EQ' #H2 destruct %[%] normalize #i @(IH ? (f' i)) //
| #v * [#o #f #x * | #v' * [#o #f #_ * | #v'' |#e #_ *] | #e #x *]
  @H
| #e * [#o #f #x * | #v #x * | #e' * [#o #f #_ * | #v #_ * | #e'']] @Herr
]
qed.


lemma rel_io_eq : ∀O,I,A.∀m,n:IO O I A.m = n → m_rel ?? (IORel O I (eq ?)) … (eq ?) m n.
#O#I#A#m#n#EQ >EQ @IORel_refl //
qed.

lemma eq_rel_io : ∀O,I,A.∀m,n:IO O I A.m_rel ?? (IORel O I (eq ?)) … (eq ?) m n → m = n.
#O#I#A#m elim m
[ #o #f #IH * [#o' #f' | #v * | #e * ]
  normalize * #EQ #H destruct @interact_proper #i @IH @H
| #v * [#o #f * | #v' | #e *]
| #e * [#o #f * | #v * | #e']
] #EQ >EQ %
qed.

coercion rel_io_to_eq nocomposites : ∀O,I,A,m,n.∀prf : m = n.
m_rel ?? (IORel O I (eq ?)) A A (eq ?) m n ≝ rel_io_eq on
  _prf : eq (IO ???) ?? to m_rel ?? (IORel ?? (eq ?)) ?? (eq ?) ??.

coercion eq_to_rel_io nocomposites : ∀O,I,A,m,n.∀prf : 
m_rel ?? (IORel O I (eq ?)) A A (eq ?) m n.m = n ≝ eq_rel_io on
  _prf : m_rel ?? (IORel ?? (eq ?)) ?? (eq ?) ?? to eq (IO ???) ??.


let rec pred_io O I A (Perr : predicate errmsg) (P : A → Prop) (v : IO O I A) on v : Prop ≝
  match v with
  [ Value x ⇒ P x
  | Wrong e ⇒ Perr e
  | Interact o f ⇒ ∀i.pred_io … Perr P (f i) 
  ].

let rec pred_io_inject O I A Perr P (a : IO O I A)
  on a : pred_io O I A Perr P a → IO O I (Σx.P x) ≝
  match a return λx.pred_io O I A Perr P x → IO O I (Σx.P x) with
  [ Interact o f ⇒ λprf.
    Interact … o (λx.pred_io_inject … Perr P (f x) (prf x))
  | Value x ⇒ λprf.return «x, prf»
  | Wrong e ⇒ λ_.Wrong … e
  ].

definition IOPred ≝ λO,I,Perr.
  mk_InjMonadPred (IOMonad O I)
    (mk_MonadPred ? (λA.pred_io O I A Perr) ???)
    (λA,P,a_sig.match a_sig with [ mk_Sig a prf ⇒ pred_io_inject O I A Perr P a prf ])
    ?. //
[ #X #P #Q #H #m elim m
  [#o #f #IH #H #i @IH @H
  | #v @H
  | #e //
  ]
| #X #Y #relin #relout #m elim m
  [ #o #f #IH whd in ⊢ (%→?); #H
    #f #G whd #i
    @(IH … (H i) f G)
  | #v #Pv #f #H normalize @H @Pv
  | #e //
  ]
| #X #P * #m elim m
  [ #o #f #IH whd in ⊢ (%→?); #H
    change with (Interact ?????) in ⊢ (???%);
    @interact_proper #i
    @(IH i (H i))
  |*: //
]
qed.

unification hint 0 ≔ O, I, Perr, A, P, v;
M ≟ IOMonad O I, Pr ≟ IOPred O I Perr
⊢ pred_io O I A Perr P v ≡ m_pred M Pr A P v.

definition err_to_io : ∀O,I,T. res T → IO O I T ≝
λO,I,T,v. match v with [ OK v' ⇒ Value O I T v' | Error m ⇒ Wrong O I T m ].

coercion err_to_io : ∀O,I,A.∀c:res A.IO O I A ≝ err_to_io on _c:res ? to IO ???.

(*
lemma res_io_pred : ∀O,I,A,Perr,P.∀m : res A.pred_res … Perr P m → pred_io O I ? Perr P m.
#O #I #A #Perr #P * /2/ qed.

lemma io_res_pred : ∀O,I,A,Perr,P.∀m : res A.pred_io O I ? Perr P m → pred_res ? Perr P m.
#O #I #A #Perr #P * /2/ qed.
*)
definition err_to_io_sig : ∀O,I,T.∀P:T → Prop. res (Sig T P) → IO O I (Sig T P) ≝
λO,I,T,P,v. match v with [ OK v' ⇒ Value O I (Sig T P) v' | Error m ⇒ Wrong O I (Sig T P) m ].
(*coercion err_to_io_sig : ∀O,I,A.∀P:A → Prop.∀c:res (Sig A P).IO O I (Sig A P) ≝ err_to_io_sig on _c:res (Sig ??) to IO ?? (Sig ??).*)

let rec P_io O I (A:Type[0]) (P:A → Prop) (v:IO O I A) on v : Prop ≝
match v return λ_.Prop with
[ Wrong _ ⇒ True
| Value z ⇒ P z
| Interact out k ⇒ ∀v'.P_io O I A P (k v')
].

let rec P_io' O I (A:Type[0]) (P:A → Prop) (v:IO O I A) on v : Prop ≝
match v return λ_.Prop with
[ Wrong _ ⇒ False
| Value z ⇒ P z
| Interact out k ⇒ ∀v'.P_io' O I A P (k v')
].

definition P_to_P_option_io : ∀O,I,A.∀P:A → Prop.option (IO O I A) → Prop ≝
  λO,I,A,P,a.match a with
   [ None ⇒ False
   | Some y ⇒ P_io O I A P y
   ].

let rec io_inject_0 O I (A:Type[0]) (P:A → Prop) (a:IO O I A) (p:P_io O I A P a) on a : IO O I (Sig A P) ≝
(match a return λa'.P_io O I A P a' → ? with
 [ Wrong m ⇒ λ_. Wrong O I ? m
 | Value c ⇒ λp'. Value ??? (mk_Sig A P c p')
 | Interact out k ⇒ λp'. Interact ??? out (λv. io_inject_0 O I A P (k v) (p' v))
 ]) p.

definition io_inject : ∀O,I,A.∀P:A → Prop.∀a:option (IO O I A).∀p:P_to_P_option_io O I A P a.IO O I (Sig A P) ≝ 
  λO,I,A.λP:A → Prop.λa:option (IO O I A).λp:P_to_P_option_io O I A P a. 
  (match a return λa'.P_to_P_option_io O I A P a' → IO O I (Sig A P) with
   [ None ⇒ λp'.?
   | Some b ⇒ λp'. io_inject_0 O I A P b p'
   ]) p.
elim p'; qed.

let rec io_eject O I (A:Type[0]) (P: A → Prop) (a:IO O I (Sig A P)) on a : IO O I A ≝ 
match a with
[ Wrong m ⇒ Wrong ??? m
| Value b ⇒ match b with [ mk_Sig w p ⇒ Value ??? w]
| Interact out k ⇒ Interact ??? out (λv. io_eject ?? A P (k v))
].

coercion io_inject : 
  ∀O,I,A.∀P:A → Prop.∀a.∀p:P_to_P_option_io O I ? P a.IO O I (Sig A P) ≝ io_inject 
  on a:option (IO ???) to IO ?? (Sig ? ?).
coercion io_eject : ∀O,I,A.∀P:A → Prop.∀c:IO O I (Sig A P).IO O I A ≝ io_eject 
  on _c:IO ?? (Sig ? ?) to IO ???.

definition opt_to_io : ∀O,I,T.errmsg → option T → IO O I T ≝
λO,I,T,m,v. match v with [ None ⇒ Wrong ?? T m | Some v' ⇒ Value ??? v' ].

lemma sig_bindIO_OK: ∀O,I,A,B. ∀P:A → Prop. ∀P':B → Prop. ∀e:IO O I (Sig A P). ∀f:Sig A P → IO O I B.
  (∀v:A. ∀p:P v. P_io O I ? P' (f (mk_Sig A P v p))) →
  P_io O I ? P' (bindIO O I (Sig A P) B e f).
#O #I #A #B #P #P' #e #f elim e;
[ #out #k #IH #IH' whd; #res @IH //;
| #v0 elim v0; #v #Hv #IH whd; @IH
| //;
] qed.

lemma sig_bindIO2_OK: ∀O,I,A,B,C. ∀P:(A×B) → Prop. ∀P':C → Prop. ∀e:IO O I (Sig (A×B) P). ∀f: A → B → IO O I C.
  (∀vA:A.∀vB:B. ∀p:P 〈vA,vB〉. P_io O I ? P' (f vA vB)) →
  P_io O I ? P' (bindIO2 O I A B C e f).
#I #O #A #B #C #P #P' #e #f elim e;
[ #out #k #IH #IH' whd; #res @IH @IH'
| #v0 elim v0; #v elim v; #vA #vB #Hv #IH @IH //;
| //;
] qed.

lemma opt_bindIO_OK: ∀O,I,A,B,m. ∀P:B → Prop. ∀e:option A. ∀f: A → IO O I B.
  (∀v:A. e = Some A v → P_io O I ? P (f v)) →
  P_io O I ? P (bindIO O I A B (opt_to_io ??? m e) f).
#I #O #A #B #m #P #e elim e; //; #v #f #H @H //;
qed.

lemma opt_bindIO2_OK: ∀O,I,A,B,C,m. ∀P:C → Prop. ∀e:option (A×B). ∀f: A → B → IO O I C.
  (∀vA:A.∀vB:B. e = Some (A×B) 〈vA,vB〉 → P_io O I ? P (f vA vB)) →
  P_io O I ? P (bindIO2 O I A B C (opt_to_io ??? m e) f).
#I #O #A #B #C #m #P #e elim e; //; #v cases v; #vA #vB #f #H @H //;
qed.

lemma res_bindIO_OK: ∀O,I,A,B. ∀P:B → Prop. ∀e:res A. ∀f: A → IO O I B.
  (∀v:A. e = OK A v → P_io O I ? P (f v)) →
  P_io O I ? P (bindIO O I A B e f).
#I #O #A #B #P #e elim e; //; #v #f #H @H //;
qed.

lemma res_bindIO2_OK: ∀O,I,A,B,C. ∀P:C → Prop. ∀e:res (A×B). ∀f: A → B → IO O I C.
  (∀vA:A.∀vB:B. e = OK (A×B) 〈vA,vB〉 → P_io O I ? P (f vA vB)) →
  P_io O I ? P (bindIO2 O I A B C e f).
#I #O #A #B #C #P #e elim e; //; #v cases v; #vA #vB #f #H @H //;
qed.

lemma bindIO_OK: ∀O,I,A,B. ∀P:B → Prop. ∀e:IO O I A. ∀f: A → IO O I B.
  (∀v:A. P_io O I ? P (f v)) →
  P_io O I ? P (bindIO O I A B e f).
#I #O #A #B #P #e elim e;
[ #out #k #IH #f #H whd; #res @IH //;
| #v #f #H @H
| //;
] qed.

lemma bindIO2_OK: ∀O,I,A,B,C. ∀P:C → Prop. ∀e:IO O I (A×B). ∀f: A → B → IO O I C.
  (∀v1:A.∀v2:B. P_io O I ? P (f v1 v2)) →
  P_io O I ? P (bindIO2 O I A B C e f).
#I #O #A #B #C #P #e elim e;
[ #out #k #IH #f #H whd; #res @IH //;
| #v cases v; #v1 #v2 #f #H @H
| //;
] qed.

lemma P_bindIO_OK: ∀O,I,A,B. ∀P':A → Prop. ∀P:B → Prop. ∀e:IO O I A. ∀f: A → IO O I B.
  P_io … P' e →
  (∀v:A. P' v → P_io O I ? P (f v)) →
  P_io O I ? P (bindIO O I A B e f).
#I #O #A #B #P' #P #e elim e;
[ #out #k #IH #f #He #H whd in He ⊢ %; #res @IH /2 by /;
| #v #f #He #H @H @He
| //;
] qed.

lemma P_bindIO2_OK: ∀O,I,A,B,C. ∀P':A×B → Prop. ∀P:C → Prop. ∀e:IO O I (A×B). ∀f: A → B → IO O I C.
  P_io … P' e →
  (∀v1:A.∀v2:B. P' 〈v1,v2〉 → P_io O I ? P (f v1 v2)) →
  P_io O I ? P (bindIO2 O I A B C e f).
#I #O #A #B #C #P' #P #e elim e;
[ #out #k #IH #f #He #H whd in He ⊢ %; #res @IH /2 by /;
| #v cases v; #v1 #v2 #f #He #H @H @He
| //;
] qed.


(* Is there a way to prove this without extensionality? *)
(*
lemma bind_assoc_r: ∀O,I,A,B,C,e,f,g.
  ∀ext:(∀T1,T2:Type[0].∀f,f':T1 → T2.(∀x.f x = f' x) → f = f').
  bindIO O I B C (bindIO O I A B e f) g = bindIO O I A C e (λx.bindIO O I B C (f x) g).
#O #I #A #B #C #e #f #g #ext elim e;
[ #o #k #IH whd in ⊢ (??%%); @eq_f
    @ext @IH
| #v @refl
| #m @refl
] qed.
*)
definition bind_assoc_r ≝ λO,I.m_bind_bind (IOMonad O I).

(*
lemma extract_subset_pair_io: ∀O,I,A,B,C,P. ∀e:{e:A×B | P e}. ∀Q:A→B→IO O I C. ∀R:C→Prop.
  (∀a,b. eject ?? e = 〈a,b〉 → P 〈a,b〉 → P_io O I ? R (Q a b)) →
  P_io O I ? R (match eject ?? e with [ pair a b ⇒ Q a b ]).
#I #O #A #B #C #P #e #Q #R cases e; #e' cases e'; normalize;
[ *;
| #e'' cases e''; #a #b #Pab #H normalize; /2 by _/;
] qed.
*)

(* Inversion when injecting errors into IO monad. *)
lemma bind_res_value : ∀O,I,A,B,e,f,v. ∀P:B → Type[0].
  (∀a. e = OK A a → f a = OK B v → P v) →
  (match e »= f with [ OK v ⇒ Value … v | Error m ⇒ Wrong … m ] = Value O I B v → P v).
#O #I #A #B *
[ #a #f #v #P #H #E @H [ @a | @refl | normalize in E; cases (f a) in E ⊢ %;
  [ #v' #E normalize in E; destruct @refl | #m #E normalize in E; destruct ] ]
| #m #f #v #P #H #E whd in E:(??%?); destruct
] qed.

lemma bindIO_value : ∀O,I,A,B,e,f,v. ∀P:B → Type[0].
  (∀a. e = Value ??? a → f a = Value ??? v → P v) →
  (bindIO O I A B e f = Value ??? v → P v).
#O #I #A #B *
[ #o #k #f #v #P #H #E whd in E:(??%?); destruct
| #a #f #v #P #H #E @H [ @a | @refl | @E ]
| #m #f #v #P #H #E whd in E:(??%?); destruct
] qed.

lemma bindIO_res_interact : ∀O,I,A,B,e,f,o,k. ∀P:∀o:O. (I o → IO O I B) → Type[0].
  (∀v. e = OK A v → f v = Interact … o k → P o k) →
  bindIO O I A B (err_to_io … e) f = Interact … o k → P o k.
#O #I #A #B *
[ #a #f #o #k #P #H #E @(H … E) @refl
| #m #f #o #k #P #_ #E whd in E:(??%?); destruct
] qed.

lemma bindIO_opt_interact : ∀O,I,A,B,m,e,f,o,k. ∀P:∀o:O. (I o → IO O I B) → Type[0].
  (∀v. e = Some A v → f v = Interact … o k → P o k) →
  bindIO O I A B (opt_to_io … m e) f = Interact … o k → P o k.
#O #I #A #B #m *
[ #f #o #k #P #_ #E whd in E:(??%?); destruct
| #a #f #o #k #P #H #E @(H … E) @refl
] qed.

lemma bindIO_inversion: ∀O,I.
  ∀A,B: Type[0]. ∀f: IO O I A. ∀g: A → IO O I B. ∀y: B.
  (f »= g = return y) →
  ∃x. (f = return x) ∧ (g x = return y).
#O #I #A #B #f #g #y cases f normalize 
[ #o #k #E destruct
| #a #e %{a} /2 by conj/
| #m #H destruct (H)
] qed.

(* When something in the error monad has found its way into the IO monad,
   ensure that we can implicitly go back. *)
lemma io_eq_to_res : ∀O,I. ∀T:Type[0]. ∀e:res T. ∀v.
  err_to_io … e = Value O I T v →
  e = OK T v.
#O #I #T *
[ #e #v #E normalize in E; destruct @refl
| #m #v #E normalize in E; destruct
]
qed.

coercion io_eq_from_res :
  ∀O,I,T,e,v. ∀E:err_to_io O I T e = Value O I T v. e = OK T v ≝ io_eq_to_res
  on _E:eq (IO ???) ?? to eq (res ?) ??.