X-Git-Url: http://git.lukelau.me/?a=blobdiff_plain;f=src%2FLanguage%2FHaskell%2FLSP%2FTest.hs;h=1c425fb89cc49aa0af526b99ab7d5753d261a679;hb=b7ee75f11c842d84221eec57715d96429eb1b689;hp=15cb2a164d4600530cbc6a626f8fedbfee7890f9;hpb=01cfd5ae44e86182adf1cee0a3e5769e3f1586d3;p=lsp-test.git

diff --git a/src/Language/Haskell/LSP/Test.hs b/src/Language/Haskell/LSP/Test.hs
index 15cb2a1..1c425fb 100644
--- a/src/Language/Haskell/LSP/Test.hs
+++ b/src/Language/Haskell/LSP/Test.hs
@@ -79,6 +79,8 @@ module Language.Haskell.LSP.Test
   , formatRange
   -- ** Edits
   , applyEdit
+  -- ** Code lenses
+  , getCodeLenses
   ) where
 
 import Control.Applicative.Combinators
@@ -92,6 +94,7 @@ import qualified Data.Text.IO as T
 import Data.Aeson
 import Data.Default
 import qualified Data.HashMap.Strict as HashMap
+import Data.IORef
 import qualified Data.Map as Map
 import Data.Maybe
 import Language.Haskell.LSP.Types
@@ -135,6 +138,8 @@ runSessionWithConfig :: SessionConfig -- ^ Configuration options for the session
                      -> Session a -- ^ The session to run.
                      -> IO a
 runSessionWithConfig config serverExe caps rootDir session = do
+  -- We use this IORef to make exception non-fatal when the server is supposed to shutdown.
+  exitOk <- newIORef False
   pid <- getCurrentProcessID
   absRootDir <- canonicalizePath rootDir
 
@@ -146,7 +151,7 @@ runSessionWithConfig config serverExe caps rootDir session = do
                                           (Just TraceOff)
                                           Nothing
   withServer serverExe (logStdErr config) $ \serverIn serverOut _ ->
-    runSessionWithHandles serverIn serverOut listenServer config caps rootDir $ do
+    runSessionWithHandles serverIn serverOut (\h c -> catchWhenTrue exitOk $ listenServer h c) config caps rootDir $ do
 
       -- Wrap the session around initialize and shutdown calls
       initRspMsg <- request Initialize initializeParams :: Session InitializeResponse
@@ -165,12 +170,22 @@ runSessionWithConfig config serverExe caps rootDir session = do
       -- Run the actual test
       result <- session
 
+      liftIO $ atomicWriteIORef exitOk True
       sendNotification Exit ExitParams
 
       return result
   where
+  catchWhenTrue :: IORef Bool -> IO () -> IO ()
+  catchWhenTrue exitOk a =
+      a `catch` (\e -> do
+          x <- readIORef exitOk
+          unless x $ throw (e :: SomeException))
+
   -- | Listens to the server output, makes sure it matches the record and
   -- signals any semaphores
+  -- Note that on Windows, we cannot kill a thread stuck in getNextMessage.
+  -- So we have to wait for the exit notification to kill the process first
+  -- and then getNextMessage will fail.
   listenServer :: Handle -> SessionContext -> IO ()
   listenServer serverOut context = do
     msgBytes <- getNextMessage serverOut
@@ -505,9 +520,12 @@ getReferences doc pos inclDecl =
 getDefinitions :: TextDocumentIdentifier -- ^ The document the term is in.
                -> Position -- ^ The position the term is at.
                -> Session [Location] -- ^ The location(s) of the definitions
-getDefinitions doc pos =
+getDefinitions doc pos = do
   let params = TextDocumentPositionParams doc pos
-  in getResponseResult <$> request TextDocumentDefinition params
+  rsp <- request TextDocumentDefinition params :: Session DefinitionResponse
+  case getResponseResult rsp of
+      SingleLoc loc -> pure [loc]
+      MultiLoc locs -> pure locs
 
 -- | Returns the type definition(s) for the term at the specified position.
 getTypeDefinitions :: TextDocumentIdentifier -- ^ The document the term is in.
@@ -564,3 +582,10 @@ applyTextEdits doc edits =
   let wEdit = WorkspaceEdit (Just (HashMap.singleton (doc ^. uri) edits)) Nothing
       req = RequestMessage "" (IdInt 0) WorkspaceApplyEdit (ApplyWorkspaceEditParams wEdit)
   in updateState (ReqApplyWorkspaceEdit req)
+
+-- | Returns the code lenses for the specified document.
+getCodeLenses :: TextDocumentIdentifier -> Session [CodeLens]
+getCodeLenses tId = do
+    rsp <- request TextDocumentCodeLens (CodeLensParams tId) :: Session CodeLensResponse
+    case getResponseResult rsp of
+        List res -> pure res