(*
    A typestate sockets library

    This is a bit more involved than the example in the paper,
    because we have error cases.  We deal with this by raising
    an exception which contains a witness that allows recovering
    the capability if presented with the corresponding socket.
*)

#load "libsocket"

module type SOCKET_CAP = sig
  (* The representation of a socket *)
  type 'a socket

  (* Socket capabilities and the socket states *)
  type 'a @ 'c : A
  type raw
  type named
  type listening
  type ready

  (* Socket operations *)
  val socket  : unit → ∃ 'a. 'a socket × 'a@raw
  val bind    : 'a socket → int → 'a@raw → 'a@named

  val connect : 'a socket → string → string →
                  'a@raw + 'a@named → 'a@ready
  val listen  : 'a socket → 'a@named → 'a@listening
  val accept  : 'a socket → 'a@listening →
                  (∃ 'b. 'b socket × 'b@ready) × 'a@listening
  val send    : 'a socket → string → 'a@ready → 'a@ready
  val recv    : 'a socket → int → 'a@ready → string × 'a@ready
  val close   :  'a socket → 'a@'c → unit

  val isSame  : 'a socket → 'b socket → ('a@'c → 'b@'c) option

  (* Operations for catching the error state associated with a given
     socket. *)
  val catchRaw       : (unit -A> `r) → 'a socket × ('a@raw -A> `r)       → `r
  val catchNamed     : (unit -A> `r) → 'a socket × ('a@named -A> `r)     → `r
  val catchListening : (unit -A> `r) → 'a socket × ('a@listening -A> `r) → `r
  val catchReady     : (unit -A> `r) → 'a socket × ('a@ready -A> `r)     → `r

  type 'a dynamicCap = Raw of 'a@raw
                     | Named of 'a@named
                     | Listening of 'a@listening
                     | Ready of 'a@ready

  exception Socket of (∃'a. 'a socket × 'a dynamicCap) option × string
end

module SocketCap : SOCKET_CAP = struct
  module S = Socket

  type 'a socket  = S.socket
  type 'a @ 'c = unit
  type raw
  type named
  type listening
  type ready

  type 'a dynamicCap = Raw of 'a@raw
                     | Named of 'a@named
                     | Listening of 'a@listening
                     | Ready of 'a@ready

  exception Socket of (∃'a. 'a socket × 'a dynamicCap) option × string

  let lift thunk sock mkcap =
    try thunk ()
    with IOError msg → raise (Socket (Some (sock, mkcap ()), msg))

  let socket _ : ∃ 'a. 'a socket × 'a@raw =
    try (S.socket (), ())
    with IOError msg → raise (Socket (None, msg))

  let bind sock port _ = lift (λ_ → S.bind sock port) sock Raw

  let connect sock host port cap =
    lift (λ_ → S.connect sock host port) sock
         (match cap with Left _ → Raw | Right _ → Named)

  let listen sock _ = lift (λ_ → S.listen sock) sock Named

  let accept sock _ =
    lift (λ_ → ((S.accept sock, ()) : ∃ 'a. 'a socket × 'a@ready, ()))
          sock Listening

  let send sock data _ = lift (λ_ → S.send sock data; ()) sock Ready

  let recv sock len _ = lift (λ_ → (S.recv sock len, ())) sock Ready

  let close sock _ =
    try S.close sock
    with IOError msg → raise (Socket (None, msg))

  let isSame sock sock' =
    if sock == sock'
      then Some (λ_ → ())
      else None

  let catchBy body state (sock', handler) =
    try body () with
    Socket ((Some (sock, dyncap), msg) as se) →
      if dyncap == state () && sock == sock'
        then handler ()
        else raise (Socket se)

  let catchRaw body       = catchBy body Raw
  let catchNamed body     = catchBy body Named
  let catchListening body = catchBy body Listening
  let catchReady body     = catchBy body Ready
end

module SocketCap2 : SOCKET_CAP = struct
  open SocketCap
  let catchBy body
              (prj : ∀'a. 'a dynamicCap →
                          'a dynamicCap + ('a@'c -A> 'a dynamicCap) × 'a@'c)
              (sock', handler) =
    try body () with
    | Socket (Some (sock, dyncap), msg) →
        match prj dyncap with
        | Left dyncap
          → raise (Socket (Some (sock, dyncap), msg))
        | Right (uncap, cap)
          → match isSame sock sock' with
            | None → raise (Socket (Some (sock, uncap cap), msg))
            | Some witness → handler (witness cap)

  let catchRaw body =
    catchBy body (function Raw cap → Right (Raw, cap)
                         | dyncap → Left dyncap)
  let catchNamed body =
    catchBy body (function Named cap → Right (Named, cap)
                         | dyncap → Left dyncap)
  let catchListening body =
    catchBy body (function Listening cap → Right (Listening, cap)
                         | dyncap → Left dyncap)
  let catchReady body =
    catchBy body (function Ready cap → Right (Ready, cap)
                         | dyncap → Left dyncap)
end