{-# LANGUAGE ExplicitNamespaces  #-}
{-# 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.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,
                                                       SemanticTokensConfig)
import           Language.LSP.Protocol.Types
import           Language.LSP.VFS                     (VirtualFile,
                                                       codePointRangeToRange)
import           Prelude                              hiding (span)

---------------------------------------------------------

-- * extract semantic map from HieAst for local variables

---------------------------------------------------------

mkLocalNameSemanticFromAst :: [Name] -> HieFunMaskKind a -> RefMap a -> NameSemanticMap
mkLocalNameSemanticFromAst :: forall a. [Name] -> HieFunMaskKind a -> RefMap a -> NameSemanticMap
mkLocalNameSemanticFromAst [Name]
names HieFunMaskKind a
hieKind RefMap a
rm = [(Name, HsSemanticTokenType)] -> NameSemanticMap
forall a. [(Name, a)] -> NameEnv a
mkNameEnv ((Name -> Maybe (Name, HsSemanticTokenType))
-> [Name] -> [(Name, HsSemanticTokenType)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (HieFunMaskKind a
-> RefMap a -> Name -> Maybe (Name, HsSemanticTokenType)
forall a.
HieFunMaskKind a
-> RefMap a -> Name -> Maybe (Name, HsSemanticTokenType)
nameNameSemanticFromHie HieFunMaskKind a
hieKind RefMap a
rm) [Name]
names)

nameNameSemanticFromHie :: forall a. HieFunMaskKind a -> RefMap a -> Name -> Maybe (Name, HsSemanticTokenType)
nameNameSemanticFromHie :: forall a.
HieFunMaskKind a
-> RefMap a -> Name -> Maybe (Name, HsSemanticTokenType)
nameNameSemanticFromHie HieFunMaskKind a
hieKind RefMap a
rm Name
ns = do
  HsSemanticTokenType
st <- RefMap a -> Name -> Maybe HsSemanticTokenType
nameSemanticFromRefMap RefMap a
rm Name
ns
  (Name, HsSemanticTokenType) -> Maybe (Name, HsSemanticTokenType)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
ns, HsSemanticTokenType
st)
  where
    nameSemanticFromRefMap :: RefMap a -> Name -> Maybe HsSemanticTokenType
    nameSemanticFromRefMap :: RefMap a -> Name -> Maybe HsSemanticTokenType
nameSemanticFromRefMap RefMap a
rm' Name
name' = do
      [(Span, IdentifierDetails a)]
spanInfos <- Either ModuleName Name
-> RefMap a -> Maybe [(Span, IdentifierDetails a)]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Name -> Either ModuleName Name
forall a b. b -> Either a b
Right Name
name') RefMap a
rm'
      let typeTokenType :: Maybe HsSemanticTokenType
typeTokenType = (a -> Maybe HsSemanticTokenType)
-> Maybe a -> Maybe HsSemanticTokenType
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (HieFunMaskKind a -> a -> Maybe HsSemanticTokenType
forall hType.
HieFunMaskKind hType -> hType -> Maybe HsSemanticTokenType
typeSemantic HieFunMaskKind a
hieKind) (Maybe a -> Maybe HsSemanticTokenType)
-> Maybe a -> Maybe HsSemanticTokenType
forall a b. (a -> b) -> a -> b
$ [a] -> Maybe a
forall a. [a] -> Maybe a
listToMaybe ([a] -> Maybe a) -> [a] -> Maybe a
forall a b. (a -> b) -> a -> b
$ ((Span, IdentifierDetails a) -> Maybe a)
-> [(Span, IdentifierDetails a)] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (IdentifierDetails a -> Maybe a
forall a. IdentifierDetails a -> Maybe a
identType (IdentifierDetails a -> Maybe a)
-> ((Span, IdentifierDetails a) -> IdentifierDetails a)
-> (Span, IdentifierDetails a)
-> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Span, IdentifierDetails a) -> IdentifierDetails a
forall a b. (a, b) -> b
snd) [(Span, IdentifierDetails a)]
spanInfos
      HsSemanticTokenType
contextInfoTokenType <- ((Span, IdentifierDetails a) -> Maybe HsSemanticTokenType)
-> [(Span, IdentifierDetails a)] -> Maybe HsSemanticTokenType
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Set ContextInfo -> Maybe HsSemanticTokenType
contextInfosMaybeTokenType (Set ContextInfo -> Maybe HsSemanticTokenType)
-> ((Span, IdentifierDetails a) -> Set ContextInfo)
-> (Span, IdentifierDetails a)
-> Maybe HsSemanticTokenType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentifierDetails a -> Set ContextInfo
forall a. IdentifierDetails a -> Set ContextInfo
identInfo (IdentifierDetails a -> Set ContextInfo)
-> ((Span, IdentifierDetails a) -> IdentifierDetails a)
-> (Span, IdentifierDetails a)
-> Set ContextInfo
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Span, IdentifierDetails a) -> IdentifierDetails a
forall a b. (a, b) -> b
snd) [(Span, IdentifierDetails a)]
spanInfos
      [Maybe HsSemanticTokenType] -> Maybe HsSemanticTokenType
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold [Maybe HsSemanticTokenType
typeTokenType, HsSemanticTokenType -> Maybe HsSemanticTokenType
forall a. a -> Maybe a
Just HsSemanticTokenType
contextInfoTokenType]

    contextInfosMaybeTokenType :: Set.Set ContextInfo -> Maybe HsSemanticTokenType
    contextInfosMaybeTokenType :: Set ContextInfo -> Maybe HsSemanticTokenType
contextInfosMaybeTokenType Set ContextInfo
details = (ContextInfo -> Maybe HsSemanticTokenType)
-> [ContextInfo] -> Maybe HsSemanticTokenType
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ContextInfo -> Maybe HsSemanticTokenType
infoTokenType (Set ContextInfo -> [ContextInfo]
forall a. Set a -> [a]
Set.toList Set ContextInfo
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 :: forall a. VirtualFile -> HieAST a -> Map Range NameSet
hieAstSpanNames VirtualFile
vf HieAST a
ast =
  if [HieAST a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
ast)
    then HieAST a -> Map Range NameSet
forall {a}. HieAST a -> Map Range NameSet
getIds HieAST a
ast
    else (NameSet -> NameSet -> NameSet)
-> [Map Range NameSet] -> Map Range NameSet
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith NameSet -> NameSet -> NameSet
unionNameSet ([Map Range NameSet] -> Map Range NameSet)
-> [Map Range NameSet] -> Map Range NameSet
forall a b. (a -> b) -> a -> b
$ (HieAST a -> Map Range NameSet)
-> [HieAST a] -> [Map Range NameSet]
forall a b. (a -> b) -> [a] -> [b]
map (VirtualFile -> HieAST a -> Map Range NameSet
forall a. VirtualFile -> HieAST a -> Map Range NameSet
hieAstSpanNames VirtualFile
vf) (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
ast)
  where
    getIds :: HieAST a -> Map Range NameSet
getIds HieAST a
ast' = Map Range NameSet -> Maybe (Map Range NameSet) -> Map Range NameSet
forall a. a -> Maybe a -> a
fromMaybe Map Range NameSet
forall a. Monoid a => a
mempty (Maybe (Map Range NameSet) -> Map Range NameSet)
-> Maybe (Map Range NameSet) -> Map Range NameSet
forall a b. (a -> b) -> a -> b
$ do
      Range
range <- VirtualFile -> CodePointRange -> Maybe Range
codePointRangeToRange VirtualFile
vf (CodePointRange -> Maybe Range) -> CodePointRange -> Maybe Range
forall a b. (a -> b) -> a -> b
$ Span -> CodePointRange
realSrcSpanToCodePointRange (Span -> CodePointRange) -> Span -> CodePointRange
forall a b. (a -> b) -> a -> b
$ HieAST a -> Span
forall a. HieAST a -> Span
nodeSpan HieAST a
ast'
      Map Range NameSet -> Maybe (Map Range NameSet)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Map Range NameSet -> Maybe (Map Range NameSet))
-> Map Range NameSet -> Maybe (Map Range NameSet)
forall a b. (a -> b) -> a -> b
$ Range -> NameSet -> Map Range NameSet
forall k a. k -> a -> Map k a
M.singleton Range
range (HieAST a -> NameSet
forall {a}. HieAST a -> NameSet
getNodeIds' HieAST a
ast')
    getNodeIds' :: HieAST a -> NameSet
getNodeIds' =
      (NameSet -> NodeInfo a -> NameSet)
-> NameSet -> Map NodeOrigin (NodeInfo a) -> NameSet
forall a b k. (a -> b -> a) -> a -> Map k b -> a
Map.foldl' NameSet -> NodeInfo a -> NameSet
forall a. NameSet -> NodeInfo a -> NameSet
combineNodeIds NameSet
forall a. Monoid a => a
mempty
        (Map NodeOrigin (NodeInfo a) -> NameSet)
-> (HieAST a -> Map NodeOrigin (NodeInfo a)) -> HieAST a -> NameSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NodeOrigin -> NodeInfo a -> Bool)
-> Map NodeOrigin (NodeInfo a) -> Map NodeOrigin (NodeInfo a)
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\NodeOrigin
k NodeInfo a
_ -> NodeOrigin
k NodeOrigin -> NodeOrigin -> Bool
forall a. Eq a => a -> a -> Bool
== NodeOrigin
SourceInfo)
        (Map NodeOrigin (NodeInfo a) -> Map NodeOrigin (NodeInfo a))
-> (HieAST a -> Map NodeOrigin (NodeInfo a))
-> HieAST a
-> Map NodeOrigin (NodeInfo a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SourcedNodeInfo a -> Map NodeOrigin (NodeInfo a)
forall a. SourcedNodeInfo a -> Map NodeOrigin (NodeInfo a)
getSourcedNodeInfo
        (SourcedNodeInfo a -> Map NodeOrigin (NodeInfo a))
-> (HieAST a -> SourcedNodeInfo a)
-> HieAST a
-> Map NodeOrigin (NodeInfo a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HieAST a -> SourcedNodeInfo a
forall a. HieAST a -> SourcedNodeInfo a
sourcedNodeInfo
    combineNodeIds :: NameSet -> NodeInfo a -> NameSet
    NameSet
ad combineNodeIds :: forall a. NameSet -> NodeInfo a -> NameSet
`combineNodeIds` (NodeInfo Set NodeAnnotation
_ [a]
_ NodeIdentifiers a
bd) = NameSet
ad NameSet -> NameSet -> NameSet
`unionNameSet` NameSet
xs
      where
        xs :: NameSet
xs = [Name] -> NameSet
mkNameSet ([Name] -> NameSet) -> [Name] -> NameSet
forall a b. (a -> b) -> a -> b
$ [Either ModuleName Name] -> [Name]
forall a b. [Either a b] -> [b]
rights ([Either ModuleName Name] -> [Name])
-> [Either ModuleName Name] -> [Name]
forall a b. (a -> b) -> a -> b
$ NodeIdentifiers a -> [Either ModuleName Name]
forall k a. Map k a -> [k]
M.keys (NodeIdentifiers a -> [Either ModuleName Name])
-> NodeIdentifiers a -> [Either ModuleName Name]
forall a b. (a -> b) -> a -> b
$ (Either ModuleName Name -> IdentifierDetails a -> Bool)
-> NodeIdentifiers a -> NodeIdentifiers a
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey Either ModuleName Name -> IdentifierDetails a -> Bool
forall a. Either ModuleName Name -> IdentifierDetails a -> Bool
inclusion NodeIdentifiers a
bd
        inclusion :: Identifier -> IdentifierDetails a -> Bool
        inclusion :: forall a. Either ModuleName Name -> IdentifierDetails a -> Bool
inclusion Either ModuleName Name
a IdentifierDetails a
b = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Either ModuleName Name -> IdentifierDetails a -> Bool
forall a. Either ModuleName Name -> IdentifierDetails a -> Bool
exclusion Either ModuleName Name
a IdentifierDetails a
b
        exclusion :: Identifier -> IdentifierDetails a -> Bool
        exclusion :: forall a. Either ModuleName Name -> IdentifierDetails a -> Bool
exclusion Either ModuleName Name
idt IdentifierDetails {identInfo :: forall a. IdentifierDetails a -> Set ContextInfo
identInfo = Set ContextInfo
infos} = case Either ModuleName Name
idt of
          Left ModuleName
_  -> Bool
True
          Right Name
_ -> (ContextInfo -> Bool) -> [ContextInfo] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ContextInfo -> Bool
isEvidenceContext (Set ContextInfo -> [ContextInfo]
forall a. Set a -> [a]
S.toList Set ContextInfo
infos)

-------------------------------------------------

-- * extract semantic tokens from NameSemanticMap

-------------------------------------------------

extractSemanticTokensFromNames :: NameSemanticMap -> M.Map Range NameSet -> M.Map Range HsSemanticTokenType
extractSemanticTokensFromNames :: NameSemanticMap
-> Map Range NameSet -> Map Range HsSemanticTokenType
extractSemanticTokensFromNames NameSemanticMap
nsm = (NameSet -> Maybe HsSemanticTokenType)
-> Map Range NameSet -> Map Range HsSemanticTokenType
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe ((Name -> Maybe HsSemanticTokenType)
-> [Name] -> Maybe HsSemanticTokenType
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (NameSemanticMap -> Name -> Maybe HsSemanticTokenType
forall a. NameEnv a -> Name -> Maybe a
lookupNameEnv NameSemanticMap
nsm) ([Name] -> Maybe HsSemanticTokenType)
-> (NameSet -> [Name]) -> NameSet -> Maybe HsSemanticTokenType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NameSet -> [Name]
nameSetElemsStable)

rangeSemanticMapSemanticTokens :: SemanticTokensConfig -> PositionMapping -> M.Map Range HsSemanticTokenType -> Either Text SemanticTokens
rangeSemanticMapSemanticTokens :: SemanticTokensConfig
-> PositionMapping
-> Map Range HsSemanticTokenType
-> Either Text SemanticTokens
rangeSemanticMapSemanticTokens SemanticTokensConfig
stc PositionMapping
mapping =
  SemanticTokensLegend
-> [SemanticTokenAbsolute] -> Either Text SemanticTokens
makeSemanticTokens SemanticTokensLegend
defaultSemanticTokensLegend
    ([SemanticTokenAbsolute] -> Either Text SemanticTokens)
-> (Map Range HsSemanticTokenType -> [SemanticTokenAbsolute])
-> Map Range HsSemanticTokenType
-> Either Text SemanticTokens
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Maybe Range, HsSemanticTokenType) -> Maybe SemanticTokenAbsolute)
-> [(Maybe Range, HsSemanticTokenType)] -> [SemanticTokenAbsolute]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(Maybe Range
range, HsSemanticTokenType
ty) -> (Range -> HsSemanticTokenType -> SemanticTokenAbsolute)
-> HsSemanticTokenType -> Range -> SemanticTokenAbsolute
forall a b c. (a -> b -> c) -> b -> a -> c
flip Range -> HsSemanticTokenType -> SemanticTokenAbsolute
toAbsSemanticToken HsSemanticTokenType
ty (Range -> SemanticTokenAbsolute)
-> Maybe Range -> Maybe SemanticTokenAbsolute
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Range
range)
    ([(Maybe Range, HsSemanticTokenType)] -> [SemanticTokenAbsolute])
-> (Map Range HsSemanticTokenType
    -> [(Maybe Range, HsSemanticTokenType)])
-> Map Range HsSemanticTokenType
-> [SemanticTokenAbsolute]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map (Maybe Range) HsSemanticTokenType
-> [(Maybe Range, HsSemanticTokenType)]
forall k a. Map k a -> [(k, a)]
Map.toAscList
    (Map (Maybe Range) HsSemanticTokenType
 -> [(Maybe Range, HsSemanticTokenType)])
-> (Map Range HsSemanticTokenType
    -> Map (Maybe Range) HsSemanticTokenType)
-> Map Range HsSemanticTokenType
-> [(Maybe Range, HsSemanticTokenType)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Range -> Maybe Range)
-> Map Range HsSemanticTokenType
-> Map (Maybe Range) HsSemanticTokenType
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys (PositionMapping -> Range -> Maybe Range
toCurrentRange PositionMapping
mapping)
  where
    toAbsSemanticToken :: Range -> HsSemanticTokenType -> SemanticTokenAbsolute
    toAbsSemanticToken :: Range -> HsSemanticTokenType -> SemanticTokenAbsolute
toAbsSemanticToken (Range (Position UInt
startLine UInt
startColumn) (Position UInt
_endLine UInt
endColumn)) HsSemanticTokenType
tokenType =
      let len :: UInt
len = UInt
endColumn UInt -> UInt -> UInt
forall a. Num a => a -> a -> a
- UInt
startColumn
       in UInt
-> UInt
-> UInt
-> SemanticTokenTypes
-> [SemanticTokenModifiers]
-> SemanticTokenAbsolute
SemanticTokenAbsolute
            (UInt -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
startLine)
            (UInt -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
startColumn)
            (UInt -> UInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral UInt
len)
            (SemanticTokensConfig -> HsSemanticTokenType -> SemanticTokenTypes
toLspTokenType SemanticTokensConfig
stc HsSemanticTokenType
tokenType)
            []