{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE LambdaCase #-}

-- So we can keep using the old prettyprinter modules (which have a better
-- compatibility range) for now.
{-# OPTIONS_GHC -Wno-deprecations #-}

module Language.LSP.Server.Control
  (
  -- * Running
    runServer
  , runServerWith
  , runServerWithHandles
  , LspServerLog (..)
  ) where

import qualified Colog.Core as L
import           Colog.Core (LogAction (..), WithSeverity (..), Severity (..), (<&))
import           Control.Concurrent
import           Control.Concurrent.STM.TChan
import           Control.Monad
import           Control.Monad.STM
import           Control.Monad.IO.Class
import qualified Data.Aeson as J
import qualified Data.Attoparsec.ByteString as Attoparsec
import Data.Attoparsec.ByteString.Char8
import qualified Data.ByteString as BS
import Data.ByteString.Builder.Extra (defaultChunkSize)
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TL
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import           Data.Text.Prettyprint.Doc
import           Data.List
import           Language.LSP.Server.Core
import qualified Language.LSP.Server.Processing as Processing
import           Language.LSP.Types
import           Language.LSP.VFS
import Language.LSP.Logging (defaultClientLogger)
import           System.IO

data LspServerLog =
  LspProcessingLog Processing.LspProcessingLog
  | DecodeInitializeError String
  | HeaderParseFail [String] String
  | EOF
  | Starting
  | ParsedMsg T.Text
  | SendMsg TL.Text
  deriving (Int -> LspServerLog -> ShowS
[LspServerLog] -> ShowS
LspServerLog -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LspServerLog] -> ShowS
$cshowList :: [LspServerLog] -> ShowS
show :: LspServerLog -> String
$cshow :: LspServerLog -> String
showsPrec :: Int -> LspServerLog -> ShowS
$cshowsPrec :: Int -> LspServerLog -> ShowS
Show)

instance Pretty LspServerLog where
  pretty :: forall ann. LspServerLog -> Doc ann
pretty (LspProcessingLog LspProcessingLog
l) = forall a ann. Pretty a => a -> Doc ann
pretty LspProcessingLog
l
  pretty (DecodeInitializeError String
err) =
    forall ann. [Doc ann] -> Doc ann
vsep [
      Doc ann
"Got error while decoding initialize:"
      , forall a ann. Pretty a => a -> Doc ann
pretty String
err
      ]
  pretty (HeaderParseFail [String]
ctxs String
err) =
    forall ann. [Doc ann] -> Doc ann
vsep [
      Doc ann
"Failed to parse message header:"
      , forall a ann. Pretty a => a -> Doc ann
pretty (forall a. [a] -> [[a]] -> [a]
intercalate String
" > " [String]
ctxs) forall a. Semigroup a => a -> a -> a
<> Doc ann
": " forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a ann. Pretty a => a -> Doc ann
pretty String
err
      ]
  pretty LspServerLog
EOF = Doc ann
"Got EOF"
  pretty LspServerLog
Starting = Doc ann
"Starting server"
  pretty (ParsedMsg Text
msg) = Doc ann
"---> " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Text
msg
  pretty (SendMsg Text
msg) = Doc ann
"<--2-- " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Text
msg

-- ---------------------------------------------------------------------

-- | Convenience function for 'runServerWithHandles' which:
--     (1) reads from stdin;
--     (2) writes to stdout; and
--     (3) logs to stderr and to the client, with some basic filtering.
runServer :: forall config . ServerDefinition config -> IO Int
runServer :: forall config. ServerDefinition config -> IO Int
runServer =
  forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> Handle
-> Handle
-> ServerDefinition config
-> IO Int
runServerWithHandles
  LogAction IO (WithSeverity LspServerLog)
ioLogger
  LogAction (LspM config) (WithSeverity LspServerLog)
lspLogger
  Handle
stdin
  Handle
stdout
  where
    prettyMsg :: WithSeverity a -> Doc ann
prettyMsg WithSeverity a
l = Doc ann
"[" forall a. Semigroup a => a -> a -> a
<> forall a ann. Show a => a -> Doc ann
viaShow (forall msg. WithSeverity msg -> Severity
L.getSeverity WithSeverity a
l) forall a. Semigroup a => a -> a -> a
<> Doc ann
"] " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty (forall msg. WithSeverity msg -> msg
L.getMsg WithSeverity a
l)
    ioLogger :: LogAction IO (WithSeverity LspServerLog)
    ioLogger :: LogAction IO (WithSeverity LspServerLog)
ioLogger = forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a} {ann}. Pretty a => WithSeverity a -> Doc ann
prettyMsg) forall (m :: * -> *). MonadIO m => LogAction m String
L.logStringStderr
    lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
    lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
lspLogger =
      let clientLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
clientLogger = forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> Text
T.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a ann. Pretty a => a -> Doc ann
pretty)) forall c (m :: * -> *).
MonadLsp c m =>
LogAction m (WithSeverity Text)
defaultClientLogger
      in LogAction (LspM config) (WithSeverity LspServerLog)
clientLogger forall a. Semigroup a => a -> a -> a
<> forall (m :: * -> *) (n :: * -> *) a.
(forall x. m x -> n x) -> LogAction m a -> LogAction n a
L.hoistLogAction forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO LogAction IO (WithSeverity LspServerLog)
ioLogger

-- | Starts a language server over the specified handles.
-- This function will return once the @exit@ notification is received.
runServerWithHandles ::
    LogAction IO (WithSeverity LspServerLog)
    -- ^ The logger to use outside the main body of the server where we can't assume the ability to send messages.
    -> LogAction (LspM config) (WithSeverity LspServerLog)
    -- ^ The logger to use once the server has started and can successfully send messages.
    -> Handle
    -- ^ Handle to read client input from.
    -> Handle
    -- ^ Handle to write output to.
    -> ServerDefinition config
    -> IO Int         -- exit code
runServerWithHandles :: forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> Handle
-> Handle
-> ServerDefinition config
-> IO Int
runServerWithHandles LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger Handle
hin Handle
hout ServerDefinition config
serverDefinition = do

  Handle -> BufferMode -> IO ()
hSetBuffering Handle
hin BufferMode
NoBuffering
  Handle -> TextEncoding -> IO ()
hSetEncoding  Handle
hin TextEncoding
utf8

  Handle -> BufferMode -> IO ()
hSetBuffering Handle
hout BufferMode
NoBuffering
  Handle -> TextEncoding -> IO ()
hSetEncoding  Handle
hout TextEncoding
utf8

  let
    clientIn :: IO ByteString
clientIn = Handle -> Int -> IO ByteString
BS.hGetSome Handle
hin Int
defaultChunkSize

    clientOut :: ByteString -> IO ()
clientOut ByteString
out = do
      Handle -> ByteString -> IO ()
BSL.hPut Handle
hout ByteString
out
      Handle -> IO ()
hFlush Handle
hout

  forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> (ByteString -> IO ())
-> ServerDefinition config
-> IO Int
runServerWith LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ByteString -> IO ()
clientOut ServerDefinition config
serverDefinition

-- | Starts listening and sending requests and responses
-- using the specified I/O.
runServerWith ::
    LogAction IO (WithSeverity LspServerLog)
    -- ^ The logger to use outside the main body of the server where we can't assume the ability to send messages.
    -> LogAction (LspM config) (WithSeverity LspServerLog)
    -- ^ The logger to use once the server has started and can successfully send messages.
    -> IO BS.ByteString
    -- ^ Client input.
    -> (BSL.ByteString -> IO ())
    -- ^ Function to provide output to.
    -> ServerDefinition config
    -> IO Int         -- exit code
runServerWith :: forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> (ByteString -> IO ())
-> ServerDefinition config
-> IO Int
runServerWith LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ByteString -> IO ()
clientOut ServerDefinition config
serverDefinition = do

  LogAction IO (WithSeverity LspServerLog)
ioLogger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& LspServerLog
Starting forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Info

  TChan Value
cout <- forall a. STM a -> IO a
atomically forall a. STM (TChan a)
newTChan :: IO (TChan J.Value)
  ThreadId
_rhpid <- IO () -> IO ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ LogAction IO (WithSeverity LspServerLog)
-> TChan Value -> (ByteString -> IO ()) -> IO ()
sendServer LogAction IO (WithSeverity LspServerLog)
ioLogger TChan Value
cout ByteString -> IO ()
clientOut

  let sendMsg :: a -> IO ()
sendMsg a
msg = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TChan a -> a -> STM ()
writeTChan TChan Value
cout forall a b. (a -> b) -> a -> b
$ forall a. ToJSON a => a -> Value
J.toJSON a
msg

  forall r. (VFS -> IO r) -> IO r
initVFS forall a b. (a -> b) -> a -> b
$ \VFS
vfs -> do
    forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
ioLoop LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ServerDefinition config
serverDefinition VFS
vfs forall {a}. ToJSON a => a -> IO ()
sendMsg

  forall (m :: * -> *) a. Monad m => a -> m a
return Int
1

-- ---------------------------------------------------------------------

ioLoop ::
  forall config
  .  LogAction IO (WithSeverity LspServerLog)
  -> LogAction (LspM config) (WithSeverity LspServerLog)
  -> IO BS.ByteString
  -> ServerDefinition config
  -> VFS
  -> (FromServerMessage -> IO ())
  -> IO ()
ioLoop :: forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
ioLoop LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
sendMsg = do
  Maybe (ByteString, ByteString)
minitialize <- forall (m :: * -> *).
MonadIO m =>
LogAction m (WithSeverity LspServerLog)
-> IO ByteString
-> Result ByteString
-> m (Maybe (ByteString, ByteString))
parseOne LogAction IO (WithSeverity LspServerLog)
ioLogger IO ByteString
clientIn (forall a. Parser a -> ByteString -> Result a
parse Parser ByteString ByteString
parser ByteString
"")
  case Maybe (ByteString, ByteString)
minitialize of
    Maybe (ByteString, ByteString)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just (ByteString
msg,ByteString
remainder) -> do
      case forall a. FromJSON a => ByteString -> Either String a
J.eitherDecode forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.fromStrict ByteString
msg of
        Left String
err -> LogAction IO (WithSeverity LspServerLog)
ioLogger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& String -> LspServerLog
DecodeInitializeError String
err forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Error
        Right RequestMessage @'FromClient 'Initialize
initialize -> do
          Maybe (LanguageContextEnv config)
mInitResp <- forall config.
ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> Message @'FromClient @'Request 'Initialize
-> IO (Maybe (LanguageContextEnv config))
Processing.initializeRequestHandler ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
sendMsg RequestMessage @'FromClient 'Initialize
initialize
          case Maybe (LanguageContextEnv config)
mInitResp of
            Maybe (LanguageContextEnv config)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just LanguageContextEnv config
env -> forall config (m :: * -> *) a.
LanguageContextEnv config -> LspT config m a -> m a
runLspT LanguageContextEnv config
env forall a b. (a -> b) -> a -> b
$ Result ByteString -> LspM config ()
loop (forall a. Parser a -> ByteString -> Result a
parse Parser ByteString ByteString
parser ByteString
remainder)
  where

    loop :: Result BS.ByteString -> LspM config ()
    loop :: Result ByteString -> LspM config ()
loop = Result ByteString -> LspM config ()
go
      where
        pLogger :: LogAction (LspM config) (WithSeverity LspProcessingLog)
pLogger =  forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LspProcessingLog -> LspServerLog
LspProcessingLog) LogAction (LspM config) (WithSeverity LspServerLog)
logger
        go :: Result ByteString -> LspM config ()
go Result ByteString
r = do
          Maybe (ByteString, ByteString)
res <- forall (m :: * -> *).
MonadIO m =>
LogAction m (WithSeverity LspServerLog)
-> IO ByteString
-> Result ByteString
-> m (Maybe (ByteString, ByteString))
parseOne LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn Result ByteString
r
          case Maybe (ByteString, ByteString)
res of
            Maybe (ByteString, ByteString)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just (ByteString
msg,ByteString
remainder) -> do
              forall (m :: * -> *) config.
((m :: (* -> *)) ~ (LspM config :: (* -> *))) =>
LogAction m (WithSeverity LspProcessingLog) -> ByteString -> m ()
Processing.processMessage LogAction (LspM config) (WithSeverity LspProcessingLog)
pLogger forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.fromStrict ByteString
msg
              Result ByteString -> LspM config ()
go (forall a. Parser a -> ByteString -> Result a
parse Parser ByteString ByteString
parser ByteString
remainder)

    parser :: Parser ByteString ByteString
parser = do
      ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
"Content-Length: "
      Int
len <- forall a. Integral a => Parser a
decimal
      ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
_TWO_CRLF
      Int -> Parser ByteString ByteString
Attoparsec.take Int
len

parseOne ::
  MonadIO m
  => LogAction m (WithSeverity LspServerLog)
  -> IO BS.ByteString
  -> Result BS.ByteString
  -> m (Maybe (BS.ByteString,BS.ByteString))
parseOne :: forall (m :: * -> *).
MonadIO m =>
LogAction m (WithSeverity LspServerLog)
-> IO ByteString
-> Result ByteString
-> m (Maybe (ByteString, ByteString))
parseOne LogAction m (WithSeverity LspServerLog)
logger IO ByteString
clientIn = Result ByteString -> m (Maybe (ByteString, ByteString))
go
  where
    go :: Result ByteString -> m (Maybe (ByteString, ByteString))
go (Fail ByteString
_ [String]
ctxs String
err) = do
      LogAction m (WithSeverity LspServerLog)
logger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& [String] -> String -> LspServerLog
HeaderParseFail [String]
ctxs String
err forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Error
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    go (Partial ByteString -> Result ByteString
c) = do
      ByteString
bs <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ByteString
clientIn
      if ByteString -> Bool
BS.null ByteString
bs
        then do
          LogAction m (WithSeverity LspServerLog)
logger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& LspServerLog
EOF forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Error
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
        else Result ByteString -> m (Maybe (ByteString, ByteString))
go (ByteString -> Result ByteString
c ByteString
bs)
    go (Done ByteString
remainder ByteString
msg) = do
      LogAction m (WithSeverity LspServerLog)
logger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text -> LspServerLog
ParsedMsg (ByteString -> Text
T.decodeUtf8 ByteString
msg) forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Debug
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (ByteString
msg,ByteString
remainder)

-- ---------------------------------------------------------------------

-- | Simple server to make sure all output is serialised
sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.ByteString -> IO ()) -> IO ()
sendServer :: LogAction IO (WithSeverity LspServerLog)
-> TChan Value -> (ByteString -> IO ()) -> IO ()
sendServer LogAction IO (WithSeverity LspServerLog)
logger TChan Value
msgChan ByteString -> IO ()
clientOut = do
  forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ do
    Value
msg <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TChan a -> STM a
readTChan TChan Value
msgChan

    -- We need to make sure we only send over the content of the message,
    -- and no other tags/wrapper stuff
    let str :: ByteString
str = forall a. ToJSON a => a -> ByteString
J.encode Value
msg

    let out :: ByteString
out = [ByteString] -> ByteString
BSL.concat
                [ Text -> ByteString
TL.encodeUtf8 forall a b. (a -> b) -> a -> b
$ String -> Text
TL.pack forall a b. (a -> b) -> a -> b
$ String
"Content-Length: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (ByteString -> Int64
BSL.length ByteString
str)
                , ByteString -> ByteString
BSL.fromStrict ByteString
_TWO_CRLF
                , ByteString
str ]

    ByteString -> IO ()
clientOut ByteString
out
    LogAction IO (WithSeverity LspServerLog)
logger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& Text -> LspServerLog
SendMsg (ByteString -> Text
TL.decodeUtf8 ByteString
str) forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Debug

-- |
--
--
_TWO_CRLF :: BS.ByteString
_TWO_CRLF :: ByteString
_TWO_CRLF = ByteString
"\r\n\r\n"