{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}

module Ide.Plugin.SelectionRange (descriptor) where

import           Control.Monad.Except                    (ExceptT (ExceptT),
                                                          runExceptT)
import           Control.Monad.IO.Class                  (liftIO)
import           Control.Monad.Reader                    (runReader)
import           Control.Monad.Trans.Maybe               (MaybeT (MaybeT),
                                                          maybeToExceptT)
import           Data.Coerce                             (coerce)
import           Data.Containers.ListUtils               (nubOrd)
import           Data.Either.Extra                       (maybeToEither)
import           Data.Foldable                           (find)
import qualified Data.Map.Strict                         as Map
import           Data.Maybe                              (fromMaybe, mapMaybe)
import qualified Data.Text                               as T
import           Development.IDE                         (GetHieAst (GetHieAst),
                                                          HieAstResult (HAR, hieAst, refMap),
                                                          IdeAction,
                                                          IdeState (shakeExtras),
                                                          Range (Range),
                                                          fromNormalizedFilePath,
                                                          ideLogger, logDebug,
                                                          realSrcSpanToRange,
                                                          runIdeAction,
                                                          toNormalizedFilePath',
                                                          uriToFilePath')
import           Development.IDE.Core.Actions            (useE)
import           Development.IDE.Core.PositionMapping    (PositionMapping,
                                                          fromCurrentPosition,
                                                          toCurrentRange)
import           Development.IDE.GHC.Compat              (HieAST (Node), Span,
                                                          getAsts)
import           Development.IDE.GHC.Compat.Util
import           Ide.Plugin.SelectionRange.ASTPreProcess (PreProcessEnv (PreProcessEnv),
                                                          preProcessAST)
import           Ide.PluginUtils                         (response)
import           Ide.Types                               (PluginDescriptor (pluginHandlers),
                                                          PluginId,
                                                          defaultPluginDescriptor,
                                                          mkPluginHandler)
import           Language.LSP.Server                     (LspM)
import           Language.LSP.Types                      (List (List),
                                                          NormalizedFilePath,
                                                          Position,
                                                          ResponseError,
                                                          SMethod (STextDocumentSelectionRange),
                                                          SelectionRange (..),
                                                          SelectionRangeParams (..),
                                                          TextDocumentIdentifier (TextDocumentIdentifier),
                                                          Uri)
import           Prelude                                 hiding (span)

descriptor :: PluginId -> PluginDescriptor IdeState
descriptor :: PluginId -> PluginDescriptor IdeState
descriptor PluginId
plId = (forall ideState. PluginId -> PluginDescriptor ideState
defaultPluginDescriptor PluginId
plId)
    { pluginHandlers :: PluginHandlers IdeState
pluginHandlers = forall (m :: Method 'FromClient 'Request) ideState.
PluginMethod m =>
SClientMethod m
-> PluginMethodHandler ideState m -> PluginHandlers ideState
mkPluginHandler SMethod 'TextDocumentSelectionRange
STextDocumentSelectionRange forall c.
IdeState
-> PluginId
-> SelectionRangeParams
-> LspM c (Either ResponseError (List SelectionRange))
selectionRangeHandler
    }

selectionRangeHandler :: IdeState -> PluginId -> SelectionRangeParams -> LspM c (Either ResponseError (List SelectionRange))
selectionRangeHandler :: forall c.
IdeState
-> PluginId
-> SelectionRangeParams
-> LspM c (Either ResponseError (List SelectionRange))
selectionRangeHandler IdeState
ide PluginId
_ SelectionRangeParams{Maybe ProgressToken
List Position
TextDocumentIdentifier
$sel:_workDoneToken:SelectionRangeParams :: SelectionRangeParams -> Maybe ProgressToken
$sel:_partialResultToken:SelectionRangeParams :: SelectionRangeParams -> Maybe ProgressToken
$sel:_textDocument:SelectionRangeParams :: SelectionRangeParams -> TextDocumentIdentifier
$sel:_positions:SelectionRangeParams :: SelectionRangeParams -> List Position
_positions :: List Position
_textDocument :: TextDocumentIdentifier
_partialResultToken :: Maybe ProgressToken
_workDoneToken :: Maybe ProgressToken
..} = do
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Logger -> Text -> IO ()
logDebug Logger
logger forall a b. (a -> b) -> a -> b
$ Text
"requesting selection range for file: " forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (forall a. Show a => a -> String
show Uri
uri)
    forall (m :: * -> *) a.
Monad m =>
ExceptT String m a -> m (Either ResponseError a)
response forall a b. (a -> b) -> a -> b
$ do
        NormalizedFilePath
filePath <- forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Maybe b -> Either a b
maybeToEither String
"fail to convert uri to file path" forall a b. (a -> b) -> a -> b
$
                String -> NormalizedFilePath
toNormalizedFilePath' forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Uri -> Maybe String
uriToFilePath' Uri
uri
        [SelectionRange]
selectionRanges <- forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. String -> ShakeExtras -> IdeAction a -> IO a
runIdeAction String
"SelectionRange" (IdeState -> ShakeExtras
shakeExtras IdeState
ide) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT forall a b. (a -> b) -> a -> b
$
            NormalizedFilePath
-> [Position] -> ExceptT String IdeAction [SelectionRange]
getSelectionRanges NormalizedFilePath
filePath [Position]
positions
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> List a
List forall a b. (a -> b) -> a -> b
$ [SelectionRange]
selectionRanges
  where
    uri :: Uri
    TextDocumentIdentifier Uri
uri = TextDocumentIdentifier
_textDocument

    positions :: [Position]
    List [Position]
positions = List Position
_positions

    logger :: Logger
logger = IdeState -> Logger
ideLogger IdeState
ide

getSelectionRanges :: NormalizedFilePath -> [Position] -> ExceptT String IdeAction [SelectionRange]
getSelectionRanges :: NormalizedFilePath
-> [Position] -> ExceptT String IdeAction [SelectionRange]
getSelectionRanges NormalizedFilePath
file [Position]
positions = do
    (HAR{HieASTs a
hieAst :: HieASTs a
hieAst :: ()
hieAst, RefMap a
refMap :: RefMap a
refMap :: ()
refMap}, PositionMapping
positionMapping) <- forall (m :: * -> *) e a.
Functor m =>
e -> MaybeT m a -> ExceptT e m a
maybeToExceptT String
"fail to get hie ast" forall a b. (a -> b) -> a -> b
$ forall k v.
IdeRule k v =>
k -> NormalizedFilePath -> MaybeT IdeAction (v, PositionMapping)
useE GetHieAst
GetHieAst NormalizedFilePath
file
    -- 'positionMapping' should be applied to the input positions before using them
    [Position]
positions' <- forall (m :: * -> *) e a.
Functor m =>
e -> MaybeT m a -> ExceptT e m a
maybeToExceptT String
"fail to apply position mapping to input positions" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (PositionMapping -> Position -> Maybe Position
fromCurrentPosition PositionMapping
positionMapping) [Position]
positions

    HieAST a
ast <- forall (m :: * -> *) e a.
Functor m =>
e -> MaybeT m a -> ExceptT e m a
maybeToExceptT String
"fail to get ast for current file" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        -- in GHC 9, the 'FastString' in 'HieASTs' is replaced by a newtype wrapper around 'LexicalFastString'
        -- so we use 'coerce' to make it work in both GHC 8 and 9
        forall a. HieASTs a -> Map HiePath (HieAST a)
getAsts HieASTs a
hieAst forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? (coerce :: forall a b. Coercible a b => a -> b
coerce forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> FastString
mkFastString forall b c a. (b -> c) -> (a -> b) -> a -> c
. NormalizedFilePath -> String
fromNormalizedFilePath) NormalizedFilePath
file

    let ast' :: HieAST a
ast' = forall r a. Reader r a -> r -> a
runReader (forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
preProcessAST HieAST a
ast) (forall a. RefMap a -> PreProcessEnv a
PreProcessEnv RefMap a
refMap)
    let selectionRanges :: [SelectionRange]
selectionRanges = [SelectionRange] -> [Position] -> [SelectionRange]
findSelectionRangesByPositions (forall a. HieAST a -> [SelectionRange]
astPathsLeafToRoot HieAST a
ast') [Position]
positions'

    -- 'positionMapping' should be applied to the output ranges before returning them
    forall (m :: * -> *) e a.
Functor m =>
e -> MaybeT m a -> ExceptT e m a
maybeToExceptT String
"fail to apply position mapping to output positions" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
         forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (PositionMapping -> SelectionRange -> Maybe SelectionRange
toCurrentSelectionRange PositionMapping
positionMapping) [SelectionRange]
selectionRanges

-- | Likes 'toCurrentPosition', but works on 'SelectionRange'
toCurrentSelectionRange :: PositionMapping -> SelectionRange -> Maybe SelectionRange
toCurrentSelectionRange :: PositionMapping -> SelectionRange -> Maybe SelectionRange
toCurrentSelectionRange PositionMapping
positionMapping SelectionRange{Maybe SelectionRange
Range
$sel:_range:SelectionRange :: SelectionRange -> Range
$sel:_parent:SelectionRange :: SelectionRange -> Maybe SelectionRange
_parent :: Maybe SelectionRange
_range :: Range
..} = do
    Range
newRange <- PositionMapping -> Range -> Maybe Range
toCurrentRange PositionMapping
positionMapping Range
_range
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ SelectionRange {
        $sel:_range:SelectionRange :: Range
_range = Range
newRange,
        $sel:_parent:SelectionRange :: Maybe SelectionRange
_parent = Maybe SelectionRange
_parent forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PositionMapping -> SelectionRange -> Maybe SelectionRange
toCurrentSelectionRange PositionMapping
positionMapping
    }

-- | Build all paths from ast leaf to root
astPathsLeafToRoot :: HieAST a -> [SelectionRange]
astPathsLeafToRoot :: forall a. HieAST a -> [SelectionRange]
astPathsLeafToRoot = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ([Span] -> Maybe SelectionRange
spansToSelectionRange forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
nubOrd) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [[Span]] -> HieAST a -> [[Span]]
go [[]]
  where
    go :: [[Span]] -> HieAST a -> [[Span]]
    go :: forall a. [[Span]] -> HieAST a -> [[Span]]
go [[Span]]
acc (Node SourcedNodeInfo a
_ Span
span [])       = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Span
spanforall a. a -> [a] -> [a]
:) [[Span]]
acc
    go [[Span]]
acc (Node SourcedNodeInfo a
_ Span
span [HieAST a]
children) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. [[Span]] -> HieAST a -> [[Span]]
go (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Span
spanforall a. a -> [a] -> [a]
:) [[Span]]
acc)) [HieAST a]
children

spansToSelectionRange :: [Span] -> Maybe SelectionRange
spansToSelectionRange :: [Span] -> Maybe SelectionRange
spansToSelectionRange [] = forall a. Maybe a
Nothing
spansToSelectionRange (Span
span:[Span]
spans) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$
    SelectionRange {$sel:_range:SelectionRange :: Range
_range = Span -> Range
realSrcSpanToRange Span
span, $sel:_parent:SelectionRange :: Maybe SelectionRange
_parent = [Span] -> Maybe SelectionRange
spansToSelectionRange [Span]
spans}

{-|
For each position, find the selection range that contains it, without taking each selection range's
parent into account. These selection ranges are un-divisible, representing the leaf nodes in original AST, so they
won't overlap.
-}
findSelectionRangesByPositions :: [SelectionRange] -- ^ all possible selection ranges
                               -> [Position] -- ^ requested positions
                               -> [SelectionRange]
findSelectionRangesByPositions :: [SelectionRange] -> [Position] -> [SelectionRange]
findSelectionRangesByPositions [SelectionRange]
selectionRanges = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Position -> SelectionRange
findByPosition
    {-
        Performance Tips:
        Doing a linear search from the first selection range for each position is not optimal.
        If it becomes too slow for a large file and many positions, you may optimize the implementation.
        Assume the number of selection range is n, then the following techniques may be applied:
            1. For each position, we may treat HieAST as a position indexed tree to search it in O(log(n)).
            2. For all positions, a searched position will narrow the search range for other positions.
    -}
  where
    findByPosition :: Position -> SelectionRange
    findByPosition :: Position -> SelectionRange
findByPosition Position
p = forall a. a -> Maybe a -> a
fromMaybe SelectionRange{$sel:_range:SelectionRange :: Range
_range = Position -> Position -> Range
Range Position
p Position
p, $sel:_parent:SelectionRange :: Maybe SelectionRange
_parent = forall a. Maybe a
Nothing} forall a b. (a -> b) -> a -> b
$
        forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (Position -> SelectionRange -> Bool
isPositionInSelectionRange Position
p) [SelectionRange]
selectionRanges

    isPositionInSelectionRange :: Position -> SelectionRange -> Bool
    isPositionInSelectionRange :: Position -> SelectionRange -> Bool
isPositionInSelectionRange Position
p SelectionRange{Range
_range :: Range
$sel:_range:SelectionRange :: SelectionRange -> Range
_range} =
        let Range Position
sp Position
ep = Range
_range in Position
sp forall a. Ord a => a -> a -> Bool
<= Position
p Bool -> Bool -> Bool
&& Position
p forall a. Ord a => a -> a -> Bool
<= Position
ep