Use STM to share state between handlers

datalog-lsp/src/Datalog/LSP.hs

@@ -1,4 +1,4 @@
-module Datalog.LSP where
+module Datalog.LSP (serverDefinition) where
 
 import Control.Monad.IO.Class
 import Language.LSP.Protocol.Message
@@ -6,24 +6,31 @@ import Language.LSP.Protocol.Types
 import Language.LSP.Server
 import Datalog.LSP.Hover (hoverHandler)
 import Datalog.LSP.Highlight (tokenHandler)
+import Datalog.LSP.Types (DLogLspM, LSPContext (LSPContext))
+import Datalog.LSP.DocChange (docChangeHandler, docOpenHandler)
+import qualified Data.Map as M
+import Control.Concurrent.STM
 
-handlers :: Handlers (LspM ())
+handlers :: Handlers DLogLspM
 handlers =
     mconcat
         [ initHandler
+        , docChangeHandler
+        , docOpenHandler
         , hoverHandler
         , tokenHandler
         ]
 
-initHandler :: Handlers (LspM ())
+initHandler :: Handlers DLogLspM
 initHandler = notificationHandler SMethod_Initialized $ \_ -> pure ()
 
-serverDefinition :: ServerDefinition ()
-serverDefinition = 
-        ServerDefinition
-            { parseConfig = const $ const $ Right ()
+serverDefinition :: IO (ServerDefinition LSPContext)
+serverDefinition = do
+        ref <- newTVarIO M.empty
+        pure $ ServerDefinition
+            { parseConfig = \c v -> Right c
             , onConfigChange = const $ pure ()
-            , defaultConfig = ()
+            , defaultConfig = LSPContext ref
             , configSection = "demo"
             , doInitialize = \env _req -> pure $ Right env
             , staticHandlers = const handlers

datalog-lsp/src/Datalog/LSP/DocChange.hs

@@ -0,0 +1,31 @@
+{-# LANGUAGE BlockArguments #-}
+
+module Datalog.LSP.DocChange (docChangeHandler, docOpenHandler) where
+
+import Control.Concurrent.STM
+import Control.Monad.Trans
+import Data.Map qualified as M
+import Datalog.LSP.Types
+import Datalog.LSP.Utils (currentBufferText, currentBufferUri)
+import Datalog.Parser (parseProgram)
+import Language.LSP.Protocol.Lens (HasParams, HasTextDocument, HasUri)
+import Language.LSP.Protocol.Message (SMethod (SMethod_TextDocumentDidChange, SMethod_TextDocumentDidOpen))
+import Language.LSP.Protocol.Types (Uri)
+import Language.LSP.Server (Handlers, MonadLsp, getConfig, notificationHandler)
+import Text.Megaparsec
+
+docOpenHandler :: Handlers DLogLspM
+docOpenHandler = notificationHandler SMethod_TextDocumentDidOpen updateState
+
+docChangeHandler :: Handlers DLogLspM
+docChangeHandler = notificationHandler SMethod_TextDocumentDidChange updateState
+
+updateState :: (HasParams s a1, MonadLsp LSPContext (t IO), HasUri a2 Uri, HasTextDocument a1 a2, MonadTrans t) => s -> t IO ()
+updateState req = do
+    let uri = currentBufferUri req
+    LSPContext parseStateRef <- getConfig
+    text <- currentBufferText req
+    lift . atomically $ do
+        v <- readTVar parseStateRef
+        writeTVar parseStateRef $
+            M.insert uri (runParser parseProgram (show uri) text) v

datalog-lsp/src/Datalog/LSP/Highlight.hs

@@ -1,32 +1,29 @@
 module Datalog.LSP.Highlight where
 
-import Control.Lens (Lens', (^.))
 import Data.Either (fromRight)
-import Data.Maybe (fromJust)
+import Data.Text qualified as T
+import Datalog.LSP.Types (DLogLspM, LSPContext (..))
 import Datalog.Parser (SrcLoc (..))
-import Datalog.Syntax (Atom' (Atom), ConId (..), Program' (..), RelId (RelId), Rule' (..), Term' (..))
-import Language.LSP.Protocol.Lens (HasParams, HasTextDocument, HasUri, params, textDocument, uri)
+import Datalog.Syntax (Atom' (Atom), Program' (..), RelId (RelId), Rule' (..), Term' (..))
 import Language.LSP.Protocol.Message
 import Language.LSP.Protocol.Types
 import Language.LSP.Server
-import Text.Megaparsec
-import Language.LSP.VFS (virtualFileText)
-import Datalog.Parser (parseProgram)
-import qualified Data.Text as T
+import Control.Monad.Trans
+import Control.Concurrent.STM
+import qualified Data.Map as M
+import Datalog.LSP.Utils (currentBufferUri)
 
-tokenHandler :: Handlers (LspM ())
+tokenHandler :: Handlers DLogLspM
 tokenHandler = requestHandler SMethod_TextDocumentSemanticTokensFull $ \req responder -> do
-    c <- virtualFileText . fromJust <$> getVirtualFile (toNormalizedUri $ req ^. docUri)
-    let
-        absTokens = case runParser parseProgram "" c of 
-            Left _ -> []
-            Right prog -> highlightProg prog
-        tokens = fromRight (error "") $ makeSemanticTokens defaultSemanticTokensLegend absTokens
-
-    responder (Right $ InL tokens)
-
-docUri :: (HasParams s s1, HasTextDocument s1 s2, HasUri s2 a) => Lens' s a
-docUri = params . textDocument . uri
+    LSPContext parseRef <- getConfig
+    p <- lift . readTVarIO $ parseRef
+    case M.lookup (currentBufferUri req) p of
+        Nothing -> responder (Left $ TResponseError (InL LSPErrorCodes_RequestFailed) "Doc not bundle" Nothing)
+        Just (Left _) ->
+            responder (Left $ TResponseError (InL LSPErrorCodes_RequestFailed) "Failed to parse" Nothing)
+        Just (Right prog) -> do
+            let semanticTokens = fromRight (error "") $ makeSemanticTokens defaultSemanticTokensLegend $ highlightProg prog
+            responder (Right $ InL semanticTokens)
 
 highlightProg :: Program' SrcLoc -> [SemanticTokenAbsolute]
 highlightProg (Program _ rs) = rs >>= highlightRule
@@ -35,41 +32,23 @@ highlightRule :: Rule' SrcLoc -> [SemanticTokenAbsolute]
 highlightRule (Rule _ a as) = a : as >>= highlightAtom
 
 highlightAtom :: Atom' SrcLoc -> [SemanticTokenAbsolute]
-highlightAtom (Atom loc (RelId relId) ts) = highlightRel ++ (ts >>= highlightTerm) where
-    highlightRel =
-        [ SemanticTokenAbsolute
-            { _line = unPos' startLine - 1
-            , _startChar = unPos' startCol - 1
-            , _length = fromIntegral . length . T.unpack $ relId
-            , _tokenType = SemanticTokenTypes_Interface
-            , _tokenModifiers = []
-            }
-        ]
-        where
-            unPos' = fromIntegral . unPos
-            startLine = sourceLine . start $ loc
-            startCol = sourceColumn . start $ loc
-            stopCol = sourceColumn . end $ loc
-
-getConLoc :: Term' SrcLoc -> SrcLoc
-getConLoc (Con loc _) = loc
-getConLoc (Var loc _) = loc
+highlightAtom (Atom loc (RelId relId) ts) = highlightRel ++ (ts >>= highlightTerm)
+  where
+    highlightRel = pure $ tokenFromSrcLoc loc' SemanticTokenTypes_Interface
+    loc' = loc{endCol = startCol loc + length (T.unpack relId)}
 
 highlightTerm :: Term' SrcLoc -> [SemanticTokenAbsolute]
-highlightTerm t =
-    [ SemanticTokenAbsolute
-        { _line = unPos' startLine - 1
-        , _startChar = unPos' startCol - 1
-        , _length = fromIntegral $ unPos stopCol - unPos startCol
-        , _tokenType = case t of
-            Con _ _ -> SemanticTokenTypes_Number
-            Var _ _ -> SemanticTokenTypes_Keyword
+highlightTerm =
+    pure <$> \case
+        Con loc _ -> tokenFromSrcLoc loc SemanticTokenTypes_Number
+        Var loc _ -> tokenFromSrcLoc loc SemanticTokenTypes_Keyword
+
+tokenFromSrcLoc :: SrcLoc -> SemanticTokenTypes -> SemanticTokenAbsolute
+tokenFromSrcLoc (SrcLoc sl sc _ ec) tokenType =
+    SemanticTokenAbsolute
+        { _line = fromIntegral $ sl - 1
+        , _startChar = fromIntegral $ sc - 1
+        , _length = fromIntegral $ ec - sc
+        , _tokenType = tokenType
         , _tokenModifiers = []
         }
-    ]
-  where
-    loc = getConLoc t
-    unPos' = fromIntegral . unPos
-    startLine = sourceLine . start $ loc
-    startCol = sourceColumn . start $ loc
-    stopCol = sourceColumn . end $ loc

datalog-lsp/src/Datalog/LSP/Hover.hs

@@ -7,8 +7,9 @@ import Language.LSP.Protocol.Message
 import Language.LSP.Protocol.Types
 import Language.LSP.Server
 import Language.LSP.VFS (virtualFileText)
+import Datalog.LSP.Types (DLogLspM)
 
-hoverHandler :: Handlers (LspM ())
+hoverHandler :: Handlers DLogLspM
 hoverHandler = requestHandler SMethod_TextDocumentHover $ \req responder -> do
     let
         openUri = toNormalizedUri $ req ^. docUri

datalog-lsp/src/Datalog/LSP/Types.hs

@@ -0,0 +1,20 @@
+module Datalog.LSP.Types (DLogLspM, LSPContext(..)) where
+
+import Control.Monad.Reader (ReaderT)
+import Language.LSP.Server (LspM)
+import Control.Concurrent.STM (TVar)
+import Text.Megaparsec
+import Data.Text (Text)
+import Data.Void (Void)
+import Data.Map (Map)
+import Language.LSP.Protocol.Types qualified as J
+import Datalog.Syntax (Program')
+import Datalog.Parser (SrcLoc)
+
+type DLogLspM = LspM LSPContext
+
+type UriBundle a = Map J.NormalizedUri a
+
+data LSPContext = LSPContext
+    { parseState :: TVar (UriBundle (Either (ParseErrorBundle Text Void) (Program' SrcLoc)))
+    }

datalog-lsp/src/Datalog/LSP/Utils.hs

@@ -0,0 +1,16 @@
+module Datalog.LSP.Utils where
+
+import Control.Lens
+import Data.Maybe
+import Data.Text qualified as T
+import Language.LSP.Protocol.Lens
+import Language.LSP.Protocol.Types
+import Language.LSP.Server
+import Language.LSP.VFS
+
+currentBufferText :: (MonadLsp config f, HasParams s a1, HasTextDocument a1 a2,  HasUri a2 Uri) => s -> f T.Text
+currentBufferText req = virtualFileText . fromJust <$> getVirtualFile (currentBufferUri req)
+
+currentBufferUri :: (HasParams s a1, HasTextDocument a1 a2, HasUri a2 Uri) => s -> NormalizedUri
+currentBufferUri req = toNormalizedUri $ req ^. (params . textDocument . uri)
+

datalog-lsp/src/Main.hs

@@ -1,7 +1,5 @@
-{-# LANGUAGE OverloadedStrings #-}
-
 import Datalog.LSP (serverDefinition)
 import Language.LSP.Server (runServer)
 
 main :: IO Int
-main = runServer serverDefinition
+main = serverDefinition >>= runServer 

datalog/src/Datalog/Parser.hs

@@ -12,8 +12,8 @@ where
 import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Void (Void)
-import Datalog.Syntax hiding (Atom, Program, Rule, Term)
 import Datalog.Syntax (Atom' (..), Program' (..), Rule' (..), Term' (..))
+import Datalog.Syntax hiding (Atom, Program, Rule, Term)
 import Text.Megaparsec
 import Text.Megaparsec.Char
 import Text.Megaparsec.Char.Lexer qualified as L
@@ -21,6 +21,14 @@ import Text.Pretty.Simple
 
 type Parser = Parsec Void Text
 
+data SrcLoc = SrcLoc
+    { startLine :: Int
+    , startCol :: Int
+    , endLine :: Int
+    , endCol :: Int
+    }
+    deriving (Show)
+
 type Atom = Atom' SrcLoc
 type Term = Term' SrcLoc
 type Rule = Rule' SrcLoc
@@ -65,7 +73,8 @@ parseQuery :: Parser [Atom]
 parseQuery = parseAtom `sepBy` comma
 
 parseRule :: Parser Rule
-parseRule = parseThingWithSub (\loc (a, as) -> Rule loc a as) $
+parseRule =
+    parseThingWithSub (\loc (a, as) -> Rule loc a as) $
         try rule1 <|> rule2
   where
     rule1 = do
@@ -83,23 +92,10 @@ parseProgram = parseThingWithSub Program (many parseRule)
 
 parseThingWithSub :: (SrcLoc -> c -> f SrcLoc) -> Parser c -> Parser (f SrcLoc)
 parseThingWithSub f parseSub = do
-    s <- getSourcePos
+    SourcePos _ sl sc <- getSourcePos
     c <- parseSub
-    e <- getSourcePos
-    pure $ f (SrcLoc s e) c
-
-annotateSrcLoc :: (Functor f) => Parser (f a) -> Parser (f SrcLoc)
-annotateSrcLoc p = do
-    s <- getSourcePos
-    res <- p
-    f <- getSourcePos
-    pure (SrcLoc s f <$ res)
-
-data SrcLoc = SrcLoc
-    { start :: SourcePos
-    , end :: SourcePos
-    }
-    deriving (Show)
+    SourcePos _ el ec <- getSourcePos
+    pure $ f (SrcLoc (unPos sl) (unPos sc) (unPos el) (unPos ec)) c
 
 test = do
     let r = runParser parseProgram "???" prog

geolog-lsp.cabal

@@ -28,6 +28,7 @@ library datalog-parser
         containers,
         megaparsec,
         pretty-simple,
+        lens
     exposed-modules:
         Datalog.Parser,
         Datalog.Syntax
@@ -44,10 +45,15 @@ executable datalog-lsp
                       containers,
                       lens,
                       megaparsec,
-                      datalog-parser
+                      datalog-parser,
+                      mtl,
+                      stm
     other-modules:
         Datalog.LSP
+        Datalog.LSP.DocChange
+        Datalog.LSP.Types
         Datalog.LSP.Highlight
         Datalog.LSP.Hover
+        Datalog.LSP.Utils
     hs-source-dirs:   datalog-lsp/src
     default-language: GHC2024