Add Method_TextDocumentSemanticTokensFullDelta (haskell#4073)

* add Method_TextDocumentSemanticTokensFullDelta

* remove persistentGetSemanticTokensRule

* add doc about semanticTokensCache location

* add Note [Semantic Tokens Cache Location]

---------

Co-authored-by: fendor <fendor@users.noreply.github.com>

ghcide/src/Development/IDE/Core/Shake.hs

@@ -164,6 +164,7 @@ import           Language.LSP.Diagnostics
 import qualified Language.LSP.Protocol.Lens             as L
 import           Language.LSP.Protocol.Message
 import           Language.LSP.Protocol.Types
+import           Language.LSP.Protocol.Types            (SemanticTokens)
 import qualified Language.LSP.Protocol.Types            as LSP
 import qualified Language.LSP.Server                    as LSP
 import           Language.LSP.VFS                       hiding (start)
@@ -243,6 +244,13 @@ data HieDbWriter
 -- with (currently) retry functionality
 type IndexQueue = TQueue (((HieDb -> IO ()) -> IO ()) -> IO ())
 
+-- Note [Semantic Tokens Cache Location]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+-- storing semantic tokens cache for each file in shakeExtras might
+-- not be ideal, since it most used in LSP request handlers
+-- instead of rules. We should consider moving it to a more
+-- appropriate place in the future if we find one, store it for now.
+
 -- information we stash inside the shakeExtra field
 data ShakeExtras = ShakeExtras
     { --eventer :: LSP.FromServerMessage -> IO ()
@@ -259,6 +267,14 @@ data ShakeExtras = ShakeExtras
     ,publishedDiagnostics :: STM.Map NormalizedUri [Diagnostic]
     -- ^ This represents the set of diagnostics that we have published.
     -- Due to debouncing not every change might get published.
+
+    ,semanticTokensCache:: STM.Map NormalizedFilePath SemanticTokens
+    -- ^ Cache of last response of semantic tokens for each file,
+    -- so we can compute deltas for semantic tokens(SMethod_TextDocumentSemanticTokensFullDelta).
+    -- putting semantic tokens cache and id in shakeExtras might not be ideal
+    -- see Note [Semantic Tokens Cache Location]
+    ,semanticTokensId :: TVar Int
+    -- ^ semanticTokensId is used to generate unique ids for each lsp response of semantic tokens.
     ,positionMapping :: STM.Map NormalizedUri (EnumMap Int32 (PositionDelta, PositionMapping))
     -- ^ Map from a text document version to a PositionMapping that describes how to map
     -- positions in a version of that document to positions in the latest version
@@ -616,12 +632,14 @@ shakeOpen recorder lspEnv defaultConfig idePlugins logger debouncer
         diagnostics <- STM.newIO
         hiddenDiagnostics <- STM.newIO
         publishedDiagnostics <- STM.newIO
+        semanticTokensCache <- STM.newIO
         positionMapping <- STM.newIO
         knownTargetsVar <- newTVarIO $ hashed HMap.empty
         let restartShakeSession = shakeRestart recorder ideState
         persistentKeys <- newTVarIO mempty
         indexPending <- newTVarIO HMap.empty
         indexCompleted <- newTVarIO 0
+        semanticTokensId <- newTVarIO 0
         indexProgressToken <- newVar Nothing
         let hiedbWriter = HieDbWriter{..}
         exportsMap <- newTVarIO mempty

haskell-language-server.cabal

@@ -1574,14 +1574,16 @@ library hls-semantic-tokens-plugin
     , hls-graph == 2.6.0.0
     , template-haskell
     , data-default
+    , stm
+    , stm-containers
 
   default-extensions: DataKinds
 
 test-suite hls-semantic-tokens-plugin-tests
   import:           defaults, pedantic, test-defaults, warnings
   type:             exitcode-stdio-1.0
   hs-source-dirs:   plugins/hls-semantic-tokens-plugin/test
-  main-is:          Main.hs
+  main-is:          SemanticTokensTest.hs
 
   build-depends:
     , aeson
@@ -1601,6 +1603,7 @@ test-suite hls-semantic-tokens-plugin-tests
     , ghcide                == 2.6.0.0
     , hls-plugin-api        == 2.6.0.0
     , data-default
+    , row-types
 
 -----------------------------
 -- HLS

hls-plugin-api/src/Ide/Plugin/ConfigUtils.hs

@@ -94,6 +94,7 @@ pluginsToDefaultConfig IdePlugins {..} =
           SMethod_TextDocumentCompletion           -> ["completionOn" A..= plcCompletionOn]
           SMethod_TextDocumentPrepareCallHierarchy -> ["callHierarchyOn" A..= plcCallHierarchyOn]
           SMethod_TextDocumentSemanticTokensFull   -> ["semanticTokensOn" A..= plcSemanticTokensOn]
+          SMethod_TextDocumentSemanticTokensFullDelta -> ["semanticTokensOn" A..= plcSemanticTokensOn]
           _                                 -> []
 
 -- | Generates json schema used in haskell vscode extension
@@ -125,6 +126,7 @@ pluginsToVSCodeExtensionSchema IdePlugins {..} = A.object $ mconcat $ singlePlug
           SMethod_TextDocumentCompletion           -> [toKey' "completionOn" A..= schemaEntry "completions" plcCompletionOn]
           SMethod_TextDocumentPrepareCallHierarchy -> [toKey' "callHierarchyOn" A..= schemaEntry "call hierarchy" plcCallHierarchyOn]
           SMethod_TextDocumentSemanticTokensFull   -> [toKey' "semanticTokensOn" A..= schemaEntry "semantic tokens" plcSemanticTokensOn]
+          SMethod_TextDocumentSemanticTokensFullDelta   -> [toKey' "semanticTokensOn" A..= schemaEntry "semantic tokens" plcSemanticTokensOn]
           _                                        -> []
         schemaEntry desc defaultVal =
           A.object

hls-plugin-api/src/Ide/Types.hs

@@ -511,6 +511,9 @@ instance PluginMethod Request Method_TextDocumentRangeFormatting where
 instance PluginMethod Request Method_TextDocumentSemanticTokensFull where
   handlesRequest = pluginEnabledWithFeature plcSemanticTokensOn
 
+instance PluginMethod Request Method_TextDocumentSemanticTokensFullDelta where
+  handlesRequest = pluginEnabledWithFeature plcSemanticTokensOn
+
 instance PluginMethod Request Method_TextDocumentPrepareCallHierarchy where
   handlesRequest = pluginEnabledWithFeature plcCallHierarchyOn
 
@@ -751,6 +754,9 @@ instance PluginRequestMethod (Method_CustomMethod m) where
 instance PluginRequestMethod Method_TextDocumentSemanticTokensFull where
   combineResponses _ _ _ _ (x :| _) = x
 
+instance PluginRequestMethod Method_TextDocumentSemanticTokensFullDelta where
+  combineResponses _ _ _ _ (x :| _) = x
+
 takeLefts :: [a |? b] -> [a]
 takeLefts = mapMaybe (\x -> [res | (InL res) <- Just x])
 

plugins/hls-semantic-tokens-plugin/src/Ide/Plugin/SemanticTokens.hs

@@ -1,5 +1,5 @@
 {-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE TemplateHaskell   #-}
+
 
 module Ide.Plugin.SemanticTokens (descriptor) where
 
@@ -12,8 +12,10 @@ import           Language.LSP.Protocol.Message
 descriptor :: Recorder (WithPriority SemanticLog) -> PluginId -> PluginDescriptor IdeState
 descriptor recorder plId =
   (defaultPluginDescriptor plId "Provides semantic tokens")
-    { Ide.Types.pluginHandlers = mkPluginHandler SMethod_TextDocumentSemanticTokensFull (Internal.semanticTokensFull recorder),
-      Ide.Types.pluginRules = Internal.getSemanticTokensRule recorder <> Internal.persistentGetSemanticTokensRule,
+    { Ide.Types.pluginHandlers =
+        mkPluginHandler SMethod_TextDocumentSemanticTokensFull (Internal.semanticTokensFull recorder)
+        <> mkPluginHandler SMethod_TextDocumentSemanticTokensFullDelta (Internal.semanticTokensFullDelta recorder),
+      Ide.Types.pluginRules = Internal.getSemanticTokensRule recorder,
       pluginConfigDescriptor =
         defaultConfigDescriptor
           { configInitialGenericConfig = (configInitialGenericConfig defaultConfigDescriptor) {plcGlobalOn = False}

plugins/hls-semantic-tokens-plugin/src/Ide/Plugin/SemanticTokens/Internal.hs

@@ -10,14 +10,19 @@
 
 -- |
 -- This module provides the core functionality of the plugin.
-module Ide.Plugin.SemanticTokens.Internal (semanticTokensFull, getSemanticTokensRule, persistentGetSemanticTokensRule, semanticConfigProperties) where
+module Ide.Plugin.SemanticTokens.Internal (semanticTokensFull, getSemanticTokensRule, semanticConfigProperties, semanticTokensFullDelta) where
 
+import           Control.Concurrent.STM                   (stateTVar)
+import           Control.Concurrent.STM.Stats             (atomically)
 import           Control.Lens                             ((^.))
 import           Control.Monad.Except                     (ExceptT, liftEither,
                                                            withExceptT)
+import           Control.Monad.IO.Class                   (MonadIO (..))
 import           Control.Monad.Trans                      (lift)
 import           Control.Monad.Trans.Except               (runExceptT)
 import qualified Data.Map.Strict                          as M
+import           Data.Text                                (Text)
+import qualified Data.Text                                as T
 import           Development.IDE                          (Action,
                                                            GetDocMap (GetDocMap),
                                                            GetHieAst (GetHieAst),
@@ -31,10 +36,10 @@ import           Development.IDE                          (Action,
                                                            hieKind, use_)
 import           Development.IDE.Core.PluginUtils         (runActionE,
                                                            useWithStaleE)
-import           Development.IDE.Core.PositionMapping     (idDelta)
 import           Development.IDE.Core.Rules               (toIdeResult)
 import           Development.IDE.Core.RuleTypes           (DocAndTyThingMap (..))
-import           Development.IDE.Core.Shake               (addPersistentRule,
+import           Development.IDE.Core.Shake               (ShakeExtras (..),
+                                                           getShakeExtras,
                                                            getVirtualFile,
                                                            useWithStale_)
 import           Development.IDE.GHC.Compat               hiding (Warning)
@@ -51,11 +56,13 @@ import           Ide.Plugin.SemanticTokens.Tokenize       (computeRangeHsSemanti
 import           Ide.Plugin.SemanticTokens.Types
 import           Ide.Types
 import qualified Language.LSP.Protocol.Lens               as L
-import           Language.LSP.Protocol.Message            (Method (Method_TextDocumentSemanticTokensFull))
+import           Language.LSP.Protocol.Message            (MessageResult,
+                                                           Method (Method_TextDocumentSemanticTokensFull, Method_TextDocumentSemanticTokensFullDelta))
 import           Language.LSP.Protocol.Types              (NormalizedFilePath,
                                                            SemanticTokens,
-                                                           type (|?) (InL))
+                                                           type (|?) (InL, InR))
 import           Prelude                                  hiding (span)
+import qualified StmContainers.Map                        as STM
 
 
 $mkSemanticConfigFunctions
@@ -68,14 +75,40 @@ computeSemanticTokens :: Recorder (WithPriority SemanticLog) -> PluginId -> IdeS
 computeSemanticTokens recorder pid _ nfp = do
   config <- lift $ useSemanticConfigAction pid
   logWith recorder Debug (LogConfig config)
+  semanticId <- lift getAndIncreaseSemanticTokensId
   (RangeHsSemanticTokenTypes {rangeSemanticList}, mapping) <- useWithStaleE GetSemanticTokens nfp
-  withExceptT PluginInternalError $ liftEither $ rangeSemanticsSemanticTokens config mapping rangeSemanticList
+  withExceptT PluginInternalError $ liftEither $ rangeSemanticsSemanticTokens semanticId config mapping rangeSemanticList
 
 semanticTokensFull :: Recorder (WithPriority SemanticLog) -> PluginMethodHandler IdeState 'Method_TextDocumentSemanticTokensFull
-semanticTokensFull recorder state pid param = do
+semanticTokensFull recorder state pid param = runActionE "SemanticTokens.semanticTokensFull" state computeSemanticTokensFull
+  where
+    computeSemanticTokensFull :: ExceptT PluginError Action (MessageResult Method_TextDocumentSemanticTokensFull)
+    computeSemanticTokensFull = do
+      nfp <- getNormalizedFilePathE (param ^. L.textDocument . L.uri)
+      items <- computeSemanticTokens recorder pid state nfp
+      lift $ setSemanticTokens nfp items
+      return $ InL items
+
+
+semanticTokensFullDelta :: Recorder (WithPriority SemanticLog) -> PluginMethodHandler IdeState 'Method_TextDocumentSemanticTokensFullDelta
+semanticTokensFullDelta recorder state pid param = do
   nfp <- getNormalizedFilePathE (param ^. L.textDocument . L.uri)
-  items <- runActionE "SemanticTokens.semanticTokensFull" state $ computeSemanticTokens recorder pid state nfp
-  return $ InL items
+  let previousVersionFromParam = param ^. L.previousResultId
+  runActionE "SemanticTokens.semanticTokensFullDelta" state $ computeSemanticTokensFullDelta recorder previousVersionFromParam  pid state nfp
+  where
+    computeSemanticTokensFullDelta :: Recorder (WithPriority SemanticLog) -> Text -> PluginId -> IdeState -> NormalizedFilePath -> ExceptT PluginError Action (MessageResult Method_TextDocumentSemanticTokensFullDelta)
+    computeSemanticTokensFullDelta recorder previousVersionFromParam  pid state nfp = do
+      semanticTokens <- computeSemanticTokens recorder pid state nfp
+      previousSemanticTokensMaybe <- lift $ getPreviousSemanticTokens nfp
+      lift $ setSemanticTokens nfp semanticTokens
+      case previousSemanticTokensMaybe of
+          Nothing -> return $ InL semanticTokens
+          Just previousSemanticTokens ->
+              if Just previousVersionFromParam == previousSemanticTokens^.L.resultId
+              then return $ InR $ InL $ makeSemanticTokensDeltaWithId (semanticTokens^.L.resultId) previousSemanticTokens semanticTokens
+              else do
+                logWith recorder Warning (LogSemanticTokensDeltaMisMatch previousVersionFromParam (previousSemanticTokens^.L.resultId))
+                return $ InL semanticTokens
 
 -- | Defines the 'getSemanticTokensRule' function, compute semantic tokens for a Haskell source file.
 --
@@ -98,9 +131,6 @@ getSemanticTokensRule recorder =
     let hsFinder = idSemantic getTyThingMap (hieKindFunMasksKind hieKind) refMap
     return $ computeRangeHsSemanticTokenTypeList hsFinder virtualFile ast
 
--- | Persistent rule to ensure that semantic tokens doesn't block on startup
-persistentGetSemanticTokensRule :: Rules ()
-persistentGetSemanticTokensRule = addPersistentRule GetSemanticTokens $ \_ -> pure $ Just (RangeHsSemanticTokenTypes mempty, idDelta, Nothing)
 
 -- taken from /haskell-language-server/plugins/hls-code-range-plugin/src/Ide/Plugin/CodeRange/Rules.hs
 
@@ -113,3 +143,22 @@ handleError recorder action' = do
       logWith recorder Warning msg
       pure $ toIdeResult (Left [])
     Right value -> pure $ toIdeResult (Right value)
+
+-----------------------
+-- helper functions
+-----------------------
+
+-- keep track of the semantic tokens response id
+-- so that we can compute the delta between two versions
+getAndIncreaseSemanticTokensId :: Action SemanticTokenId
+getAndIncreaseSemanticTokensId = do
+  ShakeExtras{semanticTokensId} <- getShakeExtras
+  liftIO $ atomically $ do
+    i <- stateTVar semanticTokensId (\val -> (val, val+1))
+    return $ T.pack $ show i
+
+getPreviousSemanticTokens :: NormalizedFilePath -> Action (Maybe SemanticTokens)
+getPreviousSemanticTokens uri = getShakeExtras >>= liftIO . atomically . STM.lookup uri . semanticTokensCache
+
+setSemanticTokens :: NormalizedFilePath -> SemanticTokens -> Action ()
+setSemanticTokens uri tokens = getShakeExtras >>= liftIO . atomically . STM.insert tokens uri . semanticTokensCache

plugins/hls-semantic-tokens-plugin/src/Ide/Plugin/SemanticTokens/Mappings.hs

@@ -1,7 +1,7 @@
 {-# LANGUAGE GADTs             #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeFamilies      #-}
-{-# LANGUAGE TypeOperators     #-}
+
 
 -- |
 -- This module provides mappings to convert token type information in the Haskell IDE plugin. It includes functions for:

plugins/hls-semantic-tokens-plugin/src/Ide/Plugin/SemanticTokens/Query.hs

@@ -1,6 +1,3 @@
-{-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE OverloadedStrings   #-}
-
 -- |
 -- The query module is used to query the semantic tokens from the AST
 module Ide.Plugin.SemanticTokens.Query where
@@ -18,13 +15,16 @@ import           Ide.Plugin.SemanticTokens.Mappings
 import           Ide.Plugin.SemanticTokens.Types      (HieFunMaskKind,
                                                        HsSemanticTokenType (TModule),
                                                        RangeSemanticTokenTypeList,
+                                                       SemanticTokenId,
                                                        SemanticTokensConfig)
 import           Language.LSP.Protocol.Types          (Position (Position),
                                                        Range (Range),
                                                        SemanticTokenAbsolute (SemanticTokenAbsolute),
-                                                       SemanticTokens,
+                                                       SemanticTokens (SemanticTokens),
+                                                       SemanticTokensDelta (SemanticTokensDelta),
                                                        defaultSemanticTokensLegend,
-                                                       makeSemanticTokens)
+                                                       makeSemanticTokens,
+                                                       makeSemanticTokensDelta)
 import           Prelude                              hiding (length, span)
 
 ---------------------------------------------------------
@@ -47,8 +47,7 @@ idSemantic tyThingMap hieKind rm (Right n) =
 ---------------------------------------------------------
 
 nameSemanticFromHie :: forall a. HieFunMaskKind a -> RefMap a -> Name -> Maybe HsSemanticTokenType
-nameSemanticFromHie hieKind rm n = do
-  idSemanticFromRefMap rm (Right n)
+nameSemanticFromHie hieKind rm n = idSemanticFromRefMap rm (Right n)
   where
     idSemanticFromRefMap :: RefMap a -> Identifier -> Maybe HsSemanticTokenType
     idSemanticFromRefMap rm' name' = do
@@ -67,10 +66,9 @@ nameSemanticFromHie hieKind rm n = do
 
 -------------------------------------------------
 
-rangeSemanticsSemanticTokens :: SemanticTokensConfig -> PositionMapping -> RangeSemanticTokenTypeList -> Either Text SemanticTokens
-rangeSemanticsSemanticTokens stc mapping =
-  makeSemanticTokens defaultSemanticTokensLegend
-    . mapMaybe (\(ran, tk) -> toAbsSemanticToken <$> toCurrentRange mapping ran <*> return tk)
+rangeSemanticsSemanticTokens :: SemanticTokenId -> SemanticTokensConfig -> PositionMapping -> RangeSemanticTokenTypeList -> Either Text SemanticTokens
+rangeSemanticsSemanticTokens sid stc mapping =
+  makeSemanticTokensWithId (Just sid) . mapMaybe (\(ran, tk) -> toAbsSemanticToken <$> toCurrentRange mapping ran <*> return tk)
   where
     toAbsSemanticToken :: Range -> HsSemanticTokenType -> SemanticTokenAbsolute
     toAbsSemanticToken (Range (Position startLine startColumn) (Position _endLine endColumn)) tokenType =
@@ -81,3 +79,14 @@ rangeSemanticsSemanticTokens stc mapping =
             (fromIntegral len)
             (toLspTokenType stc tokenType)
             []
+
+makeSemanticTokensWithId :: Maybe SemanticTokenId -> [SemanticTokenAbsolute] -> Either Text SemanticTokens
+makeSemanticTokensWithId sid tokens = do
+    (SemanticTokens _  tokens) <- makeSemanticTokens defaultSemanticTokensLegend tokens
+    return $ SemanticTokens sid tokens
+
+makeSemanticTokensDeltaWithId :: Maybe SemanticTokenId ->  SemanticTokens -> SemanticTokens -> SemanticTokensDelta
+makeSemanticTokensDeltaWithId sid previousTokens currentTokens =
+    let (SemanticTokensDelta _ stEdits) = makeSemanticTokensDelta previousTokens currentTokens
+    in SemanticTokensDelta sid stEdits
+

plugins/hls-semantic-tokens-plugin/src/Ide/Plugin/SemanticTokens/Types.hs

@@ -18,6 +18,7 @@ import           Development.IDE.Graph.Classes (Hashable)
 import           GHC.Generics                  (Generic)
 import           Language.LSP.Protocol.Types
 -- import template haskell
+import           Data.Text                     (Text)
 import           Language.Haskell.TH.Syntax    (Lift)
 
 
@@ -140,6 +141,7 @@ data SemanticLog
   | LogConfig SemanticTokensConfig
   | LogMsg String
   | LogNoVF
+  | LogSemanticTokensDeltaMisMatch Text (Maybe Text)
   deriving (Show)
 
 instance Pretty SemanticLog where
@@ -149,4 +151,9 @@ instance Pretty SemanticLog where
     LogNoVF           -> "no VirtualSourceFile exist for file"
     LogConfig config  -> "SemanticTokensConfig_: " <> pretty (show config)
     LogMsg msg        -> "SemanticLog Debug Message: " <> pretty msg
+    LogSemanticTokensDeltaMisMatch previousIdFromRequest previousIdFromCache
+                      -> "SemanticTokensDeltaMisMatch: previousIdFromRequest: " <> pretty previousIdFromRequest
+                      <> " previousIdFromCache: " <> pretty previousIdFromCache
 
+
+type SemanticTokenId = Text