Implement semantic tokens plugin to support semantic highlighting(textDocument/semanticTokens/full) (#3892)

* Implement semantic tokens lsp plugin draft

* SemanticTokens: combine information extracted from HieAst

* clean up

* map to default token types in lsp

* use lsp makeSemanticTokens to convert to lsp SemanticTokens type

* add test and cleanup

* refine semantic type to default one in lsp

* Use tokens from hieAst instead of renamedSource and add test

* use customize RefMap to get semantic type

* use refMap from useAsts

* Also compute imported names

* Also compute semantic type from TyThing

* Fix dependencies version

* fix version

* Retrieve nameSet from renamedSource to prevent names not visible(Such as by instance deriving) being handled

* add hlint config to ignore test data

* cean up test data

* revert flake.nix

* Rename query.hs to Query.hs

* Build: add semantic tokens to lts21

* Refactor and add README

* Semantic token, filter names in Ast

* CI: add consistancy check for wether semantic tokens computations is stable across different ghc versions

* Update documentation, cleanup test, remove default modifiers

* Fix: IO now classfied to TTypcon, add test for GADT and data family, Update documentation

* Restore stack.yaml

* fix stack build

* Refactor, move out ActualToken to Mappings and use ide logger

* Refactor: toLspTokenType should return Maybe type

* Stop use stale hieAst

* add getImportedNameSemanticRule rule to semantic tokens plugin

* do not retrieve hie in getImportedNameSemanticRule

* fix: add description for semantic tokens

* remove TValBind and TPaternBind and Use TFunction and TVariable instead

* cleanup

* Refactor useWithStaleMT and took care of the token range using position map

* fix build for 9.4

* refactor, use golden test

* refactor, use ExceptT for computeSemanticTokens

* Fix 9.2

* add persistentSemanticMapRule to prevent semantic tokens block on startup

* Fix, use hieKind instead of cast the type directly

* add options to turn semantic tokens on and off

* Disable stan plugin by default (#3917)

* Fix positionMapping in stale data (#3920)

* Fix positionMapping in stale data

* add test for updatePositionMapping

* add comment to demonstrate addOldDelta

* cleanup

* fix: for local variable, extract type from contextInfo instead of bind site, thus function in pattern binds can also be indentified

* clean up

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

Co-authored-by: Michael Peyton Jones <[email protected]>

* refactor: remove TNothing and compact the test output

* refactor: rename SemanticTokenType to HsSemanticTokenType to avoid confusion with lsp' SemanticTokenTypes

* refactor: push the computation of semantic token type to getSemanticTokensRule

* update documentation

* cleanup hieAstSpanNames

* remove renamed source from getSemanticTokensRule and optimize query function for semantic token type

* try to exclude names that is not visible in hie and cleanup

* add HieFunMaskKind, it is to differ wether a type at type index is a function or non-function

* expose function flag to expose (=>, ->, -=>, ==>)

* 1. Relax GetDocMap kindMap to get TyThing for more than type variables.
2. Backport isVisibleFunArg

* use customize logger, add test for unicode

* fix: handle unicode in semantic tokens

* update KindMap to TyThingMap

* cleanup

* add realSrcSpanToCodePointRange, realSrcLocToCodePointPosition to Development.IDE.GHC.Error

* add Note [Semantic information from Multiple Sources]

* move recoverFunMaskArray to Mappings.hs

* fix test, data.Set might not appear

* fix: handle semantic tokens with more than one ast

* fix: instance PluginMethod Request Method_TextDocumentSemanticTokensFull

* clean up

* turn semantic tokens off by default

* fix doc

* clean up doc

---------

Co-authored-by: fendor <[email protected]>
Co-authored-by: Michael Peyton Jones <[email protected]>

Author: 3 people

cabal.project

@@ -34,6 +34,7 @@ packages:
          ./plugins/hls-explicit-record-fields-plugin
          ./plugins/hls-refactor-plugin
          ./plugins/hls-overloaded-record-dot-plugin
+         ./plugins/hls-semantic-tokens-plugin
 
 index-state: 2023-12-13T00:00:00Z
 
@@ -55,8 +56,8 @@ constraints:
   text -simdutf,
   ghc-check -ghc-check-use-package-abis,
   ghc-lib-parser-ex -auto,
-  -- This is only present in some versions, and it's on by default since 
-  -- 0.14.5.0, but there are some versions we allow that need this 
+  -- This is only present in some versions, and it's on by default since
+  -- 0.14.5.0, but there are some versions we allow that need this
   -- setting
   stylish-haskell +ghc-lib,
   -- Centos 7 comes with an old gcc version that doesn't know about
@@ -79,8 +80,8 @@ source-repository-package
 -- END DELETE
 
 if impl(ghc >= 9.1)
-  -- ekg packagess are old and unmaintained, but we 
-  -- don't rely on them for the mainline build, so 
+  -- ekg packagess are old and unmaintained, but we
+  -- don't rely on them for the mainline build, so
   -- this is okay
   allow-newer:
     ekg-json:base,

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

@@ -238,14 +238,14 @@ type instance RuleResult GetHieAst = HieAstResult
 -- | A IntervalMap telling us what is in scope at each point
 type instance RuleResult GetBindings = Bindings
 
-data DocAndKindMap = DKMap {getDocMap :: !DocMap, getKindMap :: !KindMap}
-instance NFData DocAndKindMap where
+data DocAndTyThingMap = DKMap {getDocMap :: !DocMap, getTyThingMap :: !TyThingMap}
+instance NFData DocAndTyThingMap where
     rnf (DKMap a b) = rwhnf a `seq` rwhnf b
 
-instance Show DocAndKindMap where
+instance Show DocAndTyThingMap where
     show = const "docmap"
 
-type instance RuleResult GetDocMap = DocAndKindMap
+type instance RuleResult GetDocMap = DocAndTyThingMap
 
 -- | A GHC session that we reuse.
 type instance RuleResult GhcSession = HscEnvEq

ghcide/src/Development/IDE/GHC/Compat/Core.hs

@@ -407,6 +407,7 @@ module Development.IDE.GHC.Compat.Core (
     field_label,
 #endif
     groupOrigin,
+    isVisibleFunArg,
     ) where
 
 import qualified GHC
@@ -431,13 +432,13 @@ import           GHC.Core.DataCon             hiding (dataConExTyCoVars)
 import qualified GHC.Core.DataCon             as DataCon
 import           GHC.Core.FamInstEnv          hiding (pprFamInst)
 import           GHC.Core.InstEnv
-import           GHC.Types.Unique.FM 
+import           GHC.Types.Unique.FM
 import           GHC.Core.PatSyn
 import           GHC.Core.Predicate
 import           GHC.Core.TyCo.Ppr
 import qualified GHC.Core.TyCo.Rep            as TyCoRep
 import           GHC.Core.TyCon
-import           GHC.Core.Type                
+import           GHC.Core.Type
 import           GHC.Core.Unify
 import           GHC.Core.Utils
 import           GHC.Driver.CmdLine           (Warn (..))
@@ -489,6 +490,8 @@ import qualified GHC.Types.SrcLoc             as SrcLoc
 import           GHC.Types.Unique.Supply
 import           GHC.Types.Var                (Var (varName), setTyVarUnique,
                                                setVarUnique)
+
+import qualified GHC.Types.Var                as TypesVar
 import           GHC.Unit.Info                (PackageName (..))
 import           GHC.Unit.Module              hiding (ModLocation (..), UnitId,
                                                moduleUnit,
@@ -597,7 +600,7 @@ pattern RealSrcLoc x y = SrcLoc.RealSrcLoc x y
 pattern AvailTC :: Name -> [Name] -> [FieldLabel] -> Avail.AvailInfo
 #if __GLASGOW_HASKELL__ >= 907
 pattern AvailTC n names pieces <- Avail.AvailTC n ((,[]) -> (names,pieces))
-#else 
+#else
 pattern AvailTC n names pieces <- Avail.AvailTC n ((\gres -> foldr (\gre (names, pieces) -> case gre of
       Avail.NormalGreName name -> (name: names, pieces)
       Avail.FieldGreName label -> (names, label:pieces)) ([], []) gres) -> (names, pieces))
@@ -606,14 +609,14 @@ pattern AvailTC n names pieces <- Avail.AvailTC n ((\gres -> foldr (\gre (names,
 pattern AvailName :: Name -> Avail.AvailInfo
 #if __GLASGOW_HASKELL__ >= 907
 pattern AvailName n <- Avail.Avail n
-#else 
+#else
 pattern AvailName n <- Avail.Avail (Avail.NormalGreName n)
 #endif
 
 pattern AvailFL :: FieldLabel -> Avail.AvailInfo
 #if __GLASGOW_HASKELL__ >= 907
 pattern AvailFL fl <- (const Nothing -> Just fl) -- this pattern always fails as this field was removed in 9.7
-#else 
+#else
 pattern AvailFL fl <- Avail.Avail (Avail.FieldGreName fl)
 #endif
 
@@ -630,8 +633,17 @@ pattern ExposePackage s a mr <- DynFlags.ExposePackage s a _ mr
 pattern ExposePackage s a mr = DynFlags.ExposePackage s a mr
 #endif
 
-pattern FunTy :: Type -> Type -> Type
-pattern FunTy arg res <- TyCoRep.FunTy {ft_arg = arg, ft_res = res}
+#if __GLASGOW_HASKELL__ >= 906
+isVisibleFunArg = TypesVar.isVisibleFunArg
+type FunTyFlag = TypesVar.FunTyFlag
+#else
+isVisibleFunArg VisArg = True
+isVisibleFunArg _ = False
+type FunTyFlag = TypesVar.AnonArgFlag
+#endif
+pattern FunTy :: Development.IDE.GHC.Compat.Core.FunTyFlag -> Type -> Type -> Type
+pattern FunTy af arg res <- TyCoRep.FunTy {ft_af = af, ft_arg = arg, ft_res = res}
+
 
 -- type HasSrcSpan x a = (GenLocated SrcSpan a ~ x)
 -- type HasSrcSpan x = () :: Constraint

ghcide/src/Development/IDE/GHC/Error.hs

@@ -17,6 +17,8 @@ module Development.IDE.GHC.Error
   , realSrcSpanToRange
   , realSrcLocToPosition
   , realSrcSpanToLocation
+  , realSrcSpanToCodePointRange
+  , realSrcLocToCodePointPosition
   , srcSpanToFilename
   , rangeToSrcSpan
   , rangeToRealSrcSpan
@@ -45,6 +47,8 @@ import           Development.IDE.Types.Diagnostics as D
 import           Development.IDE.Types.Location
 import           GHC
 import           Language.LSP.Protocol.Types       (isSubrangeOf)
+import           Language.LSP.VFS                  (CodePointPosition (CodePointPosition),
+                                                    CodePointRange (CodePointRange))
 
 
 diagFromText :: T.Text -> D.DiagnosticSeverity -> SrcSpan -> T.Text -> FileDiagnostic
@@ -86,6 +90,28 @@ realSrcLocToPosition :: RealSrcLoc -> Position
 realSrcLocToPosition real =
   Position (fromIntegral $ srcLocLine real - 1) (fromIntegral $ srcLocCol real - 1)
 
+-- Note [Unicode support]
+-- the current situation is:
+-- LSP Positions use UTF-16 code units(Unicode may count as variable columns);
+-- GHC use Unicode code points(Unicode count as one column).
+-- To support unicode, ideally range should be in lsp standard,
+-- and codePoint should be in ghc standard.
+-- see https://github.com/haskell/lsp/pull/407
+
+-- | Convert a GHC SrcSpan to CodePointRange
+-- see Note [Unicode support]
+realSrcSpanToCodePointRange :: RealSrcSpan -> CodePointRange
+realSrcSpanToCodePointRange real =
+  CodePointRange
+    (realSrcLocToCodePointPosition $ Compat.realSrcSpanStart real)
+    (realSrcLocToCodePointPosition $ Compat.realSrcSpanEnd real)
+
+-- | Convert a GHC RealSrcLoc to CodePointPosition
+-- see Note [Unicode support]
+realSrcLocToCodePointPosition :: RealSrcLoc -> CodePointPosition
+realSrcLocToCodePointPosition real =
+  CodePointPosition (fromIntegral $ srcLocLine real - 1) (fromIntegral $ srcLocCol real - 1)
+
 -- | Extract a file name from a GHC SrcSpan (use message for unhelpful ones)
 -- FIXME This may not be an _absolute_ file name, needs fixing.
 srcSpanToFilename :: SrcSpan -> Maybe FilePath

ghcide/src/Development/IDE/Plugin/Completions.hs

@@ -144,8 +144,8 @@ resolveCompletion ide _pid comp@CompletionItem{_detail,_documentation,_data_} ur
 #endif
     mdkm <- liftIO $ runIdeAction "CompletionResolve.GetDocMap" (shakeExtras ide) $ useWithStaleFast GetDocMap file
     let (dm,km) = case mdkm of
-          Just (DKMap docMap kindMap, _) -> (docMap,kindMap)
-          Nothing                        -> (mempty, mempty)
+          Just (DKMap docMap tyThingMap, _) -> (docMap,tyThingMap)
+          Nothing                           -> (mempty, mempty)
     doc <- case lookupNameEnv dm name of
       Just doc -> pure $ spanDocToMarkdown doc
       Nothing -> liftIO $ spanDocToMarkdown <$> getDocumentationTryGhc (hscEnv sess) name

ghcide/src/Development/IDE/Spans/AtPoint.hs

@@ -209,7 +209,7 @@ gotoDefinition withHieDb getHieFile ideOpts imports srcSpans pos
 atPoint
   :: IdeOptions
   -> HieAstResult
-  -> DocAndKindMap
+  -> DocAndTyThingMap
   -> HscEnv
   -> Position
   -> IO (Maybe (Maybe Range, [T.Text]))
@@ -346,7 +346,7 @@ namesInType (TyVarTy n)      = [varName n]
 namesInType (AppTy a b)      = getTypes [a,b]
 namesInType (TyConApp tc ts) = tyConName tc : getTypes ts
 namesInType (ForAllTy b t)   = varName (binderVar b) : namesInType t
-namesInType (FunTy a b)      = getTypes [a,b]
+namesInType (FunTy _ a b)    = getTypes [a,b]
 namesInType (CastTy t _)     = namesInType t
 namesInType (LitTy _)        = []
 namesInType _                = []

ghcide/src/Development/IDE/Spans/Common.hs

@@ -12,7 +12,7 @@ module Development.IDE.Spans.Common (
 , spanDocToMarkdown
 , spanDocToMarkdownForTest
 , DocMap
-, KindMap
+, TyThingMap
 ) where
 
 import           Control.DeepSeq
@@ -31,7 +31,7 @@ import qualified Documentation.Haddock.Parser as H
 import qualified Documentation.Haddock.Types  as H
 
 type DocMap = NameEnv SpanDoc
-type KindMap = NameEnv TyThing
+type TyThingMap = NameEnv TyThing
 
 -- | Shows IEWrappedName, without any modifier, qualifier or unique identifier.
 #if MIN_VERSION_ghc(9,5,0)

ghcide/src/Development/IDE/Spans/Documentation.hs

@@ -39,7 +39,7 @@ mkDocMap
   :: HscEnv
   -> RefMap a
   -> TcGblEnv
-  -> IO DocAndKindMap
+  -> IO DocAndTyThingMap
 mkDocMap env rm this_mod =
   do
 #if MIN_VERSION_ghc(9,3,0)
@@ -61,8 +61,7 @@ mkDocMap env rm this_mod =
       doc <- getDocumentationTryGhc env n
       pure $ extendNameEnv nameMap n doc
     getType n nameMap
-      | isTcOcc $ occName n
-      , Nothing <- lookupNameEnv nameMap n
+      | Nothing <- lookupNameEnv nameMap n
       = do kind <- lookupKind env n
            pure $ maybe nameMap (extendNameEnv nameMap n) kind
       | otherwise = pure nameMap

haskell-language-server.cabal

@@ -164,6 +164,11 @@ flag overloadedRecordDot
   default:     True
   manual:      True
 
+flag semanticTokens
+  description: Enable semantic tokens plugin
+  default:     True
+  manual:      True
+
 -- formatters
 
 flag floskell
@@ -333,6 +338,12 @@ common refactor
     build-depends: hls-refactor-plugin == 2.5.0.0
     cpp-options: -Dhls_refactor
 
+common semanticTokens
+  if flag(semanticTokens)
+    build-depends: hls-semantic-tokens-plugin == 2.5.0.0
+    cpp-options: -Dhls_semanticTokens
+
+
 library
   import:           common-deps
                     -- configuration
@@ -365,6 +376,7 @@ library
                   , stylishHaskell
                   , refactor
                   , overloadedRecordDot
+                  , semanticTokens
 
   exposed-modules:
     Ide.Arguments

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

@@ -66,6 +66,7 @@ parsePluginConfig :: PluginConfig -> Value -> A.Parser PluginConfig
 parsePluginConfig def = A.withObject "PluginConfig" $ \o  -> PluginConfig
       <$> o .:? "globalOn"         .!= plcGlobalOn def
       <*> o .:? "callHierarchyOn"  .!= plcCallHierarchyOn def
+      <*> o .:? "semanticTokensOn" .!= plcSemanticTokensOn def
       <*> o .:? "codeActionsOn"    .!= plcCodeActionsOn def
       <*> o .:? "codeLensOn"       .!= plcCodeLensOn    def
       <*> o .:? "diagnosticsOn"    .!= plcDiagnosticsOn def -- AZ

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

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

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

@@ -230,6 +230,7 @@ data PluginConfig =
       , plcRenameOn         :: !Bool
       , plcSelectionRangeOn :: !Bool
       , plcFoldingRangeOn   :: !Bool
+      , plcSemanticTokensOn :: !Bool
       , plcConfig           :: !Object
       } deriving (Show,Eq)
 
@@ -246,11 +247,12 @@ instance Default PluginConfig where
       , plcRenameOn         = True
       , plcSelectionRangeOn = True
       , plcFoldingRangeOn   = True
+      , plcSemanticTokensOn = True
       , plcConfig           = mempty
       }
 
 instance ToJSON PluginConfig where
-    toJSON (PluginConfig g ch ca cl d h s c rn sr fr cfg) = r
+    toJSON (PluginConfig g ch ca cl d h s c rn sr fr st cfg) = r
       where
         r = object [ "globalOn"         .= g
                    , "callHierarchyOn"  .= ch
@@ -263,6 +265,7 @@ instance ToJSON PluginConfig where
                    , "renameOn"         .= rn
                    , "selectionRangeOn" .= sr
                    , "foldingRangeOn"   .= fr
+                   , "semanticTokensOn" .= st
                    , "config"           .= cfg
                    ]
 
@@ -514,6 +517,9 @@ instance PluginMethod Request Method_TextDocumentRangeFormatting where
     where
       pid = pluginId pluginDesc
 
+instance PluginMethod Request Method_TextDocumentSemanticTokensFull where
+  handlesRequest = pluginEnabledWithFeature plcSemanticTokensOn
+
 instance PluginMethod Request Method_TextDocumentPrepareCallHierarchy where
   handlesRequest = pluginEnabledWithFeature plcCallHierarchyOn
 
@@ -751,6 +757,9 @@ instance PluginRequestMethod Method_CallHierarchyOutgoingCalls where
 instance PluginRequestMethod (Method_CustomMethod m) where
   combineResponses _ _ _ _ (x :| _) = x
 
+instance PluginRequestMethod Method_TextDocumentSemanticTokensFull where
+  combineResponses _ _ _ _ (x :| _) = x
+
 takeLefts :: [a |? b] -> [a]
 takeLefts = mapMaybe (\x -> [res | (InL res) <- Just x])
 

plugins/hls-semantic-tokens-plugin/.hlint.yaml

@@ -0,0 +1 @@
+- ignore: { "within": 'test/testdata/*.hs' }