Explicit pipelining without concurrency

Pipelined.hs

{-# LANGUAGE DataKinds                #-}
{-# LANGUAGE DerivingStrategies       #-}
{-# LANGUAGE DerivingVia              #-}
{-# LANGUAGE EmptyCase                #-}
{-# LANGUAGE FlexibleInstances        #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE PatternSynonyms          #-}
{-# LANGUAGE PolyKinds                #-}
{-# LANGUAGE RankNTypes               #-}
{-# LANGUAGE StandaloneDeriving       #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeFamilies             #-}
{-# LANGUAGE TypeOperators            #-}
{-# LANGUAGE ViewPatterns             #-}

module Pipelined2 where

import           Data.Kind (Type)
-- import           Data.Typeable
import           Data.Void

--
-- Interface copied from 'typed-protocol' package
--

class Protocol ps where

  -- | The messages for this protocol. It is expected to be a GADT that is
  -- indexed by the @from@ and @to@ protocol states. That is the protocol state
  -- the message transitions from, and the protocol state it transitions into.
  -- These are the edges of the protocol state transition system.
  --
  data Message ps (st :: ps) (st' :: ps)

  -- | Tokens for those protocol states in which the client has agency.
  --
  data ClientHasAgency (st :: ps)

  -- | Tokens for those protocol states in which the server has agency.
  --
  data ServerHasAgency (st :: ps)

  -- | Tokens for terminal protocol states in which neither the client nor
  -- server has agency.
  --
  data NobodyHasAgency (st :: ps)

  -- | Lemma that if the client has agency for a state, there are no
  -- cases in which the server has agency for the same state.
  --
  exclusionLemma_ClientAndServerHaveAgency
    :: forall (st :: ps).
       ClientHasAgency st
    -> ServerHasAgency st
    -> Void

  -- | Lemma that if the nobody has agency for a state, there are no
  -- cases in which the client has agency for the same state.
  --
  exclusionLemma_NobodyAndClientHaveAgency
    :: forall (st :: ps).
       NobodyHasAgency st
    -> ClientHasAgency st
    -> Void

  -- | Lemma that if the nobody has agency for a state, there are no
  -- cases in which the server has agency for the same state.
  --
  exclusionLemma_NobodyAndServerHaveAgency
    :: forall (st :: ps).
       NobodyHasAgency st
    -> ServerHasAgency st
    -> Void

data PeerRole = AsClient | AsServer

data PeerHasAgency (pr :: PeerRole) (st :: ps) where
    ClientAgency :: !(ClientHasAgency st) -> PeerHasAgency AsClient st
    ServerAgency :: !(ServerHasAgency st) -> PeerHasAgency AsServer st

type WeHaveAgency   (pr :: PeerRole) st = PeerHasAgency             pr  st
type TheyHaveAgency (pr :: PeerRole) st = PeerHasAgency (FlipAgency pr) st
type family FlipAgency (pr :: PeerRole) where
    FlipAgency AsClient = AsServer
    FlipAgency AsServer = AsClient


--
-- Explicit pipelining without concurrency.
--

data Trans ps where
    Tr :: forall ps. ps -> ps -> Trans ps


type        Snoc :: [k]    -> k -> [k]
type family Snoc    as        a where
            Snoc    '[]       b = '[b]
            Snoc    (a ': as) b = a ': Snoc as b

type        Uncons :: ps -> Trans ps -> [Trans ps] -> [Trans ps]
type family Uncons stNext a as where
  Uncons stNext (Tr st stNext) as = as
  Uncons stNext (Tr st st')    as = Tr stNext st' : as

type PeerPipelined :: forall ps -> PeerRole -> ps -> [Trans ps] -> (Type -> Type) -> Type
                   -> Type
data PeerPipelined ps pr st tr m a where

    Effect  :: m (PeerPipelined ps pr st tr m a)
            ->    PeerPipelined ps pr st tr m a

    Yield   :: !(WeHaveAgency pr st)
            -> Message ps st st'
            -> PeerPipelined ps pr st' '[] m a
            -> PeerPipelined ps pr st  '[] m a

    Await   :: !(TheyHaveAgency pr st)
            -> (forall st'. Message ps st st'
                         -> PeerPipelined ps pr st' '[] m a)
            -> PeerPipelined              ps pr st  '[] m a

    -- | Push onto collect queue.
    --
    YieldPipelined
            :: !(WeHaveAgency pr st)
            -> Message ps st st'
            -> PeerPipelined ps pr (st'' :: ps) (Snoc tr (Tr st' st'')) m a
            -> PeerPipelined ps pr st tr m a

    -- | Collect pushed @Tr@.
    --
    -- This could be expressed as 'CollectPartial' and 'CollectDone' (see
    -- 'collect'), if we didn't use higher rank polymorphism in
    -- 'CollectParital'.
    --
    -- This is used by 'TxSubmission' mini-protocol.
    --
    Collect :: Maybe (PeerPipelined ps pr (st :: ps) (Tr st'    st'' ': tr) m a)
            -> !(TheyHaveAgency pr st')
            -> (Message ps st' st''
                   -> PeerPipelined ps pr (st :: ps)                       tr  m a)
            -> PeerPipelined        ps pr (st :: ps) (Tr st'    st'' ': tr) m a

    -- | Parially collect pushed @Tr@.
    --
    -- We use this in 'BlockFetch' mini-protocol.
    --
    CollectPartial
            :: Maybe (PeerPipelined ps pr (st :: ps) (Tr st'    st'' ': tr) m a)
            -> !(TheyHaveAgency pr st')
            -> (forall stNext. Message ps st' stNext
                   -> PeerPipelined ps pr (st :: ps) (Tr stNext st'' ': tr) m a)
            -> PeerPipelined        ps pr (st :: ps) (Tr st'    st'' ': tr) m a

    -- | Pop identity 'Tr' from the collect queue.
    --
    -- 'CollectDone' allows to defer poping @Tr ps st st@ from the queue after
    -- a message is received (in 'CollectPartial' callback), unlike 'Collect'
    -- which needs to know the transition type at compile time.
    --
    CollectDone
            :: PeerPipelined ps pr (st :: ps)                tr  m a 
            -> PeerPipelined ps pr (st :: ps) (Tr st' st' ': tr) m a

    Done    :: !(NobodyHasAgency st)
            -> a
            -> PeerPipelined ps pr st '[] m a


-- | This is enough to support pipelining in 'TxSubmission' mini-protocol,
-- but not enough to support pipelining in 'BlockFetch'.
--
-- 'collect' type  checks if 'CollectPartial' does not use ranked-n-types in
-- the continuation.
--
{--
  - collect :: Maybe (PeerPipelined ps pr (st :: ps) (Tr st' st'' ': tr) m a)
  -         -> TheyHaveAgency pr st'
  -         -> (Message ps st' st''
  -                -> PeerPipelined ps pr (st :: ps)                 tr  m a)
  -         -> PeerPipelined        ps pr (st :: ps) (Tr st' st'' ': tr) m a
  - collect more tok k =
  -     CollectPartial more tok $ \msg ->
  -       CollectDone (k msg)
  --}

--
-- Providing a 'Collect' pattern is not streight forward becuase @stNext@ is
-- umibgous in @unlift -> k@ pattern, if if 'CollectPartial' is not  
--

{--
  - unlift :: forall ps (pr :: PeerRole) (st     :: ps)
  -                                      (st'    :: ps)
  -                                      (st''   :: ps)
  -                                      (stNext :: ps)
  -                                      tr
  -                                      m a.
  -           (Typeable st'', Typeable stNext)
  -        => (Message ps st' stNext
  -             -> PeerPipelined ps pr st (Tr stNext st'' ': tr) m a)
  -        -> Message ps st' st''
  -        -> PeerPipelined ps pr st tr  m a
  - unlift k = \msg ->
  -   case eqT :: Maybe (st'' :~: stNext) of
  -     Just Refl ->
  -       case k msg of
  -         CollectDone k -> k
  --}


{--
  - pattern Collect :: Typeable st''
  -                 => Maybe (PeerPipelined ps pr (st :: ps) (Tr st' st'' ': tr) m a)
  -                 -> TheyHaveAgency pr st'
  -                 -> (Message ps st' st''
  -                        -> PeerPipelined ps pr (st :: ps)                 tr  m a)
  -                 -> PeerPipelined        ps pr (st :: ps) (Tr st' st'' ': tr) m a
  - pattern Collect more tok k <- CollectPartial more tok (unlift -> k)
  --}


--
-- PingPong Example
--

data PingPong where
  StIdle     :: PingPong
  StBusy     :: PingPong
  StVeryBusy :: PingPong
  StDone     :: PingPong

instance Protocol PingPong where
  data Message PingPong from to where
    MsgPing :: Message PingPong StIdle StBusy
    MsgPong :: Message PingPong StBusy StIdle
    MsgDone :: Message PingPong StIdle StDone

    MsgPingV    :: Message PingPong StIdle     StVeryBusy
    MsgPongV    :: Message PingPong StVeryBusy StVeryBusy
    MsgPongDone :: Message PingPong StVeryBusy StIdle

  data ClientHasAgency st where
    TokIdle :: ClientHasAgency StIdle

  data ServerHasAgency st where
    TokBusy     :: ServerHasAgency StBusy
    TokVeryBusy :: ServerHasAgency StVeryBusy

  data NobodyHasAgency st where
    TokDone :: NobodyHasAgency StDone

  exclusionLemma_ClientAndServerHaveAgency TokIdle tok = case tok of {}
  exclusionLemma_NobodyAndClientHaveAgency TokDone tok = case tok of {}
  exclusionLemma_NobodyAndServerHaveAgency TokDone tok = case tok of {}


deriving instance Show (Message PingPong from to)

instance Show (ClientHasAgency (st :: PingPong)) where
  show TokIdle = "TokIdle"

instance Show (ServerHasAgency (st :: PingPong)) where
  show TokBusy     = "TokBusy"
  show TokVeryBusy = "TokVeryBusy"


-- A non-pipelined PingPong.
--
pingPongClient :: a -> PeerPipelined PingPong AsClient StIdle '[] m a
pingPongClient a =
    Yield (ClientAgency TokIdle) MsgPing
  $ Await (ServerAgency TokBusy)
  $ \MsgPong ->
    Yield (ClientAgency TokIdle) MsgPing
  $ Await (ServerAgency TokBusy)
  $ \MsgPong ->
    Yield (ClientAgency TokIdle) MsgDone
  $ Done TokDone a


-- A pipelined 'PingPong', without partial collects.
--
-- This pipelining pattern is used by 'TxSubmission' mini-protocol.
--
pingPongClientPipelined :: a -> PeerPipelined PingPong AsClient StIdle '[] m a
pingPongClientPipelined a =
    YieldPipelined (ClientAgency TokIdle) MsgPing
  $ YieldPipelined (ClientAgency TokIdle) MsgPing
  $ Collect Nothing (ServerAgency TokBusy)
  $ \MsgPong ->
    Collect Nothing (ServerAgency TokBusy)
  $ \MsgPong ->
    Yield (ClientAgency TokIdle) MsgDone
  $ Done TokDone a


-- | A pipelined 'PingPong' which supports partial collects using the recursive 'collectV'.
--
-- This pipelining pattern is used by 'BlockFetch' mini-protocol.
--
pingPongClientPipelinedV :: a -> PeerPipelined PingPong AsClient StIdle '[] m a
pingPongClientPipelinedV a =
      YieldPipelined (ClientAgency TokIdle) MsgPingV
    $ YieldPipelined (ClientAgency TokIdle) MsgPingV
    $ YieldPipelined (ClientAgency TokIdle) MsgPingV
    $ collectV
    $ collectV
    $ collectV
    $ Yield (ClientAgency TokIdle) MsgDone
    $ Done TokDone a
  where
    -- recursievly collect responses, until 'MsgPongDone' is received
    --
    collectV :: PeerPipelined PingPong AsClient StIdle                   tk  m a
             -- ^ continuation
             -> PeerPipelined PingPong AsClient StIdle    
                                       (Tr StVeryBusy StIdle ': tk) m a
    collectV k =
        CollectPartial Nothing (ServerAgency TokVeryBusy)
      $ \msg -> case msg of
          MsgPongV    -> collectV k
          MsgPongDone -> CollectDone k