From 108001e98794f4170fa89b45b91f92761e74a436 Mon Sep 17 00:00:00 2001
From: Patrick Aldis <patricktaldis@gmail.com>
Date: Tue, 3 Mar 2026 16:00:36 +0000
Subject: [PATCH] Rudimentary highlighting

---
 datalog-lsp/src/Datalog/LSP/Highlight.hs | 62 ++++++++++++++++++------
 datalog/src/Datalog/Parser.hs            |  1 +
 flake.nix                                |  2 +-
 3 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/datalog-lsp/src/Datalog/LSP/Highlight.hs b/datalog-lsp/src/Datalog/LSP/Highlight.hs
index 0a5890e..45cb9f1 100644
--- a/datalog-lsp/src/Datalog/LSP/Highlight.hs
+++ b/datalog-lsp/src/Datalog/LSP/Highlight.hs
@@ -1,28 +1,60 @@
 module Datalog.LSP.Highlight where
 
-import Control.Lens ((^.), Lens')
+import Control.Lens (Lens', (^.))
 import Data.Either (fromRight)
 import Data.Maybe (fromJust)
-import Language.LSP.Protocol.Lens (params, textDocument, uri, HasTextDocument, HasParams, HasUri)
+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)
 
 tokenHandler :: Handlers (LspM ())
 tokenHandler = requestHandler SMethod_TextDocumentSemanticTokensFull $ \req responder -> do
-    let 
-        openUri = toNormalizedUri $ req ^. docUri
-        token =
-            SemanticTokenAbsolute
-                { _line = 0
-                , _startChar = 0
-                , _length = 5
-                , _tokenType = SemanticTokenTypes_Keyword
-                , _tokenModifiers = []
-                }
-        tokens = fromRight (error "failed") $ makeSemanticTokens defaultSemanticTokensLegend [token]
-    c <- fromJust <$> getVirtualFile openUri
+    c <- virtualFileText . fromJust <$> getVirtualFile (toNormalizedUri $ req ^. docUri)
+    let
+        parsedProgram = fromRight (error "") $ runParser parseProgram "" c
+        absTokens = highlightProg parsedProgram
+        tokens = fromRight (error "")$ makeSemanticTokens defaultSemanticTokensLegend absTokens
+
     responder (Right $ InL tokens)
 
-docUri :: (HasParams s s1,  HasTextDocument s1 s2, HasUri s2 a) => Lens' s a 
+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 = []
+
+getConLoc :: Term' SrcLoc -> SrcLoc
+getConLoc (Con loc _) = loc
+getConLoc (Var loc _) = loc
+
+highlightTerm :: Term' SrcLoc -> [SemanticTokenAbsolute]
+highlightTerm t =
+    [ SemanticTokenAbsolute
+        { _line = unPos' startLine
+        , _startChar = unPos' startCol
+        , _length = fromIntegral $ unPos stopCol - unPos startCol
+        , _tokenType = case t of
+            Con _ _ -> SemanticTokenTypes_Number
+            Var _ _ -> SemanticTokenTypes_Variable
+        , _tokenModifiers = []
+        }
+    ]
+  where
+    loc = getConLoc t
+    unPos' = fromIntegral . unPos
+    startLine = sourceLine . start $ loc
+    startCol = sourceColumn . start $ loc
+    stopCol = sourceColumn . end $ loc
diff --git a/datalog/src/Datalog/Parser.hs b/datalog/src/Datalog/Parser.hs
index 7a80615..1c3324c 100644
--- a/datalog/src/Datalog/Parser.hs
+++ b/datalog/src/Datalog/Parser.hs
@@ -5,6 +5,7 @@ module Datalog.Parser (
     parseTerm,
     parseRule,
     parseProgram,
+    SrcLoc (..),
 )
 where
 
diff --git a/flake.nix b/flake.nix
index 394def1..33a4e5a 100644
--- a/flake.nix
+++ b/flake.nix
@@ -43,7 +43,7 @@
             ghcid
           ];
         };
-        packages.default = haskell.packages."geolog-lsp:exe:geolog-lsp";
+        packages.default = haskell.packages."geolog-lsp:exe:datalog-lsp";
       }
     );
 }