include "Plogic/russell_support.ma".
include "extralib.ma".
include "Errors.ma".

(* IO monad *)

ninductive IO (input,output:Type) (T:Type) : Type ≝
| Interact : output → (input → IO input output T) → IO input output T
| Value : T → IO input output T
| Wrong : IO input output T.

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

nlet rec bindIO2 (I,O,T1,T2,T':Type) (v:IO I O (T1×T2)) (f:T1 → T2 → IO I O T') on v : IO I O T' ≝
match v with
[ Interact out k ⇒ (Interact ??? out (λres. bindIO2 ?? T1 T2 T' (k res) f))
| Value v' ⇒ match v' with [ mk_pair v1 v2 ⇒ f v1 v2 ]
| Wrong ⇒ Wrong ?? T'
].

ndefinition err_to_io : ∀I,O,T. res T → IO I O T ≝
λI,O,T,v. match v with [ OK v' ⇒ Value I O T v' | Error ⇒ Wrong I O T ].
(*ncoercion err_to_io : ∀I,O,A.∀c:res A.IO I O A ≝ err_to_io on _c:res ? to IO ???.*)
ndefinition err_to_io_sig : ∀I,O,T.∀P:T → Prop. res (sigma T P) → IO I O (sigma T P) ≝
λI,O,T,P,v. match v with [ OK v' ⇒ Value I O (sigma T P) v' | Error ⇒ Wrong I O (sigma T P) ].
ncoercion err_to_io_sig : ∀I,O,A.∀P:A → Prop.∀c:res (sigma A P).IO I O (sigma A P) ≝ err_to_io_sig on _c:res (sigma ??) to IO ?? (sigma ??).


(* If the original definitions are vague enough, do I need to do this? *)
notation > "! ident v ← e; e'" with precedence 40 for @{'bindIO ${e} (λ${ident v}.${e'})}.
notation > "! ident v : ty ← e; e'" with precedence 40 for @{'bindIO ${e} (λ${ident v} : ${ty}.${e'})}.
notation < "vbox(! \nbsp ident v ← e; break e')" with precedence 40 for @{'bindIO ${e} (λ${ident v}.${e'})}.
notation < "vbox(! \nbsp ident v : ty ← e; break e')" with precedence 40 for @{'bindIO ${e} (λ${ident v} : ${ty}.${e'})}.
notation > "! 〈ident v1, ident v2〉 ← e; e'" with precedence 40 for @{'bindIO2 ${e} (λ${ident v1}.λ${ident v2}.${e'})}.
notation > "! 〈ident v1 : ty1, ident v2 : ty2〉 ← e; e'" with precedence 40 for @{'bindIO2 ${e} (λ${ident v1} : ${ty1}.λ${ident v2} : ${ty2}.${e'})}.
notation < "vbox(! \nbsp 〈ident v1, ident v2〉 ← e; break e')" with precedence 40 for @{'bindIO2 ${e} (λ${ident v1}.λ${ident v2}.${e'})}.
notation < "vbox(! \nbsp 〈ident v1 : ty1, ident v2 : ty2〉 ← e; break e')" with precedence 40 for @{'bindIO2 ${e} (λ${ident v1} : ${ty1}.λ${ident v2} : ${ty2}.${e'})}.
interpretation "IO monad bind" 'bindIO e f = (bindIO ???? e f).
interpretation "IO monad pair bind" 'bindIO2 e f = (bindIO2 ????? e f).
(**)
nlet rec P_io (I,O,A:Type) (P:A → Prop) (v:IO I O A) on v : Prop ≝
match v return λ_.Prop with
[ Wrong ⇒ True
| Value z ⇒ P z
| Interact out k ⇒ ∀v'.P_io I O A P (k v')
].

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

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

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

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

nlet rec io_eject (I,O,A:Type) (P: A → Prop) (a:IO I O (sigma A P)) on a : IO I O A ≝ 
match a with
[ Wrong ⇒ Wrong ???
| Value b ⇒ match b with [ sig_intro w p ⇒ Value ??? w]
| Interact out k ⇒ Interact ??? out (λv. io_eject ?? A P (k v))
].

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

ndefinition opt_to_io : ∀I,O,T.option T → IO I O T ≝
λI,O,T,v. match v with [ None ⇒ Wrong I O T | Some v' ⇒ Value I O T v' ].
ncoercion opt_to_io : ∀I,O,T.∀v:option T. IO I O T ≝ opt_to_io on _v:option ? to IO ???.

nlemma sig_bindIO_OK: ∀I,O,A,B. ∀P:A → Prop. ∀P':B → Prop. ∀e:IO I O (sigma A P). ∀f:sigma A P → IO I O B.
  (∀v:A. ∀p:P v. P_io I O ? P' (f (sig_intro A P v p))) →
  P_io I O ? P' (bindIO I O (sigma A P) B e f).
#I O A B P P' e f; nelim e;
##[ #out k IH; #IH'; nwhd; #res; napply IH; //;
##| #v0; nelim v0; #v Hv IH; nwhd; napply IH;
##| //;
##] nqed.

nlemma sig_bindIO2_OK: ∀I,O,A,B,C. ∀P:(A×B) → Prop. ∀P':C → Prop. ∀e:IO I O (sigma (A×B) P). ∀f: A → B → IO I O C.
  (∀vA:A.∀vB:B. ∀p:P 〈vA,vB〉. P_io I O ? P' (f vA vB)) →
  P_io I O ? P' (bindIO2 I O A B C e f).
#I O A B C P P' e f; nelim e;
##[ #out k IH; #IH'; nwhd; #res; napply IH; napply IH';
##| #v0; nelim v0; #v; nelim v; #vA vB Hv IH; napply IH; //;
##| //;
##] nqed.

nlemma opt_bindIO_OK: ∀I,O,A,B. ∀P:B → Prop. ∀e:option A. ∀f: A → IO I O B.
  (∀v:A. e = Some A v → P_io I O ? P (f v)) →
  P_io I O ? P (bindIO I O A B e f).
#I O A B P e; nelim e; //; #v f H; napply H; //;
nqed.

nlemma opt_bindIO2_OK: ∀I,O,A,B,C. ∀P:C → Prop. ∀e:option (A×B). ∀f: A → B → IO I O C.
  (∀vA:A.∀vB:B. e = Some (A×B) 〈vA,vB〉 → P_io I O ? P (f vA vB)) →
  P_io I O ? P (bindIO2 I O A B C e f).
#I O A B C P e; nelim e; //; #v; ncases v; #vA vB f H; napply H; //;
nqed.

nlemma bindIO_OK: ∀I,O,A,B. ∀P:B → Prop. ∀e:IO I O A. ∀f: A → IO I O B.
  (∀v:A. P_io I O ? P (f v)) →
  P_io I O ? P (bindIO I O A B e f).
#I O A B P e; nelim e;
##[ #out k IH; #f H; nwhd; #res; napply IH; //;
##| #v f H; napply H;
##| //;
##] nqed.

(* TODO: is there a way to prove this without extensionality?

nlemma bind_assoc_r: ∀A,B,C,e,f,g.
  bindIO B C (bindIO A B e f) g = bindIO A C e (λx.bindIO B C (f x) g).
#A B C e f g; nelim e;
##[ #fn args k IH; nwhd in ⊢ (???%);
nnormalize; 
*)

nlemma extract_subset_pair_io: ∀I,O,A,B,C,P. ∀e:{e:A×B | P e}. ∀Q:A→B→IO I O C. ∀R:C→Prop.
  (∀a,b. eject ?? e = 〈a,b〉 → P 〈a,b〉 → P_io I O ? R (Q a b)) →
  P_io I O ? R (match eject ?? e with [ mk_pair a b ⇒ Q a b ]).
#I O A B C P e Q R; ncases e; #e'; ncases e'; nnormalize;
##[ *;
##| #e''; ncases e''; #a b Pab H; nnormalize; /2/;
##] nqed.