expand selection range by HieAST

haskell · Jan 10, 2022 · 85d0ddb · 85d0ddb
1 parent d81aad8
commit 85d0ddb
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 26 deletions.
diff --git a/ghcide/src/Development/IDE/Core/Actions.hs b/ghcide/src/Development/IDE/Core/Actions.hs
@@ -10,7 +10,6 @@ module Development.IDE.Core.Actions
 , useNoFileE
 , usesE
 , workspaceSymbols
-, getSelectionRanges
 ) where
 
 import           Control.Monad.Reader
@@ -29,10 +28,8 @@ import           Development.IDE.Graph
 import qualified Development.IDE.Spans.AtPoint        as AtPoint
 import           Development.IDE.Types.HscEnvEq       (hscEnv)
 import           Development.IDE.Types.Location
-import           Development.IDE.Types.Logger
 import qualified HieDb
 import           Language.LSP.Types                   (DocumentHighlight (..),
-                                                       SelectionRange,
                                                        SymbolInformation (..))
 
 
@@ -122,10 +119,3 @@ workspaceSymbols query = runMaybeT $ do
   ShakeExtras{withHieDb} <- ask
   res <- liftIO $ withHieDb (\hieDb -> HieDb.searchDef hieDb $ T.unpack query)
   pure $ mapMaybe AtPoint.defRowToSymbolInfo res
-
-getSelectionRanges :: NormalizedFilePath -> [Position] -> IdeAction [SelectionRange]
-getSelectionRanges file positions = fmap (fromMaybe []) <$> runMaybeT $ do
-    ShakeExtras{logger} <- ask
-    (parsedModule, _) <- useE GetParsedModuleWithComments file
-    liftIO $ logDebug logger $ T.pack (show parsedModule)
-    pure []
diff --git a/ghcide/src/Development/IDE/Plugin/SelectionRange.hs b/ghcide/src/Development/IDE/Plugin/SelectionRange.hs
@@ -1,23 +1,38 @@
+{-# LANGUAGE RankNTypes #-}
 module Development.IDE.Plugin.SelectionRange
     ( descriptor
     ) where
 
-import           Control.Monad.IO.Class       (liftIO)
-import           Development.IDE              (IdeState (shakeExtras),
-                                               runIdeAction,
-                                               toNormalizedFilePath',
-                                               uriToFilePath')
-import           Development.IDE.Core.Actions (getSelectionRanges)
-import           Ide.Types                    (PluginDescriptor (pluginHandlers),
-                                               PluginId,
-                                               defaultPluginDescriptor,
-                                               mkPluginHandler)
-import           Language.LSP.Server          (LspM)
-import           Language.LSP.Types           (List (List), ResponseError,
-                                               SMethod (STextDocumentSelectionRange),
-                                               SelectionRange,
-                                               SelectionRangeParams (..),
-                                               TextDocumentIdentifier (TextDocumentIdentifier))
+import           Control.Monad.IO.Class          (liftIO)
+import           Control.Monad.Trans.Maybe       (MaybeT (MaybeT), runMaybeT)
+import           Data.Foldable                   (find)
+import qualified Data.Map.Strict                 as Map
+import           Data.Maybe                      (fromMaybe, mapMaybe)
+import           Development.IDE                 (GetHieAst (GetHieAst),
+                                                  HieAstResult (HAR, hieAst),
+                                                  IdeAction,
+                                                  IdeState (shakeExtras),
+                                                  Range (Range),
+                                                  fromNormalizedFilePath,
+                                                  realSrcSpanToRange,
+                                                  runIdeAction,
+                                                  toNormalizedFilePath',
+                                                  uriToFilePath')
+import           Development.IDE.Core.Actions    (useE)
+import           Development.IDE.GHC.Compat      (HieAST (Node), Span, getAsts)
+import           Development.IDE.GHC.Compat.Util (mkFastString)
+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))
 
 descriptor :: PluginId -> PluginDescriptor IdeState
 descriptor plId = (defaultPluginDescriptor plId)
@@ -34,3 +49,29 @@ selectionRangeHandler ide _ SelectionRangeParams{..} = do
             let (List positions) = _positions
             selectionRanges <- runIdeAction "SelectionRange" (shakeExtras ide) $ getSelectionRanges filePath positions
             pure . Right . List $ selectionRanges
+
+getSelectionRanges :: NormalizedFilePath -> [Position] -> IdeAction [SelectionRange]
+getSelectionRanges file positions = fmap (fromMaybe []) <$> runMaybeT $ do
+    (HAR{hieAst}, _) <- useE GetHieAst file
+    ast <- MaybeT . pure $ getAsts hieAst Map.!? (mkFastString . fromNormalizedFilePath) file
+    pure $ findSelectionRangesByPositions (astPathsLeafToRoot ast) positions
+
+-- |build all paths from ast leaf to root
+astPathsLeafToRoot :: HieAST a -> [SelectionRange]
+astPathsLeafToRoot = mapMaybe spansToSelectionRange . go [[]]
+  where
+    go acc (Node _ span [])       = fmap (span:) acc
+    go acc (Node _ span children) = concatMap (go (fmap (span:) acc)) children
+
+spansToSelectionRange :: [Span] -> Maybe SelectionRange
+spansToSelectionRange [] = Nothing
+spansToSelectionRange (span:spans) = Just $
+    SelectionRange {_range = realSrcSpanToRange span, _parent = spansToSelectionRange spans}
+
+findSelectionRangesByPositions :: [SelectionRange] -> [Position] -> [SelectionRange]
+findSelectionRangesByPositions selectionRanges = fmap findByPosition
+  where
+    findByPosition p = fromMaybe SelectionRange{_range = Range p p, _parent = Nothing} $
+        find (isPositionInSelectionRange p) selectionRanges
+    isPositionInSelectionRange p SelectionRange{_range} =
+        let Range sp ep = _range in sp <= p && p <= ep