{- |
This module provides utilities to run an LSP `Session` in `IO`.
-}
module Language.LSP.Client where

import Control.Concurrent.STM
import Control.Monad (forever)
import Control.Monad.Reader (asks, runReaderT)
import Data.ByteString.Lazy qualified as LazyByteString
import Data.Dependent.Map qualified as DMap
import Data.Either (fromLeft)
import Data.Generics.Labels ()
import Language.LSP.Client.Decoding
import Language.LSP.Client.Encoding (encode)
import Language.LSP.Client.Session
import Language.LSP.Protocol.Message qualified as LSP
import Language.LSP.VFS (emptyVFS)
import System.IO (Handle)
import UnliftIO (MonadUnliftIO, concurrently_, liftIO, race)
import Prelude

{- | Starts a new session, using the specified handles to communicate with the
server.
-}
runSessionWithHandles
    :: (MonadUnliftIO io)
    => Handle
    -- ^ The input handle: messages sent from the server to the client will be read from here
    -> Handle
    -- ^ The output handle: messages sent by the client will be written here
    -> SessionT io a
    -- ^ Session actions
    -> io a
runSessionWithHandles :: forall (io :: * -> *) a.
MonadUnliftIO io =>
Handle -> Handle -> SessionT io a -> io a
runSessionWithHandles Handle
input Handle
output SessionT io a
action = do
    SessionState
initialState <- VFS -> io SessionState
forall (io :: * -> *). MonadIO io => VFS -> io SessionState
defaultSessionState VFS
emptyVFS
    (SessionT io a -> SessionState -> io a)
-> SessionState -> SessionT io a -> io a
forall a b c. (a -> b -> c) -> b -> a -> c
flip SessionT io a -> SessionState -> io a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT SessionState
initialState (SessionT io a -> io a) -> SessionT io a -> io a
forall a b. (a -> b) -> a -> b
$ do
        Either a ()
actionResult <- SessionT io a
-> ReaderT SessionState io ()
-> ReaderT SessionState io (Either a ())
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> m b -> m (Either a b)
race SessionT io a
action (ReaderT SessionState io ()
 -> ReaderT SessionState io (Either a ()))
-> ReaderT SessionState io ()
-> ReaderT SessionState io (Either a ())
forall a b. (a -> b) -> a -> b
$ do
            let send :: ReaderT SessionState io ()
send = do
                    FromClientMessage
message <- (SessionState -> TQueue FromClientMessage)
-> ReaderT SessionState io (TQueue FromClientMessage)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SessionState -> TQueue FromClientMessage
outgoing ReaderT SessionState io (TQueue FromClientMessage)
-> (TQueue FromClientMessage
    -> ReaderT SessionState io FromClientMessage)
-> ReaderT SessionState io FromClientMessage
forall a b.
ReaderT SessionState io a
-> (a -> ReaderT SessionState io b) -> ReaderT SessionState io b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO FromClientMessage -> ReaderT SessionState io FromClientMessage
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO FromClientMessage -> ReaderT SessionState io FromClientMessage)
-> (TQueue FromClientMessage -> IO FromClientMessage)
-> TQueue FromClientMessage
-> ReaderT SessionState io FromClientMessage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM FromClientMessage -> IO FromClientMessage
forall a. STM a -> IO a
atomically (STM FromClientMessage -> IO FromClientMessage)
-> (TQueue FromClientMessage -> STM FromClientMessage)
-> TQueue FromClientMessage
-> IO FromClientMessage
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TQueue FromClientMessage -> STM FromClientMessage
forall a. TQueue a -> STM a
readTQueue
                    IO () -> ReaderT SessionState io ()
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT SessionState io ())
-> IO () -> ReaderT SessionState io ()
forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
LazyByteString.hPut Handle
output (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ FromClientMessage -> ByteString
forall a. ToJSON a => a -> ByteString
encode FromClientMessage
message
            let receive :: ReaderT SessionState io ()
receive = do
                    ByteString
serverBytes <- IO ByteString -> ReaderT SessionState io ByteString
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> ReaderT SessionState io ByteString)
-> IO ByteString -> ReaderT SessionState io ByteString
forall a b. (a -> b) -> a -> b
$ Handle -> IO ByteString
getNextMessage Handle
input
                    (FromServerMessage
serverMessage, IO ()
requestCallback) <-
                        (SessionState -> TVar RequestMap)
-> ReaderT SessionState io (TVar RequestMap)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SessionState -> TVar RequestMap
pendingRequests
                            ReaderT SessionState io (TVar RequestMap)
-> (TVar RequestMap
    -> ReaderT SessionState io (FromServerMessage, IO ()))
-> ReaderT SessionState io (FromServerMessage, IO ())
forall a b.
ReaderT SessionState io a
-> (a -> ReaderT SessionState io b) -> ReaderT SessionState io b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO (FromServerMessage, IO ())
-> ReaderT SessionState io (FromServerMessage, IO ())
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO
                                (IO (FromServerMessage, IO ())
 -> ReaderT SessionState io (FromServerMessage, IO ()))
-> (TVar RequestMap -> IO (FromServerMessage, IO ()))
-> TVar RequestMap
-> ReaderT SessionState io (FromServerMessage, IO ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM (FromServerMessage, IO ()) -> IO (FromServerMessage, IO ())
forall a. STM a -> IO a
atomically
                                (STM (FromServerMessage, IO ()) -> IO (FromServerMessage, IO ()))
-> (TVar RequestMap -> STM (FromServerMessage, IO ()))
-> TVar RequestMap
-> IO (FromServerMessage, IO ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TVar RequestMap
 -> (RequestMap -> ((FromServerMessage, IO ()), RequestMap))
 -> STM (FromServerMessage, IO ()))
-> (RequestMap -> ((FromServerMessage, IO ()), RequestMap))
-> TVar RequestMap
-> STM (FromServerMessage, IO ())
forall a b c. (a -> b -> c) -> b -> a -> c
flip TVar RequestMap
-> (RequestMap -> ((FromServerMessage, IO ()), RequestMap))
-> STM (FromServerMessage, IO ())
forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar (ByteString
-> RequestMap -> ((FromServerMessage, IO ()), RequestMap)
decodeFromServerMsg ByteString
serverBytes)
                    FromServerMessage -> ReaderT SessionState io ()
forall (m :: * -> *). MonadSession m => FromServerMessage -> m ()
handleServerMessage FromServerMessage
serverMessage
                    IO () -> ReaderT SessionState io ()
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
requestCallback
                    case FromServerMessage
serverMessage of
                        LSP.FromServerMess SMethod m
smethod TMessage m
msg -> case SMethod m -> ServerNotOrReq m
forall {t :: MessageKind} (m :: Method 'ServerToClient t).
SServerMethod m -> ServerNotOrReq m
LSP.splitServerMethod SMethod m
smethod of
                            ServerNotOrReq m
LSP.IsServerNot -> do
                                NotificationMap
handlers :: NotificationMap <- (SessionState -> TVar NotificationMap)
-> ReaderT SessionState io (TVar NotificationMap)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SessionState -> TVar NotificationMap
notificationHandlers ReaderT SessionState io (TVar NotificationMap)
-> (TVar NotificationMap
    -> ReaderT SessionState io NotificationMap)
-> ReaderT SessionState io NotificationMap
forall a b.
ReaderT SessionState io a
-> (a -> ReaderT SessionState io b) -> ReaderT SessionState io b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO NotificationMap -> ReaderT SessionState io NotificationMap
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO NotificationMap -> ReaderT SessionState io NotificationMap)
-> (TVar NotificationMap -> IO NotificationMap)
-> TVar NotificationMap
-> ReaderT SessionState io NotificationMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar NotificationMap -> IO NotificationMap
forall a. TVar a -> IO a
readTVarIO
                                let NotificationCallback TNotificationMessage m -> IO ()
cb = NotificationCallback m
-> SMethod m -> NotificationMap -> NotificationCallback m
forall {k1} (k2 :: k1 -> *) (f :: k1 -> *) (v :: k1).
GCompare k2 =>
f v -> k2 v -> DMap k2 f -> f v
DMap.findWithDefault ((TNotificationMessage m -> IO ()) -> NotificationCallback m
forall (m :: Method 'ServerToClient 'Notification).
(TNotificationMessage m -> IO ()) -> NotificationCallback m
NotificationCallback (IO () -> TNotificationMessage m -> IO ()
forall a b. a -> b -> a
const (IO () -> TNotificationMessage m -> IO ())
-> IO () -> TNotificationMessage m -> IO ()
forall a b. (a -> b) -> a -> b
$ () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())) SMethod m
SMethod m
smethod NotificationMap
handlers
                                IO () -> ReaderT SessionState io ()
forall a. IO a -> ReaderT SessionState io a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT SessionState io ())
-> IO () -> ReaderT SessionState io ()
forall a b. (a -> b) -> a -> b
$ TNotificationMessage m -> IO ()
cb TMessage m
TNotificationMessage m
msg
                            ServerNotOrReq m
_ -> () -> ReaderT SessionState io ()
forall a. a -> ReaderT SessionState io a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                        FromServerMessage
_ -> () -> ReaderT SessionState io ()
forall a. a -> ReaderT SessionState io a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            ReaderT SessionState io Any
-> ReaderT SessionState io Any -> ReaderT SessionState io ()
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m ()
concurrently_ (ReaderT SessionState io () -> ReaderT SessionState io Any
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever ReaderT SessionState io ()
send) (ReaderT SessionState io () -> ReaderT SessionState io Any
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever ReaderT SessionState io ()
receive)
        pure $ a -> Either a () -> a
forall a b. a -> Either a b -> a
fromLeft ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"runSessionWithHandle: send/receive thread should not exit") Either a ()
actionResult