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

module Datalog.LSP.Highlight where

import Control.Concurrent.STM
import Control.Monad.Trans
import Data.Either (fromRight)
import Data.Map qualified as M
import Data.Text qualified as T
import Datalog.LSP.Types (DLogLspM, LSPContext (..))
import Datalog.LSP.Utils (currentBufferUri)
import Datalog.Parser (SrcLoc (..))
import Datalog.Syntax (Atom' (Atom), Program' (..), RelId (RelId), Rule' (..), Term' (..))
import Language.LSP.Protocol.Message
import Language.LSP.Protocol.Types
import Language.LSP.Server

tokenHandler :: Handlers DLogLspM
tokenHandler = requestHandler SMethod_TextDocumentSemanticTokensFull $ \req responder -> do
    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
            responder
                ( Right
                    . InL
                    . fromRight (error "")
                    . makeSemanticTokens defaultSemanticTokensLegend
                    $ highlightProg prog
                )

highlightProg :: Program' SrcLoc -> [SemanticTokenAbsolute]
highlightProg (Program _ rs) = rs >>= highlightRule

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 = pure $ tokenFromSrcLoc loc' SemanticTokenTypes_Interface
    loc' = loc{endCol = startCol loc + length (T.unpack relId)}

highlightTerm :: Term' SrcLoc -> [SemanticTokenAbsolute]
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 = []
        }