Use a command to extend imports

Author: pepeiborra

ghcide/src/Development/IDE/Plugin/Completions.hs

@@ -12,6 +12,12 @@ import Language.Haskell.LSP.Types
 import qualified Language.Haskell.LSP.Core as LSP
 import qualified Language.Haskell.LSP.VFS as VFS
 
+import Control.Monad
+import Control.Monad.Trans.Maybe
+import Data.Aeson
+import Data.List (find)
+import Data.Maybe
+import qualified Data.Text as T
 import Development.Shake.Classes
 import Development.Shake
 import GHC.Generics
@@ -22,23 +28,24 @@ import Development.IDE.Types.Location
 import Development.IDE.Core.RuleTypes
 import Development.IDE.Core.Shake
 import Development.IDE.GHC.Compat
-
+import Development.IDE.GHC.ExactPrint (Annotated (annsA), GetAnnotatedParsedSource (GetAnnotatedParsedSource))
 import Development.IDE.GHC.Util
-import TcRnDriver (tcRnImportDecls)
-import Data.Maybe
+import Development.IDE.Plugin.CodeAction.ExactPrint
+import Development.IDE.Plugin.Completions.Types
 import Ide.Plugin.Config (Config (completionSnippetsOn))
 import Ide.PluginUtils (getClientConfig)
 import Ide.Types
-
+import TcRnDriver (tcRnImportDecls)
 #if defined(GHC_LIB)
 import Development.IDE.Import.DependencyInformation
 #endif
 
 descriptor :: PluginId -> PluginDescriptor IdeState
 descriptor plId = (defaultPluginDescriptor plId)
-  { pluginRules = produceCompletions
-  , pluginCompletionProvider = Just getCompletionsLSP
-  }
+    { pluginRules = produceCompletions,
+      pluginCompletionProvider = Just (getCompletionsLSP plId),
+      pluginCommands = [extendImportCommand]
+    }
 
 produceCompletions :: Rules ()
 produceCompletions = do
@@ -48,10 +55,11 @@ produceCompletions = do
         let extract = fmap fst
         return ([], extract local <> extract nonLocal)
     define $ \LocalCompletions file -> do
+        let uri = fromNormalizedUri $ normalizedFilePathToUri file
         pm <- useWithStale GetParsedModule file
         case pm of
             Just (pm, _) -> do
-                let cdata = localCompletionsForParsedModule pm
+                let cdata = localCompletionsForParsedModule uri pm
                 return ([], Just cdata)
             _ -> return ([], Nothing)
     define $ \NonLocalCompletions file -> do
@@ -77,7 +85,8 @@ produceCompletions = do
               res <- liftIO $ tcRnImportDecls env (dropListFromImportDecl <$> imps)
               case res of
                   (_, Just rdrEnv) -> do
-                      cdata <- liftIO $ cacheDataProducer env (ms_mod ms) rdrEnv imps parsedDeps
+                      let uri = fromNormalizedUri $ normalizedFilePathToUri file
+                      cdata <- liftIO $ cacheDataProducer uri env (ms_mod ms) rdrEnv imps parsedDeps
                       return ([], Just cdata)
                   (_diag, _) ->
                       return ([], Nothing)
@@ -115,13 +124,15 @@ data NonLocalCompletions = NonLocalCompletions
 instance Hashable NonLocalCompletions
 instance NFData   NonLocalCompletions
 instance Binary   NonLocalCompletions
+
 -- | Generate code actions.
 getCompletionsLSP
-    :: LSP.LspFuncs Config
+    :: PluginId
+    -> LSP.LspFuncs Config
     -> IdeState
     -> CompletionParams
     -> IO (Either ResponseError CompletionResponseResult)
-getCompletionsLSP lsp ide
+getCompletionsLSP plId lsp ide
   CompletionParams{_textDocument=TextDocumentIdentifier uri
                   ,_position=position
                   ,_context=completionContext} = do
@@ -145,8 +156,53 @@ getCompletionsLSP lsp ide
                 let clientCaps = clientCapabilities $ shakeExtras ide
                 config <- getClientConfig lsp
                 let snippets = WithSnippets . completionSnippetsOn $ config
-                allCompletions <- getCompletions ideOpts cci' parsedMod bindMap pfix' clientCaps snippets
+                allCompletions <- getCompletions plId ideOpts cci' parsedMod bindMap pfix' clientCaps snippets
                 pure $ Completions (List allCompletions)
               _ -> return (Completions $ List [])
           _ -> return (Completions $ List [])
       _ -> return (Completions $ List [])
+
+----------------------------------------------------------------------------------------------------
+
+extendImportCommand :: PluginCommand IdeState
+extendImportCommand =
+  PluginCommand (CommandId extendImportCommandId) "additional edits for a completion" extendImportHandler
+
+extendImportHandler :: CommandFunction IdeState ExtendImport
+extendImportHandler _lsp ideState edit = do
+  res <- runMaybeT $ extendImportHandler' ideState edit
+  return (Right Null, res)
+
+extendImportHandler' :: IdeState -> ExtendImport -> MaybeT IO (ServerMethod, ApplyWorkspaceEditParams)
+extendImportHandler' ideState ExtendImport {..}
+  | Just fp <- uriToFilePath doc,
+    nfp <- toNormalizedFilePath' fp =
+    do
+      (ms, ps, imps) <- MaybeT $
+        runAction "extend import" ideState $
+          runMaybeT $ do
+            -- We want accurate edits, so do not use stale data here
+            (ms, imps) <- MaybeT $ use GetModSummaryWithoutTimestamps nfp
+            ps <- MaybeT $ use GetAnnotatedParsedSource nfp
+            return (ms, ps, imps)
+      let df = ms_hspp_opts ms
+          wantedModule = mkModuleName (T.unpack importName)
+      imp <- liftMaybe $ find (isWantedModule wantedModule) imps
+      wedit <-
+        liftEither $
+          rewriteToEdit df doc (annsA ps) $
+            extendImport (T.unpack <$> thingParent) (T.unpack newThing) imp
+      return (WorkspaceApplyEdit, ApplyWorkspaceEditParams wedit)
+  | otherwise =
+    mzero
+
+isWantedModule :: ModuleName -> GenLocated l (ImportDecl pass) -> Bool
+isWantedModule wantedModule (L _ ImportDecl {..}) = unLoc ideclName == wantedModule
+isWantedModule _ _ = False
+
+liftMaybe :: Monad m => Maybe a -> MaybeT m a
+liftMaybe a = MaybeT $ pure a
+
+liftEither :: Monad m => Either e a -> MaybeT m a
+liftEither (Left _) = mzero
+liftEither (Right x) = return x