Compile-time enforced resource exhaustion in OCaml using an indexed state monad.

Connection.hs

{-# LANGUAGE RebindableSyntax,
             KindSignatures,
             DataKinds,
             GADTs,
             GeneralizedNewtypeDeriving,
             StandaloneDeriving,
             FlexibleInstances,
             Rank2Types #-}

module Connection (
      (>>=), (>>), return, fail
    , ifThenElse
    , Client
    , ClientAction
    , Program
    , connect
    , run, fetch, next, drain
    ) where

import Prelude hiding ((>>=), (>>), fail, return, log)
import qualified Prelude as P

import Control.Applicative (Applicative)
import Control.Monad.IO.Class (MonadIO(..))

import Control.Monad.Indexed
import Control.Monad.Indexed.State (IxStateT(..), iget, iput)

-- Required for RebindableSyntax
ifThenElse :: Bool -> a -> a -> a
ifThenElse b i j = case b of
    True -> i
    False -> j

return :: IxPointed m => a -> m i i a
return = ireturn

(>>=) :: IxMonad m => m i j a -> (a -> m j k b) -> m i k b
(>>=) = (>>>=)

(>>) :: IxMonad m => m i j () -> m j k b -> m i k b
(>>) a b = a >>= \() -> b

fail :: String -> a
fail = error

-- Mostly public API
data ClientState = Ready | Producing
  deriving (Show)

-- Constructor & record selectors are private
data Client (s :: ClientState) where
    Client :: { isProducing :: Bool, stream :: [Int] } -> Client s

deriving instance Show (Client s)

-- Constructor & record selector are private
newtype ClientAction m i j a = CA { unCA :: IxStateT m i j a }

deriving instance Monad m => Functor (ClientAction m i j)
deriving instance Monad m => Applicative (ClientAction m i i)
deriving instance Monad m => Monad (ClientAction m i i)
deriving instance MonadIO m => MonadIO (ClientAction m i i)

deriving instance Monad m => IxFunctor (ClientAction m)
deriving instance Monad m => IxApplicative (ClientAction m)
deriving instance Monad m => IxPointed (ClientAction m)
deriving instance Monad m => IxMonad (ClientAction m)

-- Private API

get :: Monad m => ClientAction m i i i
get = CA iget

put :: Monad m => j -> ClientAction m i j ()
put = CA . iput

log :: MonadIO m => String -> ClientAction m i i ()
log s = liftIO $ putStrLn $ "[LOG] " ++ s

-- Public API

fetch :: MonadIO m => [Int] -> ClientAction m (Client Ready) (Client Producing) ()
fetch l = do
    log $ "Fetching " ++ show l
    _ <- get
    -- Normally, we'd use the above result somehow
    put $ Client { isProducing = not (null l)
                 , stream = l
                 }

next :: MonadIO m => ClientAction m (Client Producing) (Client Producing) (Maybe Int)
next = do
    c <- get
    if (isProducing c)
        then do
            let l = stream c
                a = head l
                t = tail l
            put $ if null t
                then Client False []
                else Client True t
            log $ "Yielding " ++ show a
            return (Just a)
        else
            return Nothing

drain :: MonadIO m => ClientAction m (Client Producing) (Client Ready) ()
drain = do
    c <- get
    log $ "Draining: " ++ show (stream c)
    put $ Client False []

connect :: Monad m => m (Client Ready)
connect = P.return (Client False [])

type Program m a = ClientAction m (Client Ready) (Client Ready) a

run :: Monad m => Program m a
               -> Client Ready
               -> m (a, Client Ready)
run = runIxStateT . unCA

demo.hs

{-# LANGUAGE RebindableSyntax #-}
module Main where

-- The 2 lines below are boilerplate, so is the LANGUAGE pragma
import Prelude hiding ((>>), (>>=), return, fail)
import qualified Prelude

import Connection

program :: Program IO Int
program = do
    fetch [1, 2, 3]
    Just 1 <- next
    Just two <- next

    -- Without call to 'drain':
    --     Couldn't match type 'Producing with 'Ready
    --     Expected type: ClientAction m (Client 'Producing) (Client 'Ready) Int
    --     Actual type: ClientAction m (Client 'Producing) (Client 'Producing) Int
    -- as desired
    drain

    -- We can re-use the Client within the same Program, as long as it's
    -- Ready again, which it is thanks to `drain` above
    fetch [4, 5]

    -- This fails to compile: we didn't drain the client yet, it's currently in
    -- Producing state, and `fetch` wants a Ready Client
    -- fetch [6, 7] >>= \() ->

    drain

    fetch []
    Nothing <- next
    drain

    return two


main :: IO ()
main = connect >>= run program >>= print
  where
    -- Use a non-indexed Monad here
    (>>=) = (Prelude.>>=)

drain.ml

module type IndexedMonad = sig
  type ('a, 'i, 'o) t

  val bind : ('a, 'i, 'o) t -> ('a -> ('b, 'o, 'j) t) -> ('b, 'i, 'j) t
  val return : 'a -> ('a, 'i, 'i) t
end

module IndexedState : sig
  type ('a, 'i, 'o) t

  include IndexedMonad with type ('a, 'i, 'o) t := ('a, 'i, 'o) t

  val get : ('i, 'i, 'i) t
  val put : 'o -> (unit, 'i, 'o) t

  val run : ('a, 'i, 'o) t -> 'i -> ('a * 'o)
end = struct
  type ('a, 'i, 'o) t = ('i -> ('a * 'o))

  let return a = fun s -> (a, s)
  let bind k m = fun s -> let (a, s') = k s in m a s'

  let get = fun x -> (x, x)
  let put x = fun _ -> ((), x)

  let run t s = t s
end

module Client : sig
  type ready
  type producing

  type 'a t

  val make : int list -> ready t

  val fetch : ready t -> producing t
  val next : producing t -> (int * producing t) option
  val drain : producing t -> ready t
end = struct
  type ready
  type producing

  type _ t =
    | Client : int list -> 'a t

  let make stream = Client stream

  let fetch (Client l) = 
    Printf.printf "Fetch\n";
    Client l
  let next (Client l) = match l with
    | [] -> Printf.printf "Next -> None\n"; None
    | (x :: xs) -> Printf.printf "Next -> Some %d\n" x; Some (x, Client xs)
  let drain (Client l) =
    let s = String.concat ", " (List.map string_of_int l) in
    Printf.printf "Draining [%s]\n" s;
    Client []
end

module ClientAction : sig
  type ('a, 'i, 'o) t

  val fetch : unit -> (unit, Client.ready Client.t, Client.producing Client.t) t
  val next : unit -> (int option, Client.producing Client.t, Client.producing Client.t) t
  val drain : unit -> (unit, Client.producing Client.t, Client.ready Client.t) t

  val run : ('a, Client.ready Client.t, Client.ready Client.t) t -> Client.ready Client.t -> 'a

  include IndexedMonad with type ('a, 'i, 'j) t := ('a, 'i, 'j) t
  val (>>=) : ('a, 'i, 'j) t -> ('a -> ('b, 'j, 'k) t) -> ('b, 'i, 'k) t

  val return : 'a -> ('a, Client.ready Client.t, Client.ready Client.t) t
end = struct
  include IndexedState

  let (>>=) = bind

  let fetch () =
    get >>= fun c ->
    put (Client.fetch c)
  let next () =
    get >>= fun c ->
    match Client.next c with
     | None -> return None
     | Some (v, c') -> put c' >>= fun () -> return (Some v)
  let drain () =
    get >>= fun c ->
    put (Client.drain c)

  let run' = run
  let run t c = let (r, _) = run' t c in r
end

module Demo = struct
  open ClientAction

  let go =
    fetch () >>= fun () ->
    next () >>= fun _ ->
    next () >>= fun two ->
    drain () >>= fun () ->

    (* Can't drain a 'ready' client
    drain () >>= fun () ->
    *)
    (* Can't next a 'ready' client
    next () >>= fun _ ->
    *)

    fetch () >>= fun () ->

    (* Can't end the action as long as the client is producing
    return two
    *)

    drain () >>= fun () ->
    return two
end

let main () =
  let client = Client.make [1; 2; 3] in
  let res = ClientAction.run Demo.go client in
  let s = match res with
   | None -> "None"
   | Some i -> Printf.sprintf "Some %d" i
  in
  Printf.printf "%s\n" s
;;

main ()