{-# 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 qualified Capnp.Basics as B
import qualified Capnp.Classes as C
import Capnp.Convert (parsedToRaw)
import Capnp.Message (Mutability (..))
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, Type)
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
_ = 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
_ = 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 :: Type -> 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
  { -- | type id for the primary interface
    MethodHandlerTree -> Word64
mhtId :: Word64,
    -- | method handlers for methods of the primary interface.
    MethodHandlerTree -> [UntypedMethodHandler]
mhtHandlers :: [UntypedMethodHandler],
    -- | Trees for parent interfaces. In the case of diamond dependencies,
    -- there may be duplicates, which are eliminated by 'mhtToCallHandler'.
    MethodHandlerTree -> [MethodHandlerTree]
mhtParents :: [MethodHandlerTree]
  }

mhtToCallHandler :: MethodHandlerTree -> CallHandler
mhtToCallHandler :: MethodHandlerTree -> CallHandler
mhtToCallHandler = CallHandler -> [MethodHandlerTree] -> CallHandler
go forall k a. Map k a
M.empty forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 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 (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (MethodHandlerTree -> Word64
mhtId MethodHandlerTree
t) (forall a. [a] -> Vector a
V.fromList (MethodHandlerTree -> [UntypedMethodHandler]
mhtHandlers MethodHandlerTree
t)) CallHandler
accum) (MethodHandlerTree -> [MethodHandlerTree]
mhtParents MethodHandlerTree
t 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 :: forall i s (m :: * -> *).
(MonadSTM m, Export i, Server i s, SomeServer s) =>
Supervisor -> s -> m (Client i)
export Supervisor
sup s
srv =
  let h :: CallHandler
h = MethodHandlerTree -> CallHandler
mhtToCallHandler (forall i s.
(Export i, Server i s) =>
Proxy i -> s -> MethodHandlerTree
methodHandlerTree (forall {k} (t :: k). Proxy t
Proxy @i) s
srv)
   in forall (m :: * -> *) a. MonadSTM m => STM a -> m a
liftSTM forall a b. (a -> b) -> a -> b
$ forall a. Client -> Client a
Client forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadSTM m =>
Supervisor -> ServerOps -> m Client
URpc.export Supervisor
sup (forall a. SomeServer a => a -> CallHandler -> ServerOps
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 <- forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Word64
interfaceId CallHandler
handler
  Vector UntypedMethodHandler
iface forall a. Vector a -> Int -> Maybe a
V.!? forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
methodId

toLegacyCallHandler ::
  CallHandler ->
  Word64 ->
  Word16 ->
  Legacy.MethodHandler (Maybe (U.Ptr 'Const)) (Maybe (U.Ptr 'Const))
toLegacyCallHandler :: CallHandler
-> Word64
-> Word16
-> MethodHandler (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
    forall a b. a -> (a -> b) -> b
& forall a. a -> Maybe a -> a
fromMaybe forall p r. MethodHandler p r
methodUnimplemented
    forall a b. a -> (a -> b) -> b
& UntypedMethodHandler
-> MethodHandler (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 :: forall p r.
(IsStruct p, IsStruct r) =>
MethodHandler p r -> UntypedMethodHandler
toUntypedMethodHandler = coerce :: forall a b. Coercible a b => a -> b
coerce

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

toLegacyServerOps :: SomeServer a => a -> CallHandler -> Legacy.ServerOps
toLegacyServerOps :: forall a. SomeServer a => a -> CallHandler -> ServerOps
toLegacyServerOps a
srv CallHandler
callHandler =
  Legacy.ServerOps
    { handleStop :: IO ()
handleStop = forall a. SomeServer a => a -> IO ()
shutdown a
srv,
      handleCast :: forall a. Typeable a => Maybe a
handleCast = forall a b. (SomeServer a, Typeable b) => a -> Maybe b
unwrap a
srv,
      handleCall :: Word64
-> Word16
-> MethodHandler (Maybe (Ptr 'Const)) (Maybe (Ptr 'Const))
handleCall = CallHandler
-> Word64
-> Word16
-> MethodHandler (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 pr,
    R.IsStruct r
  ) =>
  (pp -> IO pr) ->
  MethodHandler p r
handleParsed :: forall p pp r pr.
(Parse p pp, IsStruct p, Parse r pr, IsStruct r) =>
(pp -> IO pr) -> MethodHandler p r
handleParsed pp -> IO pr
handler Raw p 'Const
param = forall a b. (Fulfiller a -> IO b) -> Fulfiller a -> IO b
propagateExceptions forall a b. (a -> b) -> a -> b
$ \Fulfiller (Raw r 'Const)
f -> do
  pp
p <- forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
defaultLimit forall a b. (a -> b) -> a -> b
$ forall t p (m :: * -> *).
(Parse t p, ReadCtx m 'Const) =>
Raw t 'Const -> m p
C.parse Raw p 'Const
param
  pr
r <- pp -> IO pr
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 <- forall (m :: * -> *) (f :: Mutability -> *).
(MonadThrow m, MaybeMutable f) =>
WordCount -> (forall s. PureBuilder s (f ('Mut s))) -> m (f 'Const)
createPure forall a. Bounded a => a
maxBound forall a b. (a -> b) -> a -> b
$ forall a (mut :: Mutability).
Raw a mut -> Unwrapped (Untyped (ReprFor a) mut)
R.fromRaw forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (m :: * -> *) pa s.
(RWCtx m s, IsStruct a, Parse a pa) =>
pa -> m (Raw a ('Mut s))
parsedToRaw pr
r
  forall (m :: * -> *) a. MonadSTM m => Fulfiller a -> a -> m ()
fulfill Fulfiller (Raw r 'Const)
f (forall a (mut :: Mutability).
Unwrapped (Untyped (ReprFor a) mut) -> Raw a mut
R.Raw 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 :: forall p r.
(IsStruct p, IsStruct r) =>
(Raw p 'Const -> IO (Raw r 'Const)) -> MethodHandler p r
handleRaw Raw p 'Const -> IO (Raw r 'Const)
handler Raw p 'Const
param = forall a b. (Fulfiller a -> IO b) -> Fulfiller a -> IO b
propagateExceptions forall a b. (a -> b) -> a -> b
$ \Fulfiller (Raw r 'Const)
f ->
  Raw p 'Const -> IO (Raw r 'Const)
handler Raw p 'Const
param forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 :: forall a b. (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 forall (m :: * -> *) e a b.
(MonadMask m, Exception e) =>
m a -> (e -> m b) -> m a
`withException` (forall (m :: * -> *) a.
MonadSTM m =>
Fulfiller a -> Parsed Exception -> m ()
breakPromise Fulfiller a
f 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 :: forall p r. MethodHandler p r
methodUnimplemented Raw p 'Const
_ Fulfiller (Raw r 'Const)
f = 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
-}