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

module Datalog.LSP.Highlight where

import Control.Lens (Lens', (^.))
import Data.Either (fromRight)
import Data.Maybe (fromJust)
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 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

tokenHandler :: Handlers (LspM ())
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

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

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
        , _tokenModifiers = []
        }
    ]
  where
    loc = getConLoc t
    unPos' = fromIntegral . unPos
    startLine = sourceLine . start $ loc
    startCol = sourceColumn . start $ loc
    stopCol = sourceColumn . end $ loc