Semantic tokens: add module name support and improve performance and accuracy by traversing the hieAst along with source code

haskell-language-server.cabal

@@ -140,7 +140,7 @@ flag cabal
 
 common cabal
   if flag(cabal)
-    build-depends: hls-cabal-plugin 
+    build-depends: hls-cabal-plugin
     cpp-options: -Dhls_cabal
 
 library hls-cabal-plugin
@@ -223,7 +223,7 @@ flag class
 
 common class
   if flag(class)
-    build-depends: hls-class-plugin 
+    build-depends: hls-class-plugin
     cpp-options: -Dhls_class
 
 library hls-class-plugin
@@ -287,7 +287,7 @@ flag callHierarchy
 
 common callHierarchy
   if flag(callHierarchy)
-    build-depends: hls-call-hierarchy-plugin 
+    build-depends: hls-call-hierarchy-plugin
     cpp-options: -Dhls_callHierarchy
 
 library hls-call-hierarchy-plugin
@@ -348,7 +348,7 @@ flag eval
 
 common eval
   if flag(eval)
-    build-depends: hls-eval-plugin 
+    build-depends: hls-eval-plugin
     cpp-options: -Dhls_eval
 
 library hls-eval-plugin
@@ -429,7 +429,7 @@ test-suite hls-eval-plugin-tests
 
 common importLens
   if flag(importLens)
-    build-depends: hls-explicit-imports-plugin 
+    build-depends: hls-explicit-imports-plugin
     cpp-options: -Dhls_importLens
 
 flag importLens
@@ -494,7 +494,7 @@ flag rename
 
 common rename
   if flag(rename)
-    build-depends: hls-rename-plugin 
+    build-depends: hls-rename-plugin
     cpp-options: -Dhls_rename
 
 library hls-rename-plugin
@@ -550,7 +550,7 @@ flag retrie
 
 common retrie
   if flag(retrie)
-    build-depends: hls-retrie-plugin 
+    build-depends: hls-retrie-plugin
     cpp-options: -Dhls_retrie
 
 library hls-retrie-plugin
@@ -615,7 +615,7 @@ flag hlint
 
 common hlint
   if flag(hlint) && (impl(ghc < 9.8.0) || flag(ignore-plugins-ghc-bounds))
-    build-depends: hls-hlint-plugin 
+    build-depends: hls-hlint-plugin
     cpp-options: -Dhls_hlint
 
 library hls-hlint-plugin
@@ -695,7 +695,7 @@ flag stan
 
 common stan
   if flag(stan) && (impl(ghc > 8.8.1) && impl(ghc <= 9.2.3) || impl(ghc >= 9.4.0) && impl(ghc < 9.10.0))
-    build-depends: hls-stan-plugin 
+    build-depends: hls-stan-plugin
     cpp-options: -Dhls_stan
 
 library hls-stan-plugin
@@ -769,7 +769,7 @@ flag moduleName
 
 common moduleName
   if flag(moduleName)
-    build-depends: hls-module-name-plugin 
+    build-depends: hls-module-name-plugin
     cpp-options: -Dhls_moduleName
 
 library hls-module-name-plugin
@@ -814,7 +814,7 @@ flag pragmas
 
 common pragmas
   if flag(pragmas)
-    build-depends: hls-pragmas-plugin 
+    build-depends: hls-pragmas-plugin
     cpp-options: -Dhls_pragmas
 
 library hls-pragmas-plugin
@@ -862,7 +862,7 @@ flag splice
 
 common splice
   if flag(splice)
-    build-depends: hls-splice-plugin 
+    build-depends: hls-splice-plugin
     cpp-options: -Dhls_splice
 
 library hls-splice-plugin
@@ -1040,7 +1040,7 @@ flag codeRange
 
 common codeRange
   if flag(codeRange)
-    build-depends: hls-code-range-plugin 
+    build-depends: hls-code-range-plugin
     cpp-options: -Dhls_codeRange
 
 library hls-code-range-plugin
@@ -1100,7 +1100,7 @@ flag changeTypeSignature
 
 common changeTypeSignature
   if flag(changeTypeSignature)
-    build-depends: hls-change-type-signature-plugin 
+    build-depends: hls-change-type-signature-plugin
     cpp-options: -Dhls_changeTypeSignature
 
 library hls-change-type-signature-plugin
@@ -1160,7 +1160,7 @@ flag gadt
 
 common gadt
   if flag(gadt)
-    build-depends: hls-gadt-plugin 
+    build-depends: hls-gadt-plugin
     cpp-options: -Dhls_gadt
 
 library hls-gadt-plugin
@@ -1213,7 +1213,7 @@ flag explicitFixity
 
 common explicitFixity
   if flag(explicitFixity)
-    build-depends: hls-explicit-fixity-plugin 
+    build-depends: hls-explicit-fixity-plugin
     cpp-options: -DexplicitFixity
 
 library hls-explicit-fixity-plugin
@@ -1260,7 +1260,7 @@ flag explicitFields
 
 common explicitFields
   if flag(explicitFields)
-    build-depends: hls-explicit-record-fields-plugin 
+    build-depends: hls-explicit-record-fields-plugin
     cpp-options: -DexplicitFields
 
 library hls-explicit-record-fields-plugin
@@ -1284,7 +1284,7 @@ library hls-explicit-record-fields-plugin
   if flag(pedantic)
     ghc-options: -Werror
                  -Wwarn=incomplete-record-updates
-    
+
 test-suite hls-explicit-record-fields-plugin-tests
     import:           warnings
     default-language: Haskell2010
@@ -1309,7 +1309,7 @@ flag overloadedRecordDot
 
 common overloadedRecordDot
   if flag(overloadedRecordDot)
-    build-depends: hls-overloaded-record-dot-plugin 
+    build-depends: hls-overloaded-record-dot-plugin
     cpp-options: -Dhls_overloaded_record_dot
 
 library hls-overloaded-record-dot-plugin
@@ -1356,7 +1356,7 @@ flag floskell
 
 common floskell
   if flag(floskell) && (impl(ghc < 9.7) || flag(ignore-plugins-ghc-bounds))
-    build-depends: hls-floskell-plugin 
+    build-depends: hls-floskell-plugin
     cpp-options: -Dhls_floskell
 
 library hls-floskell-plugin
@@ -1398,7 +1398,7 @@ flag fourmolu
 
 common fourmolu
   if flag(fourmolu)
-    build-depends: hls-fourmolu-plugin 
+    build-depends: hls-fourmolu-plugin
     cpp-options: -Dhls_fourmolu
 
 library hls-fourmolu-plugin
@@ -1451,7 +1451,7 @@ flag ormolu
 
 common ormolu
   if flag(ormolu)
-    build-depends: hls-ormolu-plugin 
+    build-depends: hls-ormolu-plugin
     cpp-options: -Dhls_ormolu
 
 library hls-ormolu-plugin
@@ -1504,7 +1504,7 @@ flag stylishHaskell
 
 common stylishHaskell
   if flag(stylishHaskell) && (impl(ghc < 9.8.0) || flag(ignore-plugins-ghc-bounds))
-    build-depends: hls-stylish-haskell-plugin 
+    build-depends: hls-stylish-haskell-plugin
     cpp-options: -Dhls_stylishHaskell
 
 library hls-stylish-haskell-plugin
@@ -1549,7 +1549,7 @@ flag refactor
 
 common refactor
   if flag(refactor)
-    build-depends: hls-refactor-plugin 
+    build-depends: hls-refactor-plugin
     cpp-options: -Dhls_refactor
 
 library hls-refactor-plugin
@@ -1665,7 +1665,7 @@ flag semanticTokens
 
 common semanticTokens
   if flag(semanticTokens)
-    build-depends: hls-semantic-tokens-plugin 
+    build-depends: hls-semantic-tokens-plugin
     cpp-options: -Dhls_semanticTokens
 
 library hls-semantic-tokens-plugin
@@ -1679,6 +1679,7 @@ library hls-semantic-tokens-plugin
     Ide.Plugin.SemanticTokens.Query
     Ide.Plugin.SemanticTokens.SemanticConfig
     Ide.Plugin.SemanticTokens.Utils
+    Ide.Plugin.SemanticTokens.Tokenize
     Ide.Plugin.SemanticTokens.Internal
 
   hs-source-dirs:     plugins/hls-semantic-tokens-plugin/src
@@ -1688,6 +1689,7 @@ library hls-semantic-tokens-plugin
     , containers
     , extra
     , hiedb
+    , text-rope
     , mtl                   >= 2.2
     , ghcide                == 2.6.0.0
     , hls-plugin-api        == 2.6.0.0

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

@@ -15,6 +15,7 @@
 module Ide.Plugin.Properties
   ( PropertyType (..),
     ToHsType,
+    NotElem,
     MetaData (..),
     PropertyKey (..),
     SPropertyKey (..),

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

@@ -5,7 +5,9 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE NamedFieldPuns        #-}
 {-# LANGUAGE OverloadedLabels      #-}
+{-# LANGUAGE OverloadedRecordDot   #-}
 {-# LANGUAGE OverloadedStrings     #-}
+{-# LANGUAGE PolyKinds             #-}
 {-# LANGUAGE RecordWildCards       #-}
 {-# LANGUAGE ScopedTypeVariables   #-}
 {-# LANGUAGE TemplateHaskell       #-}
@@ -21,8 +23,9 @@ import           Control.Monad.Except                     (ExceptT, liftEither,
                                                            withExceptT)
 import           Control.Monad.Trans                      (lift)
 import           Control.Monad.Trans.Except               (runExceptT)
-import           Data.Aeson                               (ToJSON (toJSON))
-import qualified Data.Map                                 as Map
+import           Data.Map.Strict                          (Map)
+import qualified Data.Map.Strict                          as M
+import qualified Data.Set                                 as S
 import           Development.IDE                          (Action,
                                                            GetDocMap (GetDocMap),
                                                            GetHieAst (GetHieAst),
@@ -34,7 +37,6 @@ import           Development.IDE                          (Action,
                                                            cmapWithPrio, define,
                                                            fromNormalizedFilePath,
                                                            hieKind, logPriority,
-                                                           usePropertyAction,
                                                            use_)
 import           Development.IDE.Core.PluginUtils         (runActionE,
                                                            useWithStaleE)
@@ -54,6 +56,7 @@ import           Ide.Plugin.Error                         (PluginError (PluginIn
 import           Ide.Plugin.SemanticTokens.Mappings
 import           Ide.Plugin.SemanticTokens.Query
 import           Ide.Plugin.SemanticTokens.SemanticConfig (mkSemanticConfigFunctions)
+import           Ide.Plugin.SemanticTokens.Tokenize       (hieAstSpanIdentifiers)
 import           Ide.Plugin.SemanticTokens.Types
 import           Ide.Types
 import qualified Language.LSP.Protocol.Lens               as L
@@ -91,37 +94,36 @@ semanticTokensFull recorder state pid param = do
 -- Local names token type from 'hieAst'
 -- Name locations from 'hieAst'
 -- Visible names from 'tmrRenamed'
+
 --
 -- It then combines this information to compute the semantic tokens for the file.
 getSemanticTokensRule :: Recorder (WithPriority SemanticLog) -> Rules ()
 getSemanticTokensRule recorder =
   define (cmapWithPrio LogShake recorder) $ \GetSemanticTokens nfp -> handleError recorder $ do
     (HAR {..}) <- lift $ use_ GetHieAst nfp
     (DKMap {getTyThingMap}, _) <- lift $ useWithStale_ GetDocMap nfp
-    ast <- handleMaybe (LogNoAST $ show nfp) $ getAsts hieAst Map.!? (HiePath . mkFastString . fromNormalizedFilePath) nfp
+    ast <- handleMaybe (LogNoAST $ show nfp) $ getAsts hieAst M.!? (HiePath . mkFastString . fromNormalizedFilePath) nfp
     virtualFile <- handleMaybeM LogNoVF $ getVirtualFile nfp
     -- get current location from the old ones
-    let spanNamesMap = hieAstSpanNames virtualFile ast
-    let names = nameSetElemsStable $ unionNameSets $ Map.elems spanNamesMap
-    let localSemanticMap = mkLocalNameSemanticFromAst names (hieKindFunMasksKind hieKind) refMap
+    let spanIdMap = M.filter (not . null) $ hieAstSpanIdentifiers virtualFile ast
+    let names = S.unions $ M.elems spanIdMap
+    let localSemanticMap = mkLocalIdSemanticFromAst names (hieKindFunMasksKind hieKind) refMap
     -- get imported name semantic map
-    let importedNameSemanticMap = foldr (getTypeExclude localSemanticMap getTyThingMap) emptyNameEnv names
-    let sMap = plusNameEnv_C (<>) importedNameSemanticMap localSemanticMap
-    let rangeTokenType = extractSemanticTokensFromNames sMap spanNamesMap
+    let importedIdSemanticMap = M.mapMaybe id
+            $ M.fromSet (getTypeThing getTyThingMap) (names `S.difference` M.keysSet localSemanticMap)
+    let sMap = M.unionWith (<>) importedIdSemanticMap localSemanticMap
+    let rangeTokenType = extractSemanticTokensFromNames sMap spanIdMap
     return $ RangeHsSemanticTokenTypes rangeTokenType
   where
-    -- ignore one already in discovered in local
-    getTypeExclude ::
-      NameEnv a ->
+    getTypeThing ::
       NameEnv TyThing ->
-      Name ->
-      NameEnv HsSemanticTokenType ->
-      NameEnv HsSemanticTokenType
-    getTypeExclude localEnv tyThingMap n nameMap
-      | n `elemNameEnv` localEnv = nameMap
-      | otherwise =
-          let tyThing = lookupNameEnv tyThingMap n
-           in maybe nameMap (extendNameEnv nameMap n) (tyThing >>= tyThingSemantic)
+      Identifier ->
+      Maybe HsSemanticTokenType
+    getTypeThing tyThingMap n
+      | (Right name) <- n =
+          let tyThing = lookupNameEnv tyThingMap name
+           in (tyThing >>= tyThingSemantic)
+     | otherwise = Nothing
 
 -- | Persistent rule to ensure that semantic tokens doesn't block on startup
 persistentGetSemanticTokensRule :: Rules ()

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

@@ -13,7 +13,7 @@ module Ide.Plugin.SemanticTokens.Mappings where
 
 import qualified Data.Array                      as A
 import           Data.List.Extra                 (chunksOf, (!?))
-import qualified Data.Map                        as Map
+import qualified Data.Map.Strict                 as Map
 import           Data.Maybe                      (mapMaybe)
 import qualified Data.Set                        as Set
 import           Data.Text                       (Text, unpack)
@@ -45,6 +45,7 @@ toLspTokenType conf tk = case tk of
   TTypeFamily      -> stTypeFamily conf
   TRecordField     -> stRecordField conf
   TPatternSynonym  -> stPatternSynonym conf
+  TModule          -> stModule conf
 
 lspTokenReverseMap :: SemanticTokensConfig -> Map.Map SemanticTokenTypes HsSemanticTokenType
 lspTokenReverseMap config
@@ -114,15 +115,15 @@ recoverFunMaskArray flattened = unflattened
     -- The recursion in 'unflattened' is crucial - it's what gives us sharing
     -- function indicator check.
     unflattened :: A.Array TypeIndex Bool
-    unflattened = fmap (\flatTy -> go (fmap (unflattened A.!) flatTy)) flattened
+    unflattened = fmap (go . fmap (unflattened A.!)) flattened
 
-    -- Unfold an 'HieType' whose subterms have already been unfolded
+    -- Unfold an 'HieType' whose sub-terms have already been unfolded
     go :: HieType Bool -> Bool
     go (HTyVarTy _name)              = False
     go (HAppTy _f _x)                = False
     go (HLitTy _lit)                 = False
     go (HForAllTy ((_n, _k), _af) b) = b
-    go (HFunTy _ _ _)                = True
+    go (HFunTy {})                   = True
     go (HQualTy _constraint b)       = b
     go (HCastTy b)                   = b
     go HCoercionTy                   = False