val interp_morph (#a #b: _) (c: rwtree a) (f: (a -> rwtree b)) (p s0: _)
    : Lemma
    (interp_as_wp c s0 (fun (y, s1) -> interp_as_wp (f y) s1 p) == interp_as_wp (tbind c f) s0 p)
val interp_morph (#a #b: _) (c: rwtree a) (f: (a -> rwtree b)) (p s0: _)
    : Lemma
    (interp_as_wp c s0 (fun (y, s1) -> interp_as_wp (f y) s1 p) == interp_as_wp (tbind c f) s0 p)
let rec interp_morph #a #b (c : rwtree a) (f : a -> rwtree b) (p:_) (s0:_)
  : Lemma (interp_as_wp c s0 (fun (y, s1) -> interp_as_wp (f y) s1 p) == interp_as_wp (tbind c f) s0 p)
  = match c with
    | Return x -> ()
    | Op Read _ k ->
      let aux (o:state) : Lemma (interp_as_wp (k o) s0 (fun (y, s1) -> interp_as_wp (f y) s1 p)
                                        == interp_as_wp (tbind (k o) f) s0 p) =
        interp_morph (k o) f p s0
      in
      Classical.forall_intro aux

    | Op Write s k ->
      let aux (o:unit) : Lemma (interp_as_wp (k o) s (fun (y, s1) -> interp_as_wp (f y) s1 p)
                                        == interp_as_wp (tbind (k o) f) s p) =
        interp_morph (k o) f p s
      in
      Classical.forall_intro aux

    | _ -> ()
module AlgForAll

open FStar.Calc
module FE = FStar.FunctionalExtensionality
module F = FStar.FunctionalExtensionality
module W = FStar.WellFounded
module T = FStar.Tactics.V2
module ID5 = ID5
open Alg

type rwtree a = Alg.tree a [Read; Write]
let tbind : #a:_ -> #b:_ -> rwtree a -> (a -> rwtree b) -> rwtree b = fun c f -> Alg.bind _ _ c f

let st_wp0 (a:Type) : Type = state -> (a & state -> Type0) -> Type0

let st_monotonic #a (w : st_wp0 a) : Type0 =
  //forall s0 p1 p2. (forall r. p1 r ==> p2 r) ==> w s0 p1 ==> w s0 p2
  // ^ this version seems to be less SMT-friendly
  forall s0 p1 p2. (forall x s1. p1 (x, s1) ==> p2 (x, s1)) ==> w s0 p1 ==> w s0 p2

let st_wp (a:Type) = wp:st_wp0 a{st_monotonic wp}

unfold
let return_wp #a x : st_wp a = fun s0 p -> p (x, s0)

unfold
let bind_wp #a #b (w : st_wp a) (wf : a -> st_wp b)
  : st_wp b
  = fun s0 p -> w s0 (fun (y, s1) -> wf y s1 p)

unfold
let read_wp : st_wp state = fun s0 p -> p (s0, s0)

unfold
let write_wp : state -> st_wp unit = fun s _ p -> p ((), s)

(* Also doable with handlers *)
let rec interp_as_wp #a (t : rwtree a) : st_wp a =
  match t with
  | Return x -> return_wp x
  | Op Read _ k ->
    bind_wp read_wp (fun s -> interp_as_wp (k s))
  | Op Write s k ->
    bind_wp (write_wp s) (fun (o:unit) -> interp_as_wp (k o))

(* With handlers. Can only be done into []? See the use of `run`. *)
let interp_as_wp2 #a (t : rwtree a) : Alg (st_wp a) [] =
  handle_with #a #(st_wp a) #[Read; Write] #[]
              (fun () -> Alg?.reflect t)
              (fun x -> return_wp x)
              (function Read  -> (fun i k -> bind_wp read_wp (fun s -> run (fun () -> k s)))
                      | Write -> (fun i k -> bind_wp (write_wp i) (fun _ -> run k)))

(* Bug: defining this as a FStar.Preorder.preorder
causes stupid failures ahead *)
val stronger : (#a:Type) -> st_wp a -> st_wp a -> Type0
let stronger w1 w2 = forall p s. w1 p s ==> w2 p s

let equiv #a (w1 w2 : st_wp a) = w1 `stronger` w2 /\ w2 `stronger` w1

let (<<=) = stronger

val interp_ret (#a:Type) (x:a) : Lemma (return_wp x `stronger` interp_as_wp (Return x))
let interp_ret x = ()

let wp_is_monotonic #a (wp : st_wp a) : Type0 =
  forall p1 p2 s0. (forall x s1. p1 (x, s1) ==> p2 (x, s1)) ==> wp s0 p1 ==> wp s0 p2

let bind_preserves_mon #a #b (wp : st_wp a) (f : a -> st_wp b)
  : Lemma (requires (wp_is_monotonic wp /\ (forall x. wp_is_monotonic (f x))))
          (ensures (wp_is_monotonic (bind_wp wp f)))
  = ()

let rec interp_monotonic #a (c:rwtree a) : Lemma (wp_is_monotonic (interp_as_wp c)) =
  match c with
  | Return x -> ()
  | Op Read _ k ->
    let aux (x:state) : Lemma (wp_is_monotonic (interp_as_wp (k x))) =
      interp_monotonic (k x)
    in
    Classical.forall_intro aux;
    bind_preserves_mon read_wp (fun x -> interp_as_wp (k x))
  | Op Write s k ->
    let aux (x:unit) : Lemma (wp_is_monotonic (interp_as_wp (k x))) =
      interp_monotonic (k x)
    in
    Classical.forall_intro aux;
    bind_preserves_mon (write_wp s) (fun x -> interp_as_wp (k x))

let elim_str #a (w1 w2 : st_wp a) (p : (a & state -> Type0)) (s0:state)
  : Lemma (requires (w1 <<= w2))
          (ensures w1 s0 p ==> w2 s0 p)
  = ()
/
prims.fst.checked ID5.fst.checked FStar.WellFounded.fst.checked FStar.Tactics.V2.fst.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.Monotonic.Pure.fst.checked FStar.FunctionalExtensionality.fsti.checked FStar.Classical.fsti.checked FStar.Calc.fsti.checked Alg.fst.checked
AlgForAll.fst

    c: AlgForAll.rwtree a ->
    f: (_: a -> AlgForAll.rwtree b) ->
    p: (_: (b * Alg.state) -> Type0) ->
    s0: Alg.state
  -> FStar.Pervasives.Lemma
    (ensures
      AlgForAll.interp_as_wp c
        s0
        (fun _ ->
            (let FStar.Pervasives.Native.Mktuple2 #_ #_ y s1 = _ in
              AlgForAll.interp_as_wp (f y) s1 p)
            <:
            Type0) ==
      AlgForAll.interp_as_wp (AlgForAll.tbind c f) s0 p)
AlgForAll.rwtree FStar.Pervasives.Native.tuple2 Alg.state Alg.op_inp Alg.Read Alg.op_out Alg.tree0 FStar.Classical.forall_intro Prims.eq2 AlgForAll.interp_as_wp AlgForAll.tbind Prims.unit Prims.l_True Prims.squash Prims.Nil FStar.Pervasives.pattern AlgForAll.interp_morph Alg.Write
let rec interp_morph #a #b (c: rwtree a) (f: (a -> rwtree b)) (p: _) (s0: _)
    : Lemma
    (interp_as_wp c s0 (fun (y, s1) -> interp_as_wp (f y) s1 p) == interp_as_wp (tbind c f) s0 p) =
  
match c with
| Return x -> ()
| Op Read _ k ->
  let aux (o: state)
      : Lemma
      (interp_as_wp (k o) s0 (fun (y, s1) -> interp_as_wp (f y) s1 p) ==
        interp_as_wp (tbind (k o) f) s0 p) =
    interp_morph (k o) f p s0
  in
  Classical.forall_intro aux
| Op Write s k ->
  let aux (o: unit)
      : Lemma
      (interp_as_wp (k o) s (fun (y, s1) -> interp_as_wp (f y) s1 p) ==
        interp_as_wp (tbind (k o) f) s p) =
    interp_morph (k o) f p s
  in
  Classical.forall_intro aux
| _ -> ()