{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module TypedSession.Core where

import Control.Concurrent.Class.MonadSTM (MonadSTM (..))
import Data.IFunctor
import Data.IntMap (IntMap)
import Data.Kind

{- |

typed-session is a communication framework.
The messages received from the outside will be placed in MsgCache.
When interpreting the Peer, (Sing (r :: s)) will be generated according to the Peer's status.
SingToInt can convert Sing (r :: s) to Int.
Depending on this Int value, the required Msg is finally found from MsgCache.

In the process of multi-role communication, a message cache structure like MsgCache is needed.

Consider the following scenario

@
s1   s1   s2       Initial state
------------
a -> b             a sends message MsgA to b
------------
s3   s2   s2       State after sending
------------
     b <- c        c sends message MsgC to b
------------
s3   s4   s5       State after sending
@

For b, due to the influence of network transmission, it cannot guarantee that it can receive MsgA first and then MsgC.
If it receives MsgC first, it will directly put it in MsgCache and continue to wait until
MsgA arrives, then it will start to process MsgA first and then MsgC.

In general, dataToTag# is used directly here.

Example:

@
instance SingToInt Role where
  singToInt x = I# (dataToTag# x)

instance SingToInt PingPong where
  singToInt x = I# (dataToTag# x)
@
-}
class SingToInt s where
  singToInt :: Sing (r :: s) -> Int

{- |

Describes the type class of Msg. The core of typed-session.

@
type Done (sr :: role') :: ps
@

Describe the state of each role when it terminates.

@
data
  Msg
    role'
    ps
    (fromSt :: ps)
    (sender :: role')
    (senderNewSt :: ps)
    (receiver :: role')
    (receiverNewSt :: ps)
@
* role': the type of the role.
* ps: the type of the state machine.
* ​​fromSt: when sending a message, the sender is in this state,
where the receiver may be in this state, or a more generalized state related to this state.
For example, the sender is in state (S1 [True]), and the receiver is in state (S1 s).
* sender: the role that sends the message
* senderNewSt: the state of the role after sending the message
* receiver: the role that receives the message
* receiverNewSt: the state of the role after receiving the message

There are two principles to follow when designing the state of Msg:

1. When sending a message, the sender and receiver must be in the same state. Here the receiver may be in a more generalized state related to the state.
For example, the sender is in state (S1 [True]), and the receiver is in state (S1 s).

2. The same state can only be used for the same pair of receiver and sender.

For example, in the following example, state s1 is used for both (a -> b) and (b -> c), which is wrong.

@
s1   s1   s1
a -> b
s2   s1   s1
     b -> c
s2   s4   s5
@
-}
class (SingToInt role', SingToInt ps) => Protocol role' ps where
  type Done (sr :: role') :: ps
  data
    Msg
      role'
      ps
      (fromSt :: ps)
      (sender :: role')
      (senderNewSt :: ps)
      (receiver :: role')
      (receiverNewSt :: ps)

{- |
Messages received from the outside are placed in MsgCache. When interpreting
Peer will use the Msg in MsgCache.
-}
type MsgCache role' ps n = TVar n (IntMap (AnyMsg role' ps))

{- |
Packaging of Msg, shielding part of the type information, mainly used for serialization.
-}
data AnyMsg role' ps where
  AnyMsg
    :: ( SingI recv
       , SingI st
       , SingToInt role'
       , SingToInt ps
       )
    => Msg role' ps st send st' recv st''
    -> AnyMsg role' ps

msgFromStSing
  :: forall role' ps st send recv st' st''
   . (SingI recv, SingI st)
  => Msg role' ps st send st' recv st''
  -> Sing st
msgFromStSing :: forall role' ps (st :: ps) (send :: role') (recv :: role')
       (st' :: ps) (st'' :: ps).
(SingI recv, SingI st) =>
Msg role' ps st send st' recv st'' -> Sing st
msgFromStSing Msg role' ps st send st' recv st''
_ = forall (a :: ps). SingI a => Sing a
forall {k} (a :: k). SingI a => Sing a
sing @st

{- |
Core Ast, all we do is build this Ast and then interpret it.

@
IReturn :: ia st -> Peer role' ps r m ia st
@
IReturn indicates the termination of the continuation.

@
LiftM :: m (Peer role' ps r m ia st') -> Peer role' ps r m ia st
@

Liftm can transform state st to any state st'.
It looks a bit strange, as if it is a constructor that is not constrained by the Msg type.
Be careful when using it, it is a type breakpoint.
But some state transition functions need it, which can make the code more flexible.
Be very careful when using it!

@
Yield
  :: ( SingI recv
     , SingI from
     , SingToInt ps
     )
  => Msg role' ps from '(send, sps) '(recv, rps)
  -> Peer role' ps send m ia sps
  -> Peer role' ps send m ia from
@
Yield represents sending a message. Note that the Peer status changes from `from` to `sps`.

@
Await
  :: ( SingI recv
     , SingI from
     , SingToInt ps
     )
  => (Msg role' ps from send sps recv ~> Peer role' ps recv m ia)
  -> Peer role' ps recv m ia from
@

Await represents receiving messages.
Different messages will lead to different states.
The state is passed to the next behavior through (~>).
-}
data Peer role' ps (r :: role') (m :: Type -> Type) (ia :: ps -> Type) (st :: ps) where
  IReturn :: ia st -> Peer role' ps r m ia st
  LiftM :: m (Peer role' ps r m ia st') -> Peer role' ps r m ia st
  Yield
    :: ( SingI recv
       , SingI from
       , SingToInt ps
       )
    => Msg role' ps from send sps recv rps
    -> Peer role' ps send m ia sps
    -> Peer role' ps send m ia from
  Await
    :: ( SingI recv
       , SingI from
       , SingToInt ps
       )
    => (Msg role' ps from send sps recv ~> Peer role' ps recv m ia)
    -> Peer role' ps recv m ia from

instance (Functor m) => IMonadFail (Peer role' ps r m) where
  fail :: forall (a :: ps -> *) (ix :: ps). String -> Peer role' ps r m a ix
fail = String -> Peer role' ps r m a ix
forall a. HasCallStack => String -> a
error

instance (Functor m) => IFunctor (Peer role' ps r m) where
  imap :: forall (a :: ps -> *) (b :: ps -> *).
(a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
imap a ~> b
f = \case
    IReturn a x
ia -> b x -> Peer role' ps r m b x
forall ps (ia :: ps -> *) (st :: ps) role' (r :: role')
       (m :: * -> *).
ia st -> Peer role' ps r m ia st
IReturn (a x -> b x
a ~> b
f a x
ia)
    LiftM m (Peer role' ps r m a st')
f' -> m (Peer role' ps r m b st') -> Peer role' ps r m b x
forall (m :: * -> *) role' ps (r :: role') (ia :: ps -> *)
       (send :: ps) (st :: ps).
m (Peer role' ps r m ia send) -> Peer role' ps r m ia st
LiftM ((Peer role' ps r m a st' -> Peer role' ps r m b st')
-> m (Peer role' ps r m a st') -> m (Peer role' ps r m b st')
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
forall {k} {k1} (f :: (k -> *) -> k1 -> *) (a :: k -> *)
       (b :: k -> *).
IFunctor f =>
(a ~> b) -> f a ~> f b
forall (a :: ps -> *) (b :: ps -> *).
(a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
imap a x -> b x
a ~> b
f) m (Peer role' ps r m a st')
f')
    Yield Msg role' ps x r sps recv rps
ms Peer role' ps r m a sps
cont -> Msg role' ps x r sps recv rps
-> Peer role' ps r m b sps -> Peer role' ps r m b x
forall role' (send :: role') ps (from :: ps) (send :: role')
       (sps :: ps) (rps :: ps) (m :: * -> *) (ia :: ps -> *).
(SingI send, SingI from, SingToInt ps) =>
Msg role' ps from send sps send rps
-> Peer role' ps send m ia sps -> Peer role' ps send m ia from
Yield Msg role' ps x r sps recv rps
ms ((a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
forall {k} {k1} (f :: (k -> *) -> k1 -> *) (a :: k -> *)
       (b :: k -> *).
IFunctor f =>
(a ~> b) -> f a ~> f b
forall (a :: ps -> *) (b :: ps -> *).
(a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
imap a x -> b x
a ~> b
f Peer role' ps r m a sps
cont)
    Await Msg role' ps x send sps r ~> Peer role' ps r m a
cont -> (Msg role' ps x send sps r ~> Peer role' ps r m b)
-> Peer role' ps r m b x
forall role' (recv :: role') ps (from :: ps) (send :: role')
       (sps :: ps) (m :: * -> *) (ia :: ps -> *).
(SingI recv, SingI from, SingToInt ps) =>
(Msg role' ps from send sps recv ~> Peer role' ps recv m ia)
-> Peer role' ps recv m ia from
Await ((a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
forall {k} {k1} (f :: (k -> *) -> k1 -> *) (a :: k -> *)
       (b :: k -> *).
IFunctor f =>
(a ~> b) -> f a ~> f b
forall (a :: ps -> *) (b :: ps -> *).
(a ~> b) -> Peer role' ps r m a ~> Peer role' ps r m b
imap a x -> b x
a ~> b
f (Peer role' ps r m a x -> Peer role' ps r m b x)
-> (Msg role' ps x send sps r x -> Peer role' ps r m a x)
-> Msg role' ps x send sps r x
-> Peer role' ps r m b x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Msg role' ps x send sps r x -> Peer role' ps r m a x
Msg role' ps x send sps r ~> Peer role' ps r m a
cont)

instance (Functor m) => IMonad (Peer role' ps r m) where
  ireturn :: forall (a :: ps -> *). a ~> Peer role' ps r m a
ireturn = a x -> Peer role' ps r m a x
forall ps (ia :: ps -> *) (st :: ps) role' (r :: role')
       (m :: * -> *).
ia st -> Peer role' ps r m ia st
IReturn
  ibind :: forall (a :: ps -> *) (b :: ps -> *).
(a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
ibind a ~> Peer role' ps r m b
f = \case
    IReturn a x
ia -> (a x -> Peer role' ps r m b x
a ~> Peer role' ps r m b
f a x
ia)
    LiftM m (Peer role' ps r m a st')
f' -> m (Peer role' ps r m b st') -> Peer role' ps r m b x
forall (m :: * -> *) role' ps (r :: role') (ia :: ps -> *)
       (send :: ps) (st :: ps).
m (Peer role' ps r m ia send) -> Peer role' ps r m ia st
LiftM ((Peer role' ps r m a st' -> Peer role' ps r m b st')
-> m (Peer role' ps r m a st') -> m (Peer role' ps r m b st')
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
forall {k} (m :: (k -> *) -> k -> *) (a :: k -> *) (b :: k -> *).
IMonad m =>
(a ~> m b) -> m a ~> m b
forall (a :: ps -> *) (b :: ps -> *).
(a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
ibind a x -> Peer role' ps r m b x
a ~> Peer role' ps r m b
f) m (Peer role' ps r m a st')
f')
    Yield Msg role' ps x r sps recv rps
ms Peer role' ps r m a sps
cont -> Msg role' ps x r sps recv rps
-> Peer role' ps r m b sps -> Peer role' ps r m b x
forall role' (send :: role') ps (from :: ps) (send :: role')
       (sps :: ps) (rps :: ps) (m :: * -> *) (ia :: ps -> *).
(SingI send, SingI from, SingToInt ps) =>
Msg role' ps from send sps send rps
-> Peer role' ps send m ia sps -> Peer role' ps send m ia from
Yield Msg role' ps x r sps recv rps
ms ((a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
forall {k} (m :: (k -> *) -> k -> *) (a :: k -> *) (b :: k -> *).
IMonad m =>
(a ~> m b) -> m a ~> m b
forall (a :: ps -> *) (b :: ps -> *).
(a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
ibind a x -> Peer role' ps r m b x
a ~> Peer role' ps r m b
f Peer role' ps r m a sps
cont)
    Await Msg role' ps x send sps r ~> Peer role' ps r m a
cont -> (Msg role' ps x send sps r ~> Peer role' ps r m b)
-> Peer role' ps r m b x
forall role' (recv :: role') ps (from :: ps) (send :: role')
       (sps :: ps) (m :: * -> *) (ia :: ps -> *).
(SingI recv, SingI from, SingToInt ps) =>
(Msg role' ps from send sps recv ~> Peer role' ps recv m ia)
-> Peer role' ps recv m ia from
Await ((a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
forall {k} (m :: (k -> *) -> k -> *) (a :: k -> *) (b :: k -> *).
IMonad m =>
(a ~> m b) -> m a ~> m b
forall (a :: ps -> *) (b :: ps -> *).
(a ~> Peer role' ps r m b)
-> Peer role' ps r m a ~> Peer role' ps r m b
ibind a x -> Peer role' ps r m b x
a ~> Peer role' ps r m b
f (Peer role' ps r m a x -> Peer role' ps r m b x)
-> (Msg role' ps x send sps r x -> Peer role' ps r m a x)
-> Msg role' ps x send sps r x
-> Peer role' ps r m b x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Msg role' ps x send sps r x -> Peer role' ps r m a x
Msg role' ps x send sps r ~> Peer role' ps r m a
cont)

{- |
Send a message, the Peer status changes from `from` to `sps`.
-}
yield
  :: ( Functor m
     , SingI recv
     , SingI from
     , SingToInt ps
     )
  => Msg role' ps from send sps recv rps
  -> Peer role' ps send m (At () sps) from
yield :: forall (m :: * -> *) role' (recv :: role') ps (from :: ps)
       (send :: role') (sps :: ps) (rps :: ps).
(Functor m, SingI recv, SingI from, SingToInt ps) =>
Msg role' ps from send sps recv rps
-> Peer role' ps send m (At () sps) from
yield Msg role' ps from send sps recv rps
msg = Msg role' ps from send sps recv rps
-> Peer role' ps send m (At () sps) sps
-> Peer role' ps send m (At () sps) from
forall role' (send :: role') ps (from :: ps) (send :: role')
       (sps :: ps) (rps :: ps) (m :: * -> *) (ia :: ps -> *).
(SingI send, SingI from, SingToInt ps) =>
Msg role' ps from send sps send rps
-> Peer role' ps send m ia sps -> Peer role' ps send m ia from
Yield Msg role' ps from send sps recv rps
msg (() -> Peer role' ps send m (At () sps) sps
forall {k1} (m :: (k1 -> *) -> k1 -> *) a (k2 :: k1).
IMonad m =>
a -> m (At a k2) k2
returnAt ())

{- |
Receiving Messages.
-}
await
  :: ( Functor m
     , SingI recv
     , SingI from
     , SingToInt ps
     )
  => Peer role' ps recv m (Msg role' ps from send sps recv) from
await :: forall (m :: * -> *) role' (recv :: role') ps (from :: ps)
       (send :: role') (sps :: ps).
(Functor m, SingI recv, SingI from, SingToInt ps) =>
Peer role' ps recv m (Msg role' ps from send sps recv) from
await = (Msg role' ps from send sps recv
 ~> Peer role' ps recv m (Msg role' ps from send sps recv))
-> Peer role' ps recv m (Msg role' ps from send sps recv) from
forall role' (recv :: role') ps (from :: ps) (send :: role')
       (sps :: ps) (m :: * -> *) (ia :: ps -> *).
(SingI recv, SingI from, SingToInt ps) =>
(Msg role' ps from send sps recv ~> Peer role' ps recv m ia)
-> Peer role' ps recv m ia from
Await Msg role' ps from send sps recv x
-> Peer role' ps recv m (Msg role' ps from send sps recv) x
Msg role' ps from send sps recv
~> Peer role' ps recv m (Msg role' ps from send sps recv)
forall {k} (m :: (k -> *) -> k -> *) (a :: k -> *).
IMonad m =>
a ~> m a
forall (a :: ps -> *). a ~> Peer role' ps recv m a
ireturn

{- |
Lift any m to Peer role' ps r m, which is an application of LiftM.
Note that the state of `ts` has not changed.
-}
liftm :: (Functor m) => m a -> Peer role' ps r m (At a ts) ts
liftm :: forall (m :: * -> *) a role' ps (r :: role') (ts :: ps).
Functor m =>
m a -> Peer role' ps r m (At a ts) ts
liftm m a
m = m (Peer role' ps r m (At a ts) ts)
-> Peer role' ps r m (At a ts) ts
forall (m :: * -> *) role' ps (r :: role') (ia :: ps -> *)
       (send :: ps) (st :: ps).
m (Peer role' ps r m ia send) -> Peer role' ps r m ia st
LiftM (a -> Peer role' ps r m (At a ts) ts
forall {k1} (m :: (k1 -> *) -> k1 -> *) a (k2 :: k1).
IMonad m =>
a -> m (At a k2) k2
returnAt (a -> Peer role' ps r m (At a ts) ts)
-> m a -> m (Peer role' ps r m (At a ts) ts)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
m)

liftConstructor :: (Applicative m) => ia st' -> Peer role' k r m ia st
liftConstructor :: forall (m :: * -> *) k (ia :: k -> *) (st' :: k) role' (r :: role')
       (st :: k).
Applicative m =>
ia st' -> Peer role' k r m ia st
liftConstructor ia st'
a = m (Peer role' k r m ia st') -> Peer role' k r m ia st
forall (m :: * -> *) role' ps (r :: role') (ia :: ps -> *)
       (send :: ps) (st :: ps).
m (Peer role' ps r m ia send) -> Peer role' ps r m ia st
LiftM (m (Peer role' k r m ia st') -> Peer role' k r m ia st)
-> m (Peer role' k r m ia st') -> Peer role' k r m ia st
forall a b. (a -> b) -> a -> b
$ Peer role' k r m ia st' -> m (Peer role' k r m ia st')
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Peer role' k r m ia st' -> m (Peer role' k r m ia st'))
-> Peer role' k r m ia st' -> m (Peer role' k r m ia st')
forall a b. (a -> b) -> a -> b
$ ia st' -> Peer role' k r m ia st'
forall {k} (m :: (k -> *) -> k -> *) (a :: k -> *).
IMonad m =>
a ~> m a
forall (a :: k -> *). a ~> Peer role' k r m a
ireturn ia st'
a