-- |
-- The query module is used to query the semantic tokens from the AST
module Ide.Plugin.SemanticTokens.Query where

import           Control.Applicative                  ((<|>))
import           Data.Foldable                        (fold)
import qualified Data.Map.Strict                      as M
import           Data.Maybe                           (listToMaybe, mapMaybe)
import qualified Data.Set                             as Set
import           Data.Text                            (Text)
import           Development.IDE.Core.PositionMapping (PositionMapping,
                                                       toCurrentRange)
import           Development.IDE.GHC.Compat
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),
                                                       SemanticTokensDelta (SemanticTokensDelta),
                                                       defaultSemanticTokensLegend,
                                                       makeSemanticTokens,
                                                       makeSemanticTokensDelta)
import           Prelude                              hiding (length, span)

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

-- * extract semantic

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

idSemantic :: forall a. NameEnv TyThing -> HieFunMaskKind a -> RefMap a -> Identifier -> Maybe HsSemanticTokenType
idSemantic :: forall a.
NameEnv TyThing
-> HieFunMaskKind a
-> RefMap a
-> Identifier
-> Maybe HsSemanticTokenType
idSemantic NameEnv TyThing
_ HieFunMaskKind a
_ RefMap a
_ (Left ModuleName
_) = HsSemanticTokenType -> Maybe HsSemanticTokenType
forall a. a -> Maybe a
Just HsSemanticTokenType
TModule
idSemantic NameEnv TyThing
tyThingMap HieFunMaskKind a
hieKind RefMap a
rm (Right Name
n) =
    HieFunMaskKind a -> RefMap a -> Name -> Maybe HsSemanticTokenType
forall a.
HieFunMaskKind a -> RefMap a -> Name -> Maybe HsSemanticTokenType
nameSemanticFromHie HieFunMaskKind a
hieKind RefMap a
rm Name
n -- local name
    Maybe HsSemanticTokenType
-> Maybe HsSemanticTokenType -> Maybe HsSemanticTokenType
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (NameEnv TyThing -> Name -> Maybe TyThing
forall a. NameEnv a -> Name -> Maybe a
lookupNameEnv NameEnv TyThing
tyThingMap Name
n Maybe TyThing
-> (TyThing -> Maybe HsSemanticTokenType)
-> Maybe HsSemanticTokenType
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TyThing -> Maybe HsSemanticTokenType
tyThingSemantic) -- global name


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

-- * extract semantic from HieAst for local variables

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

nameSemanticFromHie :: forall a. HieFunMaskKind a -> RefMap a -> Name -> Maybe HsSemanticTokenType
nameSemanticFromHie :: forall a.
HieFunMaskKind a -> RefMap a -> Name -> Maybe HsSemanticTokenType
nameSemanticFromHie HieFunMaskKind a
hieKind RefMap a
rm Name
n = RefMap a -> Identifier -> Maybe HsSemanticTokenType
idSemanticFromRefMap RefMap a
rm (Name -> Identifier
forall a b. b -> Either a b
Right Name
n)
  where
    idSemanticFromRefMap :: RefMap a -> Identifier -> Maybe HsSemanticTokenType
    idSemanticFromRefMap :: RefMap a -> Identifier -> Maybe HsSemanticTokenType
idSemanticFromRefMap RefMap a
rm' Identifier
name' = do
      [(Span, IdentifierDetails a)]
spanInfos <- Identifier -> RefMap a -> Maybe [(Span, IdentifierDetails a)]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Identifier
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, Name -> Maybe HsSemanticTokenType
nameInfixOperator Name
n]

    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 lsp semantic tokens from RangeSemanticTokenTypeList

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

rangeSemanticsSemanticTokens :: SemanticTokenId -> SemanticTokensConfig -> PositionMapping -> RangeSemanticTokenTypeList -> Either Text SemanticTokens
rangeSemanticsSemanticTokens :: Text
-> SemanticTokensConfig
-> PositionMapping
-> RangeSemanticTokenTypeList
-> Either Text SemanticTokens
rangeSemanticsSemanticTokens Text
sid SemanticTokensConfig
stc PositionMapping
mapping =
  Maybe Text -> [SemanticTokenAbsolute] -> Either Text SemanticTokens
makeSemanticTokensWithId (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
sid) ([SemanticTokenAbsolute] -> Either Text SemanticTokens)
-> (RangeSemanticTokenTypeList -> [SemanticTokenAbsolute])
-> RangeSemanticTokenTypeList
-> Either Text SemanticTokens
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Range, HsSemanticTokenType) -> Maybe SemanticTokenAbsolute)
-> RangeSemanticTokenTypeList -> [SemanticTokenAbsolute]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(Range
ran, HsSemanticTokenType
tk) -> Range -> HsSemanticTokenType -> SemanticTokenAbsolute
toAbsSemanticToken (Range -> HsSemanticTokenType -> SemanticTokenAbsolute)
-> Maybe Range
-> Maybe (HsSemanticTokenType -> SemanticTokenAbsolute)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PositionMapping -> Range -> Maybe Range
toCurrentRange PositionMapping
mapping Range
ran Maybe (HsSemanticTokenType -> SemanticTokenAbsolute)
-> Maybe HsSemanticTokenType -> Maybe SemanticTokenAbsolute
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> HsSemanticTokenType -> Maybe HsSemanticTokenType
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return HsSemanticTokenType
tk)
  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)
            []

makeSemanticTokensWithId :: Maybe SemanticTokenId -> [SemanticTokenAbsolute] -> Either Text SemanticTokens
makeSemanticTokensWithId :: Maybe Text -> [SemanticTokenAbsolute] -> Either Text SemanticTokens
makeSemanticTokensWithId Maybe Text
sid [SemanticTokenAbsolute]
tokens = do
    (SemanticTokens Maybe Text
_  [UInt]
tokens) <- SemanticTokensLegend
-> [SemanticTokenAbsolute] -> Either Text SemanticTokens
makeSemanticTokens SemanticTokensLegend
defaultSemanticTokensLegend [SemanticTokenAbsolute]
tokens
    SemanticTokens -> Either Text SemanticTokens
forall a. a -> Either Text a
forall (m :: * -> *) a. Monad m => a -> m a
return (SemanticTokens -> Either Text SemanticTokens)
-> SemanticTokens -> Either Text SemanticTokens
forall a b. (a -> b) -> a -> b
$ Maybe Text -> [UInt] -> SemanticTokens
SemanticTokens Maybe Text
sid [UInt]
tokens

makeSemanticTokensDeltaWithId :: Maybe SemanticTokenId ->  SemanticTokens -> SemanticTokens -> SemanticTokensDelta
makeSemanticTokensDeltaWithId :: Maybe Text
-> SemanticTokens -> SemanticTokens -> SemanticTokensDelta
makeSemanticTokensDeltaWithId Maybe Text
sid SemanticTokens
previousTokens SemanticTokens
currentTokens =
    let (SemanticTokensDelta Maybe Text
_ [SemanticTokensEdit]
stEdits) = SemanticTokens -> SemanticTokens -> SemanticTokensDelta
makeSemanticTokensDelta SemanticTokens
previousTokens SemanticTokens
currentTokens
    in Maybe Text -> [SemanticTokensEdit] -> SemanticTokensDelta
SemanticTokensDelta Maybe Text
sid [SemanticTokensEdit]
stEdits