Internal.hs

{-# LANGUAGE CPP               #-}
{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}

module Ide.Plugin.CallHierarchy.Internal (
  prepareCallHierarchy
, incomingCalls
, outgoingCalls
) where

import           Control.Lens                   (Lens', (^.))
import           Control.Monad.IO.Class
import           Data.Aeson                     as A
import           Data.Functor                   ((<&>))
import           Data.List                      (groupBy, sortBy)
import qualified Data.Map                       as M
import           Data.Maybe
import           Data.Ord                       (comparing)
import qualified Data.Set                       as S
import qualified Data.Text                      as T
import           Data.Tuple.Extra
import           Development.IDE
import           Development.IDE.Core.Shake
import           Development.IDE.GHC.Compat     as Compat
import           Development.IDE.Spans.AtPoint
import           HieDb                          (Symbol (Symbol))
import qualified Ide.Plugin.CallHierarchy.Query as Q
import           Ide.Plugin.CallHierarchy.Types
import           Ide.Plugin.Error
import           Ide.Types
import qualified Language.LSP.Protocol.Lens     as L
import           Language.LSP.Protocol.Message
import           Language.LSP.Protocol.Types
import           Prelude                        hiding (mod, span)
import           Text.Read                      (readMaybe)

-- | Render prepare call hierarchy request.
prepareCallHierarchy :: PluginMethodHandler IdeState Method_TextDocumentPrepareCallHierarchy
prepareCallHierarchy state _ param = do
    nfp <- getNormalizedFilePathE (param ^. (L.textDocument . L.uri))
    items <- liftIO
        $ runAction "CallHierarchy.prepareHierarchy" state
        $ prepareCallHierarchyItem nfp (param ^. L.position)
    pure $ InL items

prepareCallHierarchyItem :: NormalizedFilePath -> Position -> Action [CallHierarchyItem]
prepareCallHierarchyItem nfp pos = use GetHieAst nfp <&> \case
    Nothing               -> mempty
    Just (HAR _ hf _ _ _) -> prepareByAst hf pos nfp

prepareByAst :: HieASTs a -> Position -> NormalizedFilePath -> [CallHierarchyItem]
prepareByAst hf pos nfp =
    case listToMaybe $ pointCommand hf pos extract of
        Nothing    -> mempty
        Just infos -> mapMaybe (construct nfp hf) infos

extract :: HieAST a -> [(Identifier, [ContextInfo], Span)]
extract ast = let span = nodeSpan ast
                  infos = M.toList $ M.map (S.toList . identInfo) (Compat.getNodeIds ast)
              in  [(ident, contexts, span) | (ident, contexts) <- infos]

recFieldInfo, declInfo, valBindInfo, classTyDeclInfo,
    useInfo, patternBindInfo, tyDeclInfo, matchBindInfo :: [ContextInfo] -> Maybe ContextInfo
recFieldInfo    ctxs = listToMaybe [ctx       | ctx@RecField{}    <- ctxs]
declInfo        ctxs = listToMaybe [ctx       | ctx@Decl{}        <- ctxs]
valBindInfo     ctxs = listToMaybe [ctx       | ctx@ValBind{}     <- ctxs]
classTyDeclInfo ctxs = listToMaybe [ctx       | ctx@ClassTyDecl{} <- ctxs]
useInfo         ctxs = listToMaybe [Use       | Use               <- ctxs]
patternBindInfo ctxs = listToMaybe [ctx       | ctx@PatternBind{} <- ctxs]
tyDeclInfo      ctxs = listToMaybe [TyDecl    | TyDecl            <- ctxs]
matchBindInfo   ctxs = listToMaybe [MatchBind | MatchBind         <- ctxs]

construct :: NormalizedFilePath -> HieASTs a -> (Identifier, [ContextInfo], Span) -> Maybe CallHierarchyItem
construct nfp hf (ident, contexts, ssp)
    | isInternalIdentifier ident = Nothing

    | Just (RecField RecFieldDecl _) <- recFieldInfo contexts
        -- ignored type span
        = Just $ mkCallHierarchyItem' ident SymbolKind_Field ssp ssp

    | isJust (matchBindInfo contexts) && isNothing (valBindInfo contexts)
        = Just $ mkCallHierarchyItem' ident SymbolKind_Function ssp ssp

    | Just ctx <- valBindInfo contexts
        = Just $ case ctx of
            ValBind _ _ span -> mkCallHierarchyItem' ident SymbolKind_Function (renderSpan span) ssp
            _                -> mkCallHierarchyItem' ident skUnknown ssp ssp

    | Just ctx <- declInfo contexts
        = Just $ case ctx of
            Decl ClassDec span -> mkCallHierarchyItem' ident SymbolKind_Interface     (renderSpan span) ssp
            Decl ConDec   span -> mkCallHierarchyItem' ident SymbolKind_Constructor   (renderSpan span) ssp
            Decl DataDec  span -> mkCallHierarchyItem' ident SymbolKind_Struct        (renderSpan span) ssp
            Decl FamDec   span -> mkCallHierarchyItem' ident SymbolKind_Function      (renderSpan span) ssp
            Decl InstDec  span -> mkCallHierarchyItem' ident SymbolKind_Interface     (renderSpan span) ssp
            Decl SynDec   span -> mkCallHierarchyItem' ident SymbolKind_TypeParameter (renderSpan span) ssp
            _ -> mkCallHierarchyItem' ident skUnknown ssp ssp

    | Just (ClassTyDecl span) <- classTyDeclInfo contexts
        = Just $ mkCallHierarchyItem' ident SymbolKind_Method (renderSpan span) ssp

    | Just (PatternBind _ _ span) <- patternBindInfo contexts
        = Just $ mkCallHierarchyItem' ident SymbolKind_Function (renderSpan span) ssp

    | Just _ <- useInfo contexts = Just $ mkCallHierarchyItem' ident SymbolKind_Interface ssp ssp

    | Just _ <- tyDeclInfo contexts = renderTyDecl

    | otherwise = Nothing
    where
        renderSpan (Just span) = span
        renderSpan _           = ssp

        -- https://github.com/haskell/lsp/blob/e11b7c09658610f6d815d04db08a64e7cf6b4467/lsp-types/src/Language/LSP/Types/DocumentSymbol.hs#L97
        -- There is no longer an unknown symbol, thus using SymbolKind_Function
        -- as this is the call-hierarchy plugin
        skUnknown = SymbolKind_Function

        mkCallHierarchyItem' = mkCallHierarchyItem nfp

        isInternalIdentifier = \case
            Left _     -> False
            Right name -> isInternalName name

        renderTyDecl = case ident of
            Left _ -> Nothing
            Right name -> case getNameBinding name (getAsts hf) of
                Nothing -> Nothing
                Just sp -> listToMaybe $ prepareByAst hf (realSrcSpanToRange sp ^. L.start) nfp

mkCallHierarchyItem :: NormalizedFilePath -> Identifier -> SymbolKind -> Span -> Span -> CallHierarchyItem
mkCallHierarchyItem nfp ident kind span selSpan =
    CallHierarchyItem
        (T.pack $ optimizeDisplay $ identifierName ident)
        kind
        Nothing
        (Just $ T.pack $ identifierToDetail ident)
        (fromNormalizedUri $ normalizedFilePathToUri nfp)
        (realSrcSpanToRange span)
        (realSrcSpanToRange selSpan)
        (toJSON . show <$> mkSymbol ident)
    where
        identifierToDetail :: Identifier -> String
        identifierToDetail = \case
            Left modName -> moduleNameString modName
            Right name   -> (moduleNameString . moduleName . nameModule) name

        identifierName :: Identifier -> String
        identifierName = \case
            Left modName -> moduleNameString modName
            Right name   -> occNameString $ nameOccName name

        optimizeDisplay :: String -> String
        optimizeDisplay name -- Optimize display for DuplicateRecordFields
            | "$sel:" == take 5 name = drop 5 name
            | otherwise = name

mkSymbol :: Identifier -> Maybe Symbol
mkSymbol = \case
    Left _     -> Nothing
    Right name -> Just $ Symbol (occName name) (nameModule name)

----------------------------------------------------------------------
-------------- Incoming calls and outgoing calls ---------------------
----------------------------------------------------------------------

-- | Render incoming calls request.
incomingCalls :: PluginMethodHandler IdeState Method_CallHierarchyIncomingCalls
incomingCalls state _pluginId param = do
    calls <- liftIO
        $ runAction "CallHierarchy.incomingCalls" state
        $ queryCalls
            (param ^. L.item)
            Q.incomingCalls
            mkCallHierarchyIncomingCall
            (mergeCalls CallHierarchyIncomingCall L.from)
    pure $ InL calls
    where
        mkCallHierarchyIncomingCall :: Vertex -> Action (Maybe CallHierarchyIncomingCall)
        mkCallHierarchyIncomingCall = mkCallHierarchyCall CallHierarchyIncomingCall

-- | Render outgoing calls request.
outgoingCalls :: PluginMethodHandler IdeState Method_CallHierarchyOutgoingCalls
outgoingCalls state _pluginId param = do
    calls <- liftIO
        $ runAction "CallHierarchy.outgoingCalls" state
        $ queryCalls
            (param ^. L.item)
            Q.outgoingCalls
            mkCallHierarchyOutgoingCall
            (mergeCalls CallHierarchyOutgoingCall L.to)
    pure $ InL calls
    where
        mkCallHierarchyOutgoingCall :: Vertex -> Action (Maybe CallHierarchyOutgoingCall)
        mkCallHierarchyOutgoingCall = mkCallHierarchyCall CallHierarchyOutgoingCall

-- | Merge calls from the same place
mergeCalls ::
    L.HasFromRanges s [Range]
    => (CallHierarchyItem -> [Range] -> s)
    -> Lens' s CallHierarchyItem
    -> [s]
    -> [s]
mergeCalls constructor target =
    concatMap merge
        . groupBy (\a b -> a ^. target == b ^. target)
        . sortBy (comparing (^. target))
    where
        merge [] = []
        merge calls@(call:_) =
            let ranges = concatMap (^. L.fromRanges) calls
            in  [constructor (call ^. target) ranges]

mkCallHierarchyCall :: (CallHierarchyItem ->  [Range] -> a) -> Vertex -> Action (Maybe a)
mkCallHierarchyCall mk v@Vertex{..} = do
    let pos = Position (fromIntegral $ sl - 1) (fromIntegral $ sc - 1)
        nfp = toNormalizedFilePath' hieSrc
        range = mkRange
                    (fromIntegral $ casl - 1)
                    (fromIntegral $ casc - 1)
                    (fromIntegral $ cael - 1)
                    (fromIntegral $ caec - 1)

    prepareCallHierarchyItem nfp pos >>=
        \case
            [item] -> pure $ Just $ mk item [range]
            _      -> do
                ShakeExtras{withHieDb} <- getShakeExtras
                sps <- liftIO (withHieDb (`Q.getSymbolPosition` v))
                case sps of
                    (x:_) -> do
                        items <- prepareCallHierarchyItem
                                    nfp
                                    (Position (fromIntegral $ psl x - 1) (fromIntegral $ psc x - 1))
                        case items of
                            [item] -> pure $ Just $ mk item [range]
                            _      -> pure Nothing
                    []     -> pure Nothing

-- | Unified queries include incoming calls and outgoing calls.
queryCalls ::
    CallHierarchyItem
    -> (HieDb -> Symbol -> IO [Vertex])
    -> (Vertex -> Action (Maybe a))
    -> ([a] -> [a])
    -> Action [a]
queryCalls item queryFunc makeFunc merge
    | Just nfp <- uriToNormalizedFilePath $ toNormalizedUri uri = do
        ShakeExtras{withHieDb} <- getShakeExtras
        maySymbol <- getSymbol nfp
        case maySymbol of
            Nothing -> pure mempty
            Just symbol -> do
                vs <- liftIO $ withHieDb (`queryFunc` symbol)
                items <- catMaybes <$> mapM makeFunc vs
                pure $ merge items
    | otherwise = pure mempty
    where
        uri = item ^. L.uri
        pos = item ^. (L.selectionRange . L.start)

        getSymbol nfp = case item ^. L.data_ of
            Just xdata -> case fromJSON xdata of
                A.Success (symbolStr :: String) -> maybe (getSymbolFromAst nfp pos) (pure . pure) $ readMaybe symbolStr
                A.Error _ -> getSymbolFromAst nfp pos
            Nothing -> getSymbolFromAst nfp pos -- Fallback if xdata lost, some editor(VSCode) will drop it

        getSymbolFromAst :: NormalizedFilePath -> Position -> Action (Maybe Symbol)
        getSymbolFromAst nfp pos_ = use GetHieAst nfp <&> \case
            Nothing -> Nothing
            Just (HAR _ hf _ _ _) -> do
                case listToMaybe $ pointCommand hf pos_ extract of
                    Just infos -> mkSymbol . fst3 =<< listToMaybe infos
                    Nothing    -> Nothing