{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
module Capnp.New.Rpc.Server
    ( CallHandler
    , MethodHandler
    , UntypedMethodHandler
    , Export(..)
    , export
    , findMethod

    , SomeServer(..)

    -- * Helpers for writing method handlers
    , handleParsed
    , handleRaw
    , methodUnimplemented

    , toUntypedMethodHandler

    -- * Internals; exposed only for use by generated code.
    , MethodHandlerTree(..)
    ) where

import           Capnp.Convert           (parsedToRaw)
import           Capnp.Message           (Mutability(..))
import qualified Capnp.New.Basics        as B
import qualified Capnp.New.Classes       as C
import qualified Capnp.Repr              as R
import           Capnp.Repr.Methods      (Client(..))
import           Capnp.Rpc.Errors
    (eFailed, eMethodUnimplemented, wrapException)
import           Capnp.Rpc.Promise
    (Fulfiller, breakPromise, fulfill, newCallback)
import qualified Capnp.Rpc.Server        as Legacy
import qualified Capnp.Rpc.Untyped       as URpc
import           Capnp.TraversalLimit    (defaultLimit, evalLimitT)
import qualified Capnp.Untyped           as U
import           Control.Exception.Safe  (withException)
import           Control.Monad.STM.Class (MonadSTM(..))
import           Data.Function           ((&))
import           Data.Kind               (Constraint)
import qualified Data.Map.Strict         as M
import           Data.Maybe              (fromMaybe)
import           Data.Proxy              (Proxy(..))
import           Data.Typeable           (Typeable)
import qualified Data.Vector             as V
import           Data.Word
import           GHC.Prim                (coerce)
import           Internal.BuildPure      (createPure)
import           Supervisors             (Supervisor)

-- | A handler for arbitrary RPC calls. Maps (interfaceId, methodId) pairs to
-- 'UntypedMethodHandler's.
type CallHandler = M.Map Word64 (V.Vector UntypedMethodHandler)

-- | Type alias for a handler for a particular rpc method.
type MethodHandler p r
    = R.Raw p 'Const
    -> Fulfiller (R.Raw r 'Const)
    -> IO ()

-- | Type alias for a handler for an untyped RPC method.
type UntypedMethodHandler = MethodHandler B.AnyStruct B.AnyStruct

-- | Base class for things that can act as capnproto servers.
class SomeServer a where
    -- | Called when the last live reference to a server is dropped.
    shutdown :: a -> IO ()
    shutdown a
_ = () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    -- | Try to extract a value of a given type. The default implementation
    -- always fails (returns 'Nothing'). If an instance chooses to implement
    -- this, it will be possible to use "reflection" on clients that point
    -- at local servers to dynamically unwrap the server value. A typical
    -- implementation will just call Typeable's @cast@ method, but this
    -- needn't be the case -- a server may wish to allow local peers to
    -- unwrap some value that is not exactly the data the server has access
    -- to.
    unwrap :: Typeable b => a -> Maybe b
    unwrap a
_ = Maybe b
forall a. Maybe a
Nothing

-- | Generated interface types have instances of 'Export', which allows a server
-- for that interface to be exported as a 'Client'.
class (R.IsCap i, C.HasTypeId i) => Export i where
    -- | The constraint needed for a server to implement an interface;
    -- if @'Server' i s@ is satisfied, @s@ is a server for interface @i@.
    -- The code generator generates a type class for each interface, and
    -- this will aways be an alias for that type class.
    type Server i :: * -> Constraint

    -- | Convert the server to a 'MethodHandlerTree' populated with appropriate
    -- 'MethodHandler's for the interface. This is really only exported for use
    -- by generated code; users of the library will generally prefer to use
    -- 'export'.
    methodHandlerTree :: Server i s => Proxy i -> s -> MethodHandlerTree
    -- NB: the proxy helps disambiguate types; for some reason TypeApplications
    -- doesn't seem to be enough in the face of a type alias of kind 'Constraint'.
    -- the inconsistency is a bit ugly, but this method isn't intended to called
    -- by users directly, only by generated code and our helper in this module,
    -- so it's less of a big deal.

-- | Lazily computed tree of the method handlers exposed by an interface. Only
-- of interest to generated code.
data MethodHandlerTree = MethodHandlerTree
    { MethodHandlerTree -> Word64
mhtId       :: Word64
    -- ^ type id for the primary interface
    , MethodHandlerTree -> [UntypedMethodHandler]
mhtHandlers :: [UntypedMethodHandler]
    -- ^ method handlers for methods of the primary interface.
    , MethodHandlerTree -> [MethodHandlerTree]
mhtParents  :: [MethodHandlerTree]
    -- ^ Trees for parent interfaces. In the case of diamond dependencies,
    -- there may be duplicates, which are eliminated by 'mhtToCallHandler'.
    }

mhtToCallHandler :: MethodHandlerTree -> CallHandler
mhtToCallHandler :: MethodHandlerTree -> CallHandler
mhtToCallHandler = CallHandler -> [MethodHandlerTree] -> CallHandler
go CallHandler
forall k a. Map k a
M.empty ([MethodHandlerTree] -> CallHandler)
-> (MethodHandlerTree -> [MethodHandlerTree])
-> MethodHandlerTree
-> CallHandler
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MethodHandlerTree -> [MethodHandlerTree]
forall (f :: * -> *) a. Applicative f => a -> f a
pure where
    go :: CallHandler -> [MethodHandlerTree] -> CallHandler
go CallHandler
accum [] = CallHandler
accum
    go CallHandler
accum (MethodHandlerTree
t : [MethodHandlerTree]
ts)
        | MethodHandlerTree -> Word64
mhtId MethodHandlerTree
t Word64 -> CallHandler -> Bool
forall k a. Ord k => k -> Map k a -> Bool
`M.member` CallHandler
accum = CallHandler -> [MethodHandlerTree] -> CallHandler
go CallHandler
accum [MethodHandlerTree]
ts -- dedup diamond dependencies
        | Bool
otherwise =
            CallHandler -> [MethodHandlerTree] -> CallHandler
go (Word64 -> Vector UntypedMethodHandler -> CallHandler -> CallHandler
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (MethodHandlerTree -> Word64
mhtId MethodHandlerTree
t) ([UntypedMethodHandler] -> Vector UntypedMethodHandler
forall a. [a] -> Vector a
V.fromList (MethodHandlerTree -> [UntypedMethodHandler]
mhtHandlers MethodHandlerTree
t)) CallHandler
accum) (MethodHandlerTree -> [MethodHandlerTree]
mhtParents MethodHandlerTree
t [MethodHandlerTree] -> [MethodHandlerTree] -> [MethodHandlerTree]
forall a. [a] -> [a] -> [a]
++ [MethodHandlerTree]
ts)

-- | Export the server as a client for interface @i@. Spawns a server thread
-- with its lifetime bound to the supervisor.
export :: forall i s m. (MonadSTM m, Export i, Server i s, SomeServer s) => Supervisor -> s -> m (Client i)
export :: Supervisor -> s -> m (Client i)
export Supervisor
sup s
srv =
    let h :: CallHandler
h = MethodHandlerTree -> CallHandler
mhtToCallHandler (Proxy i -> s -> MethodHandlerTree
forall i s.
(Export i, Server i s) =>
Proxy i -> s -> MethodHandlerTree
methodHandlerTree (Proxy i
forall k (t :: k). Proxy t
Proxy @i) s
srv) in
    STM (Client i) -> m (Client i)
forall (m :: * -> *) a. MonadSTM m => STM a -> m a
liftSTM (STM (Client i) -> m (Client i)) -> STM (Client i) -> m (Client i)
forall a b. (a -> b) -> a -> b
$ Client -> Client i
forall a. Client -> Client a
Client (Client -> Client i) -> STM Client -> STM (Client i)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Supervisor -> ServerOps IO -> STM Client
forall (m :: * -> *).
MonadSTM m =>
Supervisor -> ServerOps IO -> m Client
URpc.export Supervisor
sup (s -> CallHandler -> ServerOps IO
forall a. SomeServer a => a -> CallHandler -> ServerOps IO
toLegacyServerOps s
srv CallHandler
h)

-- | Look up a particlar 'MethodHandler' in the 'CallHandler'.
findMethod :: Word64 -> Word16 -> CallHandler -> Maybe UntypedMethodHandler
findMethod :: Word64 -> Word16 -> CallHandler -> Maybe UntypedMethodHandler
findMethod Word64
interfaceId Word16
methodId CallHandler
handler = do
    Vector UntypedMethodHandler
iface <- Word64 -> CallHandler -> Maybe (Vector UntypedMethodHandler)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Word64
interfaceId CallHandler
handler
    Vector UntypedMethodHandler
iface Vector UntypedMethodHandler -> Int -> Maybe UntypedMethodHandler
forall a. Vector a -> Int -> Maybe a
V.!? Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
methodId

toLegacyCallHandler
    :: CallHandler
    -> Word64
    -> Word16
    -> Legacy.MethodHandler IO (Maybe (U.Ptr 'Const)) (Maybe (U.Ptr 'Const))
toLegacyCallHandler :: CallHandler
-> Word64
-> Word16
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
toLegacyCallHandler CallHandler
callHandler Word64
interfaceId Word16
methodId =
    Word64 -> Word16 -> CallHandler -> Maybe UntypedMethodHandler
findMethod Word64
interfaceId Word16
methodId CallHandler
callHandler
    Maybe UntypedMethodHandler
-> (Maybe UntypedMethodHandler -> UntypedMethodHandler)
-> UntypedMethodHandler
forall a b. a -> (a -> b) -> b
& UntypedMethodHandler
-> Maybe UntypedMethodHandler -> UntypedMethodHandler
forall a. a -> Maybe a -> a
fromMaybe UntypedMethodHandler
forall p r. MethodHandler p r
methodUnimplemented
    UntypedMethodHandler
-> (UntypedMethodHandler
    -> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const)))
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
forall a b. a -> (a -> b) -> b
& UntypedMethodHandler
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
toLegacyMethodHandler


-- | Convert a typed method handler to an untyped one. Mostly intended for
-- use by generated code.
toUntypedMethodHandler
    :: forall p r. (R.IsStruct p, R.IsStruct r)
    => MethodHandler p r
    -> UntypedMethodHandler
toUntypedMethodHandler :: MethodHandler p r -> UntypedMethodHandler
toUntypedMethodHandler = MethodHandler p r -> UntypedMethodHandler
coerce

toLegacyMethodHandler :: UntypedMethodHandler -> Legacy.MethodHandler IO (Maybe (U.Ptr 'Const)) (Maybe (U.Ptr 'Const))
toLegacyMethodHandler :: UntypedMethodHandler
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
toLegacyMethodHandler UntypedMethodHandler
handler =
    (Maybe (Ptr 'Const) -> Fulfiller (Maybe (Ptr 'Const)) -> IO ())
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
forall (m :: * -> *).
(Maybe (Ptr 'Const) -> Fulfiller (Maybe (Ptr 'Const)) -> m ())
-> MethodHandler m (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
Legacy.untypedHandler ((Maybe (Ptr 'Const) -> Fulfiller (Maybe (Ptr 'Const)) -> IO ())
 -> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const)))
-> (Maybe (Ptr 'Const) -> Fulfiller (Maybe (Ptr 'Const)) -> IO ())
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
forall a b. (a -> b) -> a -> b
$ \Maybe (Ptr 'Const)
args Fulfiller (Maybe (Ptr 'Const))
respond -> do
        Fulfiller (Raw AnyStruct 'Const)
respond' <- (Either (Parsed Exception) (Raw AnyStruct 'Const) -> STM ())
-> IO (Fulfiller (Raw AnyStruct 'Const))
forall (m :: * -> *) a.
MonadSTM m =>
(Either (Parsed Exception) a -> STM ()) -> m (Fulfiller a)
newCallback ((Either (Parsed Exception) (Raw AnyStruct 'Const) -> STM ())
 -> IO (Fulfiller (Raw AnyStruct 'Const)))
-> (Either (Parsed Exception) (Raw AnyStruct 'Const) -> STM ())
-> IO (Fulfiller (Raw AnyStruct 'Const))
forall a b. (a -> b) -> a -> b
$ \case
            Left Parsed Exception
e ->
                Fulfiller (Maybe (Ptr 'Const)) -> Parsed Exception -> STM ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Parsed Exception -> m ()
breakPromise Fulfiller (Maybe (Ptr 'Const))
respond Parsed Exception
e
            Right (R.Raw Unwrapped (Untyped (ReprFor AnyStruct) 'Const)
s) ->
                Fulfiller (Maybe (Ptr 'Const)) -> Maybe (Ptr 'Const) -> STM ()
forall (m :: * -> *) a. MonadSTM m => Fulfiller a -> a -> m ()
fulfill Fulfiller (Maybe (Ptr 'Const))
respond (Ptr 'Const -> Maybe (Ptr 'Const)
forall a. a -> Maybe a
Just (Struct 'Const -> Ptr 'Const
forall (mut :: Mutability). Struct mut -> Ptr mut
U.PtrStruct Unwrapped (Untyped (ReprFor AnyStruct) 'Const)
Struct 'Const
s))
        case Maybe (Ptr 'Const)
args of
            Just (U.PtrStruct Struct 'Const
argStruct) ->
                UntypedMethodHandler
handler (Unwrapped (Untyped (ReprFor AnyStruct) 'Const)
-> Raw AnyStruct 'Const
forall a (mut :: Mutability).
Unwrapped (Untyped (ReprFor a) mut) -> Raw a mut
R.Raw Unwrapped (Untyped (ReprFor AnyStruct) 'Const)
Struct 'Const
argStruct) Fulfiller (Raw AnyStruct 'Const)
respond'
            Maybe (Ptr 'Const)
_ ->
                Fulfiller (Maybe (Ptr 'Const)) -> Parsed Exception -> IO ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Parsed Exception -> m ()
breakPromise Fulfiller (Maybe (Ptr 'Const))
respond (Parsed Exception -> IO ()) -> Parsed Exception -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> Parsed Exception
eFailed Text
"Argument was not a struct"

toLegacyServerOps :: SomeServer a => a -> CallHandler -> Legacy.ServerOps IO
toLegacyServerOps :: a -> CallHandler -> ServerOps IO
toLegacyServerOps a
srv CallHandler
callHandler = ServerOps :: forall (m :: * -> *).
(Word64
 -> Word16
 -> MethodHandler m (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const)))
-> m () -> (forall a. Typeable a => Maybe a) -> ServerOps m
Legacy.ServerOps
    { handleStop :: IO ()
handleStop = a -> IO ()
forall a. SomeServer a => a -> IO ()
shutdown a
srv
    , handleCast :: forall a. Typeable a => Maybe a
handleCast = a -> Maybe a
forall a b. (SomeServer a, Typeable b) => a -> Maybe b
unwrap a
srv
    , handleCall :: Word64
-> Word16
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
handleCall = CallHandler
-> Word64
-> Word16
-> MethodHandler IO (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
toLegacyCallHandler CallHandler
callHandler
    }

-- Helpers for writing method handlers

-- | Handle a method, working with the parsed form of parameters and
-- results.
handleParsed ::
    ( C.Parse p pp, R.IsStruct p
    , C.Parse r rr, R.IsStruct r
    ) => (pp -> IO rr) -> MethodHandler p r
handleParsed :: (pp -> IO rr) -> MethodHandler p r
handleParsed pp -> IO rr
handler Raw p 'Const
param = (Fulfiller (Raw r 'Const) -> IO ())
-> Fulfiller (Raw r 'Const) -> IO ()
forall a b. (Fulfiller a -> IO b) -> Fulfiller a -> IO b
propagateExceptions ((Fulfiller (Raw r 'Const) -> IO ())
 -> Fulfiller (Raw r 'Const) -> IO ())
-> (Fulfiller (Raw r 'Const) -> IO ())
-> Fulfiller (Raw r 'Const)
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Fulfiller (Raw r 'Const)
f -> do
    pp
p <- WordCount -> LimitT IO pp -> IO pp
forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
defaultLimit (LimitT IO pp -> IO pp) -> LimitT IO pp -> IO pp
forall a b. (a -> b) -> a -> b
$ Raw p 'Const -> LimitT IO pp
forall t p (m :: * -> *).
(Parse t p, ReadCtx m 'Const) =>
Raw t 'Const -> m p
C.parse Raw p 'Const
param
    rr
r <- pp -> IO rr
handler pp
p
    -- TODO: Figure out how to add an instance of Thaw for
    -- Raw so we can skip the (un)wrapping here.
    Struct 'Const
struct <- WordCount
-> (forall s. PureBuilder s (Struct ('Mut s)))
-> IO (Struct 'Const)
forall (m :: * -> *) (f :: Mutability -> *).
(MonadThrow m, MaybeMutable f) =>
WordCount -> (forall s. PureBuilder s (f ('Mut s))) -> m (f 'Const)
createPure WordCount
forall a. Bounded a => a
maxBound ((forall s. PureBuilder s (Struct ('Mut s))) -> IO (Struct 'Const))
-> (forall s. PureBuilder s (Struct ('Mut s)))
-> IO (Struct 'Const)
forall a b. (a -> b) -> a -> b
$ Raw r ('Mut s) -> Struct ('Mut s)
forall a (mut :: Mutability).
Raw a mut -> Unwrapped (Untyped (ReprFor a) mut)
R.fromRaw (Raw r ('Mut s) -> Struct ('Mut s))
-> PureBuilder s (Raw r ('Mut s))
-> PureBuilder s (Struct ('Mut s))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> rr -> PureBuilder s (Raw r ('Mut s))
forall a (m :: * -> *) pa s.
(RWCtx m s, IsStruct a, Parse a pa) =>
pa -> m (Raw a ('Mut s))
parsedToRaw rr
r
    Fulfiller (Raw r 'Const) -> Raw r 'Const -> IO ()
forall (m :: * -> *) a. MonadSTM m => Fulfiller a -> a -> m ()
fulfill Fulfiller (Raw r 'Const)
f (Unwrapped (Untyped (ReprFor r) 'Const) -> Raw r 'Const
forall a (mut :: Mutability).
Unwrapped (Untyped (ReprFor a) mut) -> Raw a mut
R.Raw Unwrapped (Untyped (ReprFor r) 'Const)
Struct 'Const
struct)

-- | Handle a method, working with the raw (unparsed) form of
-- parameters and results.
handleRaw
    :: (R.IsStruct p, R.IsStruct r)
    => (R.Raw p 'Const -> IO (R.Raw r 'Const)) -> MethodHandler p r
handleRaw :: (Raw p 'Const -> IO (Raw r 'Const)) -> MethodHandler p r
handleRaw Raw p 'Const -> IO (Raw r 'Const)
handler Raw p 'Const
param = (Fulfiller (Raw r 'Const) -> IO ())
-> Fulfiller (Raw r 'Const) -> IO ()
forall a b. (Fulfiller a -> IO b) -> Fulfiller a -> IO b
propagateExceptions ((Fulfiller (Raw r 'Const) -> IO ())
 -> Fulfiller (Raw r 'Const) -> IO ())
-> (Fulfiller (Raw r 'Const) -> IO ())
-> Fulfiller (Raw r 'Const)
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Fulfiller (Raw r 'Const)
f ->
    Raw p 'Const -> IO (Raw r 'Const)
handler Raw p 'Const
param IO (Raw r 'Const) -> (Raw r 'Const -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Fulfiller (Raw r 'Const) -> Raw r 'Const -> IO ()
forall (m :: * -> *) a. MonadSTM m => Fulfiller a -> a -> m ()
fulfill Fulfiller (Raw r 'Const)
f

-- Helper for handle*; breaks the promise if the handler throws.
propagateExceptions :: (Fulfiller a -> IO b) -> Fulfiller a -> IO b
propagateExceptions :: (Fulfiller a -> IO b) -> Fulfiller a -> IO b
propagateExceptions Fulfiller a -> IO b
h Fulfiller a
f =
    Fulfiller a -> IO b
h Fulfiller a
f IO b -> (SomeException -> IO ()) -> IO b
forall (m :: * -> *) e a b.
(MonadMask m, Exception e) =>
m a -> (e -> m b) -> m a
`withException` (Fulfiller a -> Parsed Exception -> IO ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Parsed Exception -> m ()
breakPromise Fulfiller a
f (Parsed Exception -> IO ())
-> (SomeException -> Parsed Exception) -> SomeException -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> SomeException -> Parsed Exception
wrapException Bool
False)

-- | 'MethodHandler' that always throws unimplemented.
methodUnimplemented :: MethodHandler p r
methodUnimplemented :: MethodHandler p r
methodUnimplemented Raw p 'Const
_ Fulfiller (Raw r 'Const)
f = Fulfiller (Raw r 'Const) -> Parsed Exception -> IO ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Parsed Exception -> m ()
breakPromise Fulfiller (Raw r 'Const)
f Parsed Exception
eMethodUnimplemented

{-
Sketch of future Async API, might take a bit of internals work to make
this possible:

-- | Handle a method call asynchronously.
--
-- When invoked, the handleer will be run synchronously, blocking further
-- method calls until the 'IO' returns. The method call does not return
-- until the 'Async' resolves, but further method calls can be serviced in
-- the meantime.
--
-- If a Finish message is received before the Async resolves, it will be
-- 'cancel'ed.
handleRawAsync
    :: (R.IsStruct p, R.IsStruct r)
    => (R.Raw 'Const p -> IO (Async (R.Raw 'Const r)))
    -> MethodHandler IO p r

-- | Like 'handleRawAsync', but accepts and returns parsed values.
handleParsedAsync  ::
    ( C.Parse p pp, R.IsStruct p
    , C.Parse r, rr, R.IsStruct r
    )
    => (pp -> IO (Async rr))
    -> MethodHandler IO p r
-}