{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Network.TypedProtocol.ReqResp.Client
  ( -- * Normal client
    ReqRespClient (..)
  , reqRespClientPeer
    -- * Pipelined client
  , ReqRespClientPipelined (..)
  , reqRespClientPeerPipelined
  , ReqRespIdle (..)
  , reqRespClientPeerIdle
    -- * Request once
  , requestOnce
  ) where

import           Network.TypedProtocol.Core
import           Network.TypedProtocol.Peer.Client
import           Network.TypedProtocol.Peer.Server (Server)
import           Network.TypedProtocol.Proofs (connect)
import           Network.TypedProtocol.ReqResp.Type

data ReqRespClient req resp m a where
  SendMsgReq     :: req
                 -> (resp -> m (ReqRespClient req resp m a))
                 -> ReqRespClient req resp m a

  SendMsgDone    :: m a -> ReqRespClient req resp m a


-- | Interpret a particular client action sequence into the client side of the
-- 'ReqResp' protocol.
--
reqRespClientPeer
  :: Monad m
  => ReqRespClient req resp m a
  -> Client (ReqResp req resp) NonPipelined StIdle m a

reqRespClientPeer :: forall (m :: * -> *) req resp a.
Monad m =>
ReqRespClient req resp m a
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
reqRespClientPeer (SendMsgDone m a
result) =
    -- We do an actual transition using 'yield', to go from the 'StIdle' to
    -- 'StDone' state. Once in the 'StDone' state we can actually stop using
    -- 'done', with a return value.
    m (Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a.
m (Client ps pl st m a) -> Client ps pl st m a
Effect (m (Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
 -> Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
-> m (Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
forall a b. (a -> b) -> a -> b
$
      Message (ReqResp req resp) 'StIdle 'StDone
-> Client (ReqResp req resp) 'NonPipelined 'StDone m a
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a
       (st' :: ps).
(StateTokenI st, StateTokenI st', StateAgency st ~ 'ClientAgency,
 Outstanding pl ~ 'Z) =>
Message ps st st' -> Client ps pl st' m a -> Client ps pl st m a
Yield Message (ReqResp req resp) 'StIdle 'StDone
forall {k} {k1} (req :: k) (resp :: k1).
Message (ReqResp req resp) 'StIdle 'StDone
MsgDone (Client (ReqResp req resp) 'NonPipelined 'StDone m a
 -> Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
-> (a -> Client (ReqResp req resp) 'NonPipelined 'StDone m a)
-> a
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Client (ReqResp req resp) 'NonPipelined 'StDone m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a.
(StateTokenI st, StateAgency st ~ 'NobodyAgency,
 Outstanding pl ~ 'Z) =>
a -> Client ps pl st m a
Done (a -> Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
-> m a -> m (Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
result

reqRespClientPeer (SendMsgReq req
req resp -> m (ReqRespClient req resp m a)
next) =

    -- Send our message.
    Message (ReqResp req resp) 'StIdle 'StBusy
-> Client (ReqResp req resp) 'NonPipelined 'StBusy m a
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a
       (st' :: ps).
(StateTokenI st, StateTokenI st', StateAgency st ~ 'ClientAgency,
 Outstanding pl ~ 'Z) =>
Message ps st st' -> Client ps pl st' m a -> Client ps pl st m a
Yield (req -> Message (ReqResp req resp) 'StIdle 'StBusy
forall {k1} req1 (resp :: k1).
req1 -> Message (ReqResp req1 resp) 'StIdle 'StBusy
MsgReq req
req) (Client (ReqResp req resp) 'NonPipelined 'StBusy m a
 -> Client (ReqResp req resp) 'NonPipelined 'StIdle m a)
-> Client (ReqResp req resp) 'NonPipelined 'StBusy m a
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
forall a b. (a -> b) -> a -> b
$

    -- The type of our protocol means that we're now into the 'StBusy' state
    -- and the only thing we can do next is local effects or wait for a reply.
    -- We'll wait for a reply.
    (forall (st' :: ReqResp req resp).
 Message (ReqResp req resp) 'StBusy st'
 -> Client (ReqResp req resp) 'NonPipelined st' m a)
-> Client (ReqResp req resp) 'NonPipelined 'StBusy m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a.
(StateTokenI st, StateAgency st ~ 'ServerAgency,
 Outstanding pl ~ 'Z) =>
(forall (st' :: ps). Message ps st st' -> Client ps pl st' m a)
-> Client ps pl st m a
Await ((forall (st' :: ReqResp req resp).
  Message (ReqResp req resp) 'StBusy st'
  -> Client (ReqResp req resp) 'NonPipelined st' m a)
 -> Client (ReqResp req resp) 'NonPipelined 'StBusy m a)
-> (forall (st' :: ReqResp req resp).
    Message (ReqResp req resp) 'StBusy st'
    -> Client (ReqResp req resp) 'NonPipelined st' m a)
-> Client (ReqResp req resp) 'NonPipelined 'StBusy m a
forall a b. (a -> b) -> a -> b
$ \(MsgResp resp1
resp) ->

    -- Now in this case there is only one possible response, and we have
    -- one corresponding continuation 'kPong' to handle that response.
    -- The pong reply has no content so there's nothing to pass to our
    -- continuation, but if there were we would.
      m (Client (ReqResp req resp) 'NonPipelined st' m a)
-> Client (ReqResp req resp) 'NonPipelined st' m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a.
m (Client ps pl st m a) -> Client ps pl st m a
Effect (m (Client (ReqResp req resp) 'NonPipelined st' m a)
 -> Client (ReqResp req resp) 'NonPipelined st' m a)
-> m (Client (ReqResp req resp) 'NonPipelined st' m a)
-> Client (ReqResp req resp) 'NonPipelined st' m a
forall a b. (a -> b) -> a -> b
$ do
        client <- resp -> m (ReqRespClient req resp m a)
next resp
resp1
resp
        pure $ reqRespClientPeer client



requestOnce :: forall req resp m.
               Monad m
            => (forall x. Server (ReqResp req resp) NonPipelined StIdle m x)
            -> (req -> m resp)
requestOnce :: forall req resp (m :: * -> *).
Monad m =>
(forall x. Server (ReqResp req resp) 'NonPipelined 'StIdle m x)
-> req -> m resp
requestOnce forall x. Server (ReqResp req resp) 'NonPipelined 'StIdle m x
server req
req = (\(resp
resp, Any
_, TerminalStates (ReqResp req resp)
_) -> resp
resp)
                     ((resp, Any, TerminalStates (ReqResp req resp)) -> resp)
-> m (resp, Any, TerminalStates (ReqResp req resp)) -> m resp
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ReqRespClient req resp m resp
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m resp
forall (m :: * -> *) req resp a.
Monad m =>
ReqRespClient req resp m a
-> Client (ReqResp req resp) 'NonPipelined 'StIdle m a
reqRespClientPeer ReqRespClient req resp m resp
client Client (ReqResp req resp) 'NonPipelined 'StIdle m resp
-> Peer
     (ReqResp req resp)
     (FlipAgency 'AsClient)
     'NonPipelined
     'StIdle
     m
     Any
-> m (resp, Any, TerminalStates (ReqResp req resp))
forall ps (pr :: PeerRole) (initSt :: ps) (m :: * -> *) a b.
(Monad m, SingI pr) =>
Peer ps pr 'NonPipelined initSt m a
-> Peer ps (FlipAgency pr) 'NonPipelined initSt m b
-> m (a, b, TerminalStates ps)
`connect` Peer
  (ReqResp req resp)
  (FlipAgency 'AsClient)
  'NonPipelined
  'StIdle
  m
  Any
Server (ReqResp req resp) 'NonPipelined 'StIdle m Any
forall x. Server (ReqResp req resp) 'NonPipelined 'StIdle m x
server
  where
    client :: ReqRespClient req resp m resp
    client :: ReqRespClient req resp m resp
client = req
-> (resp -> m (ReqRespClient req resp m resp))
-> ReqRespClient req resp m resp
forall req resp (m :: * -> *) a.
req
-> (resp -> m (ReqRespClient req resp m a))
-> ReqRespClient req resp m a
SendMsgReq req
req ((resp -> m (ReqRespClient req resp m resp))
 -> ReqRespClient req resp m resp)
-> (resp -> m (ReqRespClient req resp m resp))
-> ReqRespClient req resp m resp
forall a b. (a -> b) -> a -> b
$ \resp
resp -> ReqRespClient req resp m resp -> m (ReqRespClient req resp m resp)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ReqRespClient req resp m resp
 -> m (ReqRespClient req resp m resp))
-> ReqRespClient req resp m resp
-> m (ReqRespClient req resp m resp)
forall a b. (a -> b) -> a -> b
$ m resp -> ReqRespClient req resp m resp
forall (m :: * -> *) a req resp. m a -> ReqRespClient req resp m a
SendMsgDone (resp -> m resp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure resp
resp)


--
-- Pipelined client
--

-- | A request-response client designed for running the 'ReqResp' protocol in
-- a pipelined way.
--
data ReqRespClientPipelined req resp m a where
  -- | A 'PingPongSender', but starting with zero outstanding pipelined
  -- responses, and for any internal collect type @c@.
  ReqRespClientPipelined ::
      ReqRespIdle            req resp Z c m a
   -> ReqRespClientPipelined req resp     m a


data ReqRespIdle req resp n c m a where
  -- | Send a `Req` message but alike in `ReqRespClient` do not await for the
  -- resopnse, instead supply a monadic action which will run on a received
  -- `Pong` message.
  SendMsgReqPipelined
    :: req
    -> (resp -> m c)                     -- receive action
    -> ReqRespIdle req resp (S n) c m a  -- continuation
    -> ReqRespIdle req resp    n  c m a

  CollectPipelined
    :: Maybe   (ReqRespIdle req resp (S n) c m a)
    -> (c -> m (ReqRespIdle req resp    n  c m a))
    ->          ReqRespIdle req resp (S n) c m a

  -- | Termination of the req-resp protocol.
  SendMsgDonePipelined
    :: a -> ReqRespIdle req resp Z c m a


-- | Interpret a pipelined client as a 'Peer' on the client side of
-- the 'ReqResp' protocol.
--
reqRespClientPeerPipelined
  :: Functor m
  => ReqRespClientPipelined req resp           m a
  -> ClientPipelined (ReqResp req resp) StIdle m a
reqRespClientPeerPipelined :: forall (m :: * -> *) req resp a.
Functor m =>
ReqRespClientPipelined req resp m a
-> ClientPipelined (ReqResp req resp) 'StIdle m a
reqRespClientPeerPipelined (ReqRespClientPipelined ReqRespIdle req resp 'Z c m a
peer) =
    Client (ReqResp req resp) ('Pipelined 'Z c) 'StIdle m a
-> ClientPipelined (ReqResp req resp) 'StIdle m a
forall ps (st :: ps) (m :: * -> *) a c.
Client ps ('Pipelined 'Z c) st m a -> ClientPipelined ps st m a
ClientPipelined (Client (ReqResp req resp) ('Pipelined 'Z c) 'StIdle m a
 -> ClientPipelined (ReqResp req resp) 'StIdle m a)
-> Client (ReqResp req resp) ('Pipelined 'Z c) 'StIdle m a
-> ClientPipelined (ReqResp req resp) 'StIdle m a
forall a b. (a -> b) -> a -> b
$ ReqRespIdle req resp 'Z c m a
-> Client (ReqResp req resp) ('Pipelined 'Z c) 'StIdle m a
forall req resp (n :: N) c (m :: * -> *) a.
Functor m =>
ReqRespIdle req resp n c m a
-> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a
reqRespClientPeerIdle ReqRespIdle req resp 'Z c m a
peer


reqRespClientPeerIdle
  :: forall req resp n c m a.
     Functor m
  => ReqRespIdle   req resp               n c          m a
  -> Client (ReqResp req resp) (Pipelined n c) StIdle  m a

reqRespClientPeerIdle :: forall req resp (n :: N) c (m :: * -> *) a.
Functor m =>
ReqRespIdle req resp n c m a
-> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a
reqRespClientPeerIdle = ReqRespIdle req resp n c m a
-> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a
forall (n' :: N).
ReqRespIdle req resp n' c m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
go
  where
    go :: forall n'.
          ReqRespIdle   req resp               n' c         m a
       -> Client (ReqResp req resp) (Pipelined n' c) StIdle m a

    go :: forall (n' :: N).
ReqRespIdle req resp n' c m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
go (SendMsgReqPipelined req
req resp -> m c
receive ReqRespIdle req resp ('S n') c m a
next) =
      -- Pipelined yield: send `MsgReq`, immediately follow with the next step.
      -- Await for a response in a continuation.
      Message (ReqResp req resp) 'StIdle 'StBusy
-> Receiver (ReqResp req resp) 'StBusy 'StIdle m c
-> Client (ReqResp req resp) ('Pipelined ('S n') c) 'StIdle m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
forall ps (st :: ps) (n :: N) c (m :: * -> *) a (st' :: ps)
       (st'' :: ps).
(StateTokenI st, StateTokenI st',
 StateAgency st ~ 'ClientAgency) =>
Message ps st st'
-> Receiver ps st' st'' m c
-> Client ps ('Pipelined ('S n) c) st'' m a
-> Client ps ('Pipelined n c) st m a
YieldPipelined
        (req -> Message (ReqResp req resp) 'StIdle 'StBusy
forall {k1} req1 (resp :: k1).
req1 -> Message (ReqResp req1 resp) 'StIdle 'StBusy
MsgReq req
req)
        ((forall (st' :: ReqResp req resp).
 Message (ReqResp req resp) 'StBusy st'
 -> Receiver (ReqResp req resp) st' 'StIdle m c)
-> Receiver (ReqResp req resp) 'StBusy 'StIdle m c
forall ps (st :: ps) (stdone :: ps) (m :: * -> *) c.
(StateTokenI st, ActiveState st, StateAgency st ~ 'ServerAgency) =>
(forall (st' :: ps).
 Message ps st st' -> Receiver ps st' stdone m c)
-> Receiver ps st stdone m c
ReceiverAwait  ((forall (st' :: ReqResp req resp).
  Message (ReqResp req resp) 'StBusy st'
  -> Receiver (ReqResp req resp) st' 'StIdle m c)
 -> Receiver (ReqResp req resp) 'StBusy 'StIdle m c)
-> (forall (st' :: ReqResp req resp).
    Message (ReqResp req resp) 'StBusy st'
    -> Receiver (ReqResp req resp) st' 'StIdle m c)
-> Receiver (ReqResp req resp) 'StBusy 'StIdle m c
forall a b. (a -> b) -> a -> b
$ \(MsgResp resp1
resp) ->
          m (Receiver (ReqResp req resp) st' 'StIdle m c)
-> Receiver (ReqResp req resp) st' 'StIdle m c
forall ps (st :: ps) (stdone :: ps) (m :: * -> *) c.
m (Receiver ps st stdone m c) -> Receiver ps st stdone m c
ReceiverEffect (m (Receiver (ReqResp req resp) st' 'StIdle m c)
 -> Receiver (ReqResp req resp) st' 'StIdle m c)
-> m (Receiver (ReqResp req resp) st' 'StIdle m c)
-> Receiver (ReqResp req resp) st' 'StIdle m c
forall a b. (a -> b) -> a -> b
$
            c -> Receiver (ReqResp req resp) st' st' m c
c -> Receiver (ReqResp req resp) st' 'StIdle m c
forall ps (stdone :: ps) (m :: * -> *) c.
c -> Receiver ps stdone stdone m c
ReceiverDone (c -> Receiver (ReqResp req resp) st' 'StIdle m c)
-> m c -> m (Receiver (ReqResp req resp) st' 'StIdle m c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> resp -> m c
receive resp
resp1
resp
        )
        (ReqRespIdle req resp ('S n') c m a
-> Client (ReqResp req resp) ('Pipelined ('S n') c) 'StIdle m a
forall (n' :: N).
ReqRespIdle req resp n' c m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
go ReqRespIdle req resp ('S n') c m a
next)

    go (CollectPipelined Maybe (ReqRespIdle req resp ('S n) c m a)
mNone c -> m (ReqRespIdle req resp n c m a)
collect) =
      Maybe (Client (ReqResp req resp) ('Pipelined ('S n) c) 'StIdle m a)
-> (c -> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
-> Client (ReqResp req resp) ('Pipelined ('S n) c) 'StIdle m a
forall ps (st :: ps) (n :: N) c (m :: * -> *) a.
(StateTokenI st, ActiveState st) =>
Maybe (Client ps ('Pipelined ('S n) c) st m a)
-> (c -> Client ps ('Pipelined n c) st m a)
-> Client ps ('Pipelined ('S n) c) st m a
Collect
        (ReqRespIdle req resp ('S n) c m a
-> Client (ReqResp req resp) ('Pipelined ('S n) c) 'StIdle m a
forall (n' :: N).
ReqRespIdle req resp n' c m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
go (ReqRespIdle req resp ('S n) c m a
 -> Client (ReqResp req resp) ('Pipelined ('S n) c) 'StIdle m a)
-> Maybe (ReqRespIdle req resp ('S n) c m a)
-> Maybe
     (Client (ReqResp req resp) ('Pipelined ('S n) c) 'StIdle m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (ReqRespIdle req resp ('S n) c m a)
mNone)
        (\c
c -> m (Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
-> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a.
m (Client ps pl st m a) -> Client ps pl st m a
Effect (m (Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
 -> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
-> m (Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
-> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a
forall a b. (a -> b) -> a -> b
$ ReqRespIdle req resp n c m a
-> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a
forall (n' :: N).
ReqRespIdle req resp n' c m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
go (ReqRespIdle req resp n c m a
 -> Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
-> m (ReqRespIdle req resp n c m a)
-> m (Client (ReqResp req resp) ('Pipelined n c) 'StIdle m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> c -> m (ReqRespIdle req resp n c m a)
collect c
c)

    go (SendMsgDonePipelined a
result) =
      -- Send `MsgDone` and complete the protocol
      Message (ReqResp req resp) 'StIdle 'StDone
-> Client (ReqResp req resp) ('Pipelined n' c) 'StDone m a
-> Client (ReqResp req resp) ('Pipelined n' c) 'StIdle m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a
       (st' :: ps).
(StateTokenI st, StateTokenI st', StateAgency st ~ 'ClientAgency,
 Outstanding pl ~ 'Z) =>
Message ps st st' -> Client ps pl st' m a -> Client ps pl st m a
Yield
        Message (ReqResp req resp) 'StIdle 'StDone
forall {k} {k1} (req :: k) (resp :: k1).
Message (ReqResp req resp) 'StIdle 'StDone
MsgDone
        (a -> Client (ReqResp req resp) ('Pipelined n' c) 'StDone m a
forall ps (pl :: IsPipelined) (st :: ps) (m :: * -> *) a.
(StateTokenI st, StateAgency st ~ 'NobodyAgency,
 Outstanding pl ~ 'Z) =>
a -> Client ps pl st m a
Done a
result)