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

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/hls-semantic-tokens-plugin.cabal

@@ -20,7 +20,7 @@ source-repository head
     location: https://github.com/haskell/haskell-language-server.git
 
 library
-  ghc-options: -Wall
+  ghc-options: -Wall  -Wincomplete-patterns -Werror=incomplete-patterns
   buildable: True
   exposed-modules:
     Ide.Plugin.SemanticTokens
@@ -30,13 +30,15 @@ library
     Ide.Plugin.SemanticTokens.Query
     Ide.Plugin.SemanticTokens.SemanticConfig
     Ide.Plugin.SemanticTokens.Utils
+    Ide.Plugin.SemanticTokens.Tokenize
     Ide.Plugin.SemanticTokens.Internal
 
   hs-source-dirs:     src
   build-depends:
     , aeson
     , base                  >=4.12 && <5
     , containers
+    , text-rope
     , extra
     , hiedb
     , mtl                   >= 2.2
@@ -61,7 +63,7 @@ library
 
 test-suite tests
   type:             exitcode-stdio-1.0
-  ghc-options: -Wall
+  ghc-options: -Wall  -Wincomplete-patterns -Werror=incomplete-patterns
   default-language: Haskell2010
   hs-source-dirs:   test
   main-is:          Main.hs

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,37 @@ 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 (getTypeExclude 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 ->
       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
+    getTypeExclude 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

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

@@ -1,106 +1,75 @@
 {-# LANGUAGE ExplicitNamespaces  #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE OverloadedStrings   #-}
 {-# LANGUAGE RankNTypes          #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
 -- |
 -- The query module is used to query the semantic tokens from the AST
 module Ide.Plugin.SemanticTokens.Query where
 
-import           Data.Either                          (rights)
 import           Data.Foldable                        (fold)
-import qualified Data.Map                             as M
-import qualified Data.Map                             as Map
-import           Data.Maybe                           (fromMaybe, listToMaybe,
-                                                       mapMaybe)
-import qualified Data.Set                             as S
+import qualified Data.Map.Strict                      as M
+import           Data.Maybe                           (listToMaybe, mapMaybe)
+import           Data.Set                             (Set)
 import qualified Data.Set                             as Set
 import           Data.Text                            (Text)
 import           Development.IDE.Core.PositionMapping (PositionMapping,
                                                        toCurrentRange)
 import           Development.IDE.GHC.Compat
-import           Development.IDE.GHC.Error            (realSrcSpanToCodePointRange)
 import           Ide.Plugin.SemanticTokens.Mappings
 import           Ide.Plugin.SemanticTokens.Types      (HieFunMaskKind,
-                                                       HsSemanticTokenType,
-                                                       NameSemanticMap,
+                                                       HsSemanticTokenType (TModule),
+                                                       IdSemanticMap,
+                                                       RangeIdSetMap,
                                                        SemanticTokensConfig)
-import           Language.LSP.Protocol.Types
-import           Language.LSP.VFS                     (VirtualFile,
-                                                       codePointRangeToRange)
-import           Prelude                              hiding (span)
+import           Language.LSP.Protocol.Types          (Position (Position),
+                                                       Range (Range),
+                                                       SemanticTokenAbsolute (SemanticTokenAbsolute),
+                                                       SemanticTokens,
+                                                       defaultSemanticTokensLegend,
+                                                       makeSemanticTokens)
+import           Prelude                              hiding (length, span)
 
 ---------------------------------------------------------
 
 -- * extract semantic map from HieAst for local variables
 
 ---------------------------------------------------------
 
-mkLocalNameSemanticFromAst :: [Name] -> HieFunMaskKind a -> RefMap a -> NameSemanticMap
-mkLocalNameSemanticFromAst names hieKind rm = mkNameEnv (mapMaybe (nameNameSemanticFromHie hieKind rm) names)
+mkLocalIdSemanticFromAst :: Set Identifier -> HieFunMaskKind a -> RefMap a -> IdSemanticMap
+mkLocalIdSemanticFromAst names hieKind rm = M.mapMaybe (idIdSemanticFromHie hieKind rm) $ M.fromSet id names
 
-nameNameSemanticFromHie :: forall a. HieFunMaskKind a -> RefMap a -> Name -> Maybe (Name, HsSemanticTokenType)
-nameNameSemanticFromHie hieKind rm ns = do
-  st <- nameSemanticFromRefMap rm ns
-  return (ns, st)
+idIdSemanticFromHie :: forall a. HieFunMaskKind a -> RefMap a -> Identifier -> Maybe HsSemanticTokenType
+idIdSemanticFromHie _ _ (Left _) = Just TModule
+idIdSemanticFromHie hieKind rm ns = do
+  idSemanticFromRefMap rm ns
   where
-    nameSemanticFromRefMap :: RefMap a -> Name -> Maybe HsSemanticTokenType
-    nameSemanticFromRefMap rm' name' = do
-      spanInfos <- Map.lookup (Right name') rm'
+    idSemanticFromRefMap :: RefMap a -> Identifier -> Maybe HsSemanticTokenType
+    idSemanticFromRefMap rm' name' = do
+      spanInfos <- M.lookup name' rm'
       let typeTokenType = foldMap (typeSemantic hieKind) $ listToMaybe $ mapMaybe (identType . snd) spanInfos
       contextInfoTokenType <- foldMap (contextInfosMaybeTokenType . identInfo . snd) spanInfos
       fold [typeTokenType, Just contextInfoTokenType]
 
     contextInfosMaybeTokenType :: Set.Set ContextInfo -> Maybe HsSemanticTokenType
     contextInfosMaybeTokenType details = foldMap infoTokenType (Set.toList details)
 
------------------------------------
-
--- * extract location from HieAST a
-
------------------------------------
-
--- | get only visible names from HieAST
--- we care only the leaf node of the AST
--- and filter out the derived and evidence names
-hieAstSpanNames :: VirtualFile -> HieAST a -> M.Map Range NameSet
-hieAstSpanNames vf ast =
-  if null (nodeChildren ast)
-    then getIds ast
-    else M.unionsWith unionNameSet $ map (hieAstSpanNames vf) (nodeChildren ast)
-  where
-    getIds ast' = fromMaybe mempty $ do
-      range <- codePointRangeToRange vf $ realSrcSpanToCodePointRange $ nodeSpan ast'
-      return $ M.singleton range (getNodeIds' ast')
-    getNodeIds' =
-      Map.foldl' combineNodeIds mempty
-        . Map.filterWithKey (\k _ -> k == SourceInfo)
-        . getSourcedNodeInfo
-        . sourcedNodeInfo
-    combineNodeIds :: NameSet -> NodeInfo a -> NameSet
-    ad `combineNodeIds` (NodeInfo _ _ bd) = ad `unionNameSet` xs
-      where
-        xs = mkNameSet $ rights $ M.keys $ M.filterWithKey inclusion bd
-        inclusion :: Identifier -> IdentifierDetails a -> Bool
-        inclusion a b = not $ exclusion a b
-        exclusion :: Identifier -> IdentifierDetails a -> Bool
-        exclusion idt IdentifierDetails {identInfo = infos} = case idt of
-          Left _  -> True
-          Right _ -> any isEvidenceContext (S.toList infos)
 
 -------------------------------------------------
 
--- * extract semantic tokens from NameSemanticMap
+-- * extract semantic tokens from IdSemanticMap
 
 -------------------------------------------------
 
-extractSemanticTokensFromNames :: NameSemanticMap -> M.Map Range NameSet -> M.Map Range HsSemanticTokenType
-extractSemanticTokensFromNames nsm = Map.mapMaybe (foldMap (lookupNameEnv nsm) . nameSetElemsStable)
+extractSemanticTokensFromNames :: IdSemanticMap -> RangeIdSetMap -> M.Map Range HsSemanticTokenType
+extractSemanticTokensFromNames nsm = M.mapMaybe (foldMap (`M.lookup` nsm))
 
 rangeSemanticMapSemanticTokens :: SemanticTokensConfig -> PositionMapping -> M.Map Range HsSemanticTokenType -> Either Text SemanticTokens
 rangeSemanticMapSemanticTokens stc mapping =
   makeSemanticTokens defaultSemanticTokensLegend
     . mapMaybe (\(range, ty) -> flip toAbsSemanticToken ty <$> range)
-    . Map.toAscList
+    . M.toAscList
     . M.mapKeys (toCurrentRange mapping)
   where
     toAbsSemanticToken :: Range -> HsSemanticTokenType -> SemanticTokenAbsolute