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

module Datalog.LSP.Highlight where

import Data.Either (fromRight)
import Data.Text qualified as T
import Datalog.LSP.Types (DLogLspM, LSPContext (..))
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
import Control.Monad.Trans
import Control.Concurrent.STM
import qualified Data.Map as M
import Datalog.LSP.Utils (currentBufferUri)

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
            let semanticTokens = fromRight (error "") $ makeSemanticTokens defaultSemanticTokensLegend $ highlightProg prog
            responder (Right $ InL semanticTokens)

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 = []
        }