{-# 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)
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 'Const p
    -> Fulfiller (R.Raw 'Const r)
    -> 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 'Const AnyStruct)
respond' <- (Either Exception (Raw 'Const AnyStruct) -> STM ())
-> IO (Fulfiller (Raw 'Const AnyStruct))
forall (m :: * -> *) a.
MonadSTM m =>
(Either Exception a -> STM ()) -> m (Fulfiller a)
newCallback ((Either Exception (Raw 'Const AnyStruct) -> STM ())
 -> IO (Fulfiller (Raw 'Const AnyStruct)))
-> (Either Exception (Raw 'Const AnyStruct) -> STM ())
-> IO (Fulfiller (Raw 'Const AnyStruct))
forall a b. (a -> b) -> a -> b
$ \case
            Left Exception
e ->
                Fulfiller (Maybe (Ptr 'Const)) -> Exception -> STM ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Exception -> m ()
breakPromise Fulfiller (Maybe (Ptr 'Const))
respond Exception
e
            Right (R.Raw Untyped 'Const (ReprFor AnyStruct)
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 Struct 'Const
Untyped 'Const (ReprFor AnyStruct)
s))
        case Maybe (Ptr 'Const)
args of
            Just (U.PtrStruct Struct 'Const
argStruct) ->
                UntypedMethodHandler
handler (Untyped 'Const (ReprFor AnyStruct) -> Raw 'Const AnyStruct
forall (mut :: Mutability) a. Untyped mut (ReprFor a) -> Raw mut a
R.Raw Struct 'Const
Untyped 'Const (ReprFor AnyStruct)
argStruct) Fulfiller (Raw 'Const AnyStruct)
respond'
            Maybe (Ptr 'Const)
_ ->
                Fulfiller (Maybe (Ptr 'Const)) -> Exception -> IO ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Exception -> m ()
breakPromise Fulfiller (Maybe (Ptr 'Const))
respond (Exception -> IO ()) -> Exception -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> 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 'Const p -> IO (Raw 'Const r)) -> MethodHandler p r
forall p r.
(IsStruct p, IsStruct r) =>
(Raw 'Const p -> IO (Raw 'Const r)) -> MethodHandler p r
handleRaw ((Raw 'Const p -> IO (Raw 'Const r)) -> MethodHandler p r)
-> (Raw 'Const p -> IO (Raw 'Const r)) -> MethodHandler p r
forall a b. (a -> b) -> a -> b
$ \Raw 'Const p
param -> 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 'Const p -> LimitT IO pp
forall t p (m :: * -> *).
(Parse t p, ReadCtx m 'Const) =>
Raw 'Const t -> m p
C.parse Raw 'Const p
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 (Mutable s (Struct 'Const)))
-> IO (Struct 'Const)
forall (m :: * -> *) a.
(MonadThrow m, Thaw a) =>
WordCount -> (forall s. PureBuilder s (Mutable s a)) -> m a
createPure WordCount
forall a. Bounded a => a
maxBound ((forall s. PureBuilder s (Mutable s (Struct 'Const)))
 -> IO (Struct 'Const))
-> (forall s. PureBuilder s (Mutable s (Struct 'Const)))
-> IO (Struct 'Const)
forall a b. (a -> b) -> a -> b
$ Raw ('Mut s) r -> Struct ('Mut s)
forall (mut :: Mutability) a. Raw mut a -> Untyped mut (ReprFor a)
R.fromRaw (Raw ('Mut s) r -> Struct ('Mut s))
-> PureBuilder s (Raw ('Mut s) r)
-> PureBuilder s (Struct ('Mut s))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> rr -> PureBuilder s (Raw ('Mut s) r)
forall a (m :: * -> *) pa s.
(RWCtx m s, IsStruct a, Parse a pa) =>
pa -> m (Raw ('Mut s) a)
parsedToRaw rr
r
        Raw 'Const r -> IO (Raw 'Const r)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Untyped 'Const (ReprFor r) -> Raw 'Const r
forall (mut :: Mutability) a. Untyped mut (ReprFor a) -> Raw mut a
R.Raw Struct 'Const
Untyped 'Const (ReprFor r)
struct)

-- | Handle a method, working with the raw (unparsed) form of
-- parameters and results.
handleRaw
    :: (R.IsStruct p, R.IsStruct r)
    => (R.Raw 'Const p -> IO (R.Raw 'Const r)) -> MethodHandler p r
handleRaw :: (Raw 'Const p -> IO (Raw 'Const r)) -> MethodHandler p r
handleRaw Raw 'Const p -> IO (Raw 'Const r)
handler Raw 'Const p
param Fulfiller (Raw 'Const r)
f = do
    Raw 'Const r
res <- Raw 'Const p -> IO (Raw 'Const r)
handler Raw 'Const p
param IO (Raw 'Const r) -> (Exception -> IO ()) -> IO (Raw 'Const r)
forall (m :: * -> *) e a b.
(MonadMask m, Exception e) =>
m a -> (e -> m b) -> m a
`withException` Fulfiller (Raw 'Const r) -> Exception -> IO ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Exception -> m ()
breakPromise Fulfiller (Raw 'Const r)
f
    Fulfiller (Raw 'Const r) -> Raw 'Const r -> IO ()
forall (m :: * -> *) a. MonadSTM m => Fulfiller a -> a -> m ()
fulfill Fulfiller (Raw 'Const r)
f Raw 'Const r
res


-- | 'MethodHandler' that always throws unimplemented.
methodUnimplemented :: MethodHandler p r
methodUnimplemented :: MethodHandler p r
methodUnimplemented Raw 'Const p
_ Fulfiller (Raw 'Const r)
f = Fulfiller (Raw 'Const r) -> Exception -> IO ()
forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Exception -> m ()
breakPromise Fulfiller (Raw 'Const r)
f 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
-}