From 35b2fa4282b001ee19726e5fcf816840d7a33f38 Mon Sep 17 00:00:00 2001 From: Patrick Aldis Date: Thu, 5 Mar 2026 15:46:01 +0000 Subject: [PATCH] Use STM to share state between handlers --- datalog-lsp/src/Datalog/LSP.hs | 23 ++++--- datalog-lsp/src/Datalog/LSP/DocChange.hs | 31 +++++++++ datalog-lsp/src/Datalog/LSP/Highlight.hs | 87 +++++++++--------------- datalog-lsp/src/Datalog/LSP/Hover.hs | 3 +- datalog-lsp/src/Datalog/LSP/Types.hs | 20 ++++++ datalog-lsp/src/Datalog/LSP/Utils.hs | 16 +++++ datalog-lsp/src/Main.hs | 4 +- datalog/src/Datalog/Parser.hs | 32 ++++----- geolog-lsp.cabal | 8 ++- 9 files changed, 139 insertions(+), 85 deletions(-) create mode 100644 datalog-lsp/src/Datalog/LSP/DocChange.hs create mode 100644 datalog-lsp/src/Datalog/LSP/Types.hs create mode 100644 datalog-lsp/src/Datalog/LSP/Utils.hs diff --git a/datalog-lsp/src/Datalog/LSP.hs b/datalog-lsp/src/Datalog/LSP.hs index 6fe6131..c512ac1 100644 --- a/datalog-lsp/src/Datalog/LSP.hs +++ b/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 diff --git a/datalog-lsp/src/Datalog/LSP/DocChange.hs b/datalog-lsp/src/Datalog/LSP/DocChange.hs new file mode 100644 index 0000000..67fdd31 --- /dev/null +++ b/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 diff --git a/datalog-lsp/src/Datalog/LSP/Highlight.hs b/datalog-lsp/src/Datalog/LSP/Highlight.hs index ba8a34f..d627bed 100644 --- a/datalog-lsp/src/Datalog/LSP/Highlight.hs +++ b/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 diff --git a/datalog-lsp/src/Datalog/LSP/Hover.hs b/datalog-lsp/src/Datalog/LSP/Hover.hs index ef11663..98a3b3f 100644 --- a/datalog-lsp/src/Datalog/LSP/Hover.hs +++ b/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 diff --git a/datalog-lsp/src/Datalog/LSP/Types.hs b/datalog-lsp/src/Datalog/LSP/Types.hs new file mode 100644 index 0000000..e35e50e --- /dev/null +++ b/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))) + } diff --git a/datalog-lsp/src/Datalog/LSP/Utils.hs b/datalog-lsp/src/Datalog/LSP/Utils.hs new file mode 100644 index 0000000..143ebec --- /dev/null +++ b/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) + diff --git a/datalog-lsp/src/Main.hs b/datalog-lsp/src/Main.hs index ee612d0..a165817 100644 --- a/datalog-lsp/src/Main.hs +++ b/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 diff --git a/datalog/src/Datalog/Parser.hs b/datalog/src/Datalog/Parser.hs index 1c3324c..3d19bb9 100644 --- a/datalog/src/Datalog/Parser.hs +++ b/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 diff --git a/geolog-lsp.cabal b/geolog-lsp.cabal index 7bc6a4b..d7cb21f 100644 --- a/geolog-lsp.cabal +++ b/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