Support resolve in type lenses (#3743)

* Support resolve in type lenses

* Fix tests

* Dump local signature code

* Update comments

* Address michealpj's suggestions

* Always do position mapping

Author: joyfulmantis

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

@@ -15,37 +15,38 @@ module Development.IDE.Plugin.TypeLenses (
 
 import           Control.Concurrent.STM.Stats         (atomically)
 import           Control.DeepSeq                      (rwhnf)
+import           Control.Lens                         ((?~))
 import           Control.Monad                        (mzero)
 import           Control.Monad.Extra                  (whenMaybe)
 import           Control.Monad.IO.Class               (MonadIO (liftIO))
 import           Control.Monad.Trans.Class            (MonadTrans (lift))
-import           Data.Aeson.Types                     (Value, toJSON)
+import           Data.Aeson.Types                     (toJSON)
 import qualified Data.Aeson.Types                     as A
 import           Data.List                            (find)
 import qualified Data.Map                             as Map
-import           Data.Maybe                           (catMaybes, mapMaybe)
+import           Data.Maybe                           (catMaybes, fromMaybe,
+                                                       maybeToList)
 import qualified Data.Text                            as T
 import           Development.IDE                      (GhcSession (..),
                                                        HscEnvEq (hscEnv),
-                                                       RuleResult, Rules,
+                                                       RuleResult, Rules, Uri,
                                                        define, srcSpanToRange,
                                                        usePropertyAction)
 import           Development.IDE.Core.Compile         (TcModuleResult (..))
 import           Development.IDE.Core.PluginUtils
 import           Development.IDE.Core.PositionMapping (PositionMapping,
+                                                       fromCurrentRange,
                                                        toCurrentRange)
 import           Development.IDE.Core.Rules           (IdeState, runAction)
-import           Development.IDE.Core.RuleTypes       (GetBindings (GetBindings),
-                                                       TypeCheck (TypeCheck))
+import           Development.IDE.Core.RuleTypes       (TypeCheck (TypeCheck))
 import           Development.IDE.Core.Service         (getDiagnostics)
 import           Development.IDE.Core.Shake           (getHiddenDiagnostics,
                                                        use)
 import qualified Development.IDE.Core.Shake           as Shake
 import           Development.IDE.GHC.Compat
 import           Development.IDE.GHC.Util             (printName)
 import           Development.IDE.Graph.Classes
-import           Development.IDE.Spans.LocalBindings  (Bindings, getFuzzyScope)
-import           Development.IDE.Types.Location       (Position (Position, _character, _line),
+import           Development.IDE.Types.Location       (Position (Position, _line),
                                                        Range (Range, _end, _start))
 import           GHC.Generics                         (Generic)
 import           Ide.Logger                           (Pretty (pretty),
@@ -60,38 +61,43 @@ import           Ide.Types                            (CommandFunction,
                                                        PluginDescriptor (..),
                                                        PluginId,
                                                        PluginMethodHandler,
+                                                       ResolveFunction,
                                                        configCustomConfig,
                                                        defaultConfigDescriptor,
                                                        defaultPluginDescriptor,
                                                        mkCustomConfig,
-                                                       mkPluginHandler)
-import           Language.LSP.Protocol.Message        (Method (Method_TextDocumentCodeLens),
+                                                       mkPluginHandler,
+                                                       mkResolveHandler)
+import qualified Language.LSP.Protocol.Lens           as L
+import           Language.LSP.Protocol.Message        (Method (Method_CodeLensResolve, Method_TextDocumentCodeLens),
                                                        SMethod (..))
 import           Language.LSP.Protocol.Types          (ApplyWorkspaceEditParams (ApplyWorkspaceEditParams),
-                                                       CodeLens (CodeLens),
+                                                       CodeLens (..),
                                                        CodeLensParams (CodeLensParams, _textDocument),
-                                                       Diagnostic (..),
+                                                       Command, Diagnostic (..),
                                                        Null (Null),
                                                        TextDocumentIdentifier (TextDocumentIdentifier),
                                                        TextEdit (TextEdit),
                                                        WorkspaceEdit (WorkspaceEdit),
                                                        type (|?) (..))
 import qualified Language.LSP.Server                  as LSP
-import           Text.Regex.TDFA                      ((=~), (=~~))
+import           Text.Regex.TDFA                      ((=~))
 
 data Log = LogShake Shake.Log deriving Show
 
 instance Pretty Log where
   pretty = \case
     LogShake log -> pretty log
 
+
 typeLensCommandId :: T.Text
 typeLensCommandId = "typesignature.add"
 
 descriptor :: Recorder (WithPriority Log) -> PluginId -> PluginDescriptor IdeState
 descriptor recorder plId =
   (defaultPluginDescriptor plId)
     { pluginHandlers = mkPluginHandler SMethod_TextDocumentCodeLens codeLensProvider
+                    <> mkResolveHandler SMethod_CodeLensResolve codeLensResolveProvider
     , pluginCommands = [PluginCommand (CommandId typeLensCommandId) "adds a signature" commandHandler]
     , pluginRules = rules recorder
     , pluginConfigDescriptor = defaultConfigDescriptor {configCustomConfig = mkCustomConfig properties}
@@ -109,97 +115,115 @@ codeLensProvider :: PluginMethodHandler IdeState Method_TextDocumentCodeLens
 codeLensProvider ideState pId CodeLensParams{_textDocument = TextDocumentIdentifier uri} = do
     mode <- liftIO $ runAction "codeLens.config" ideState $ usePropertyAction #mode pId properties
     nfp <- getNormalizedFilePathE uri
-    env <- hscEnv . fst <$>
-      runActionE "codeLens.GhcSession" ideState
-        (useWithStaleE GhcSession nfp)
-
-    (tmr, _) <- runActionE "codeLens.TypeCheck" ideState
-      (useWithStaleE TypeCheck nfp)
-
-    (bindings, _) <- runActionE "codeLens.GetBindings" ideState
-      (useWithStaleE GetBindings nfp)
-
-    (gblSigs@(GlobalBindingTypeSigsResult gblSigs'), gblSigsMp) <-
-      runActionE "codeLens.GetGlobalBindingTypeSigs" ideState
-        (useWithStaleE GetGlobalBindingTypeSigs nfp)
-
-    diag <- liftIO $ atomically $ getDiagnostics ideState
-    hDiag <- liftIO $ atomically $ getHiddenDiagnostics ideState
-
-    let toWorkSpaceEdit tedit = WorkspaceEdit (Just $ Map.singleton uri $ tedit) Nothing Nothing
-        generateLensForGlobal mp sig@GlobalBindingTypeSig{gbRendered} = do
-            range <- toCurrentRange mp =<< srcSpanToRange (gbSrcSpan sig)
-            tedit <- gblBindingTypeSigToEdit sig (Just gblSigsMp)
-            let wedit = toWorkSpaceEdit [tedit]
-            pure $ generateLens pId range (T.pack gbRendered) wedit
-        generateLensFromDiags f =
-              [ generateLens pId _range title edit
-              | (dFile, _, dDiag@Diagnostic{_range = _range}) <- diag ++ hDiag
-              , dFile == nfp
-              , (title, tedit) <- f dDiag
-              , let edit = toWorkSpaceEdit tedit
-              ]
-    -- `suggestLocalSignature` relies on diagnostic, if diagnostics don't have the local signature warning,
-    -- the `bindings` is useless, and if diagnostic has, that means we parsed success, and the `bindings` is fresh.
-    pure $ InL $ case mode of
-        Always ->
-          mapMaybe (generateLensForGlobal gblSigsMp) gblSigs'
-            <> generateLensFromDiags
-                (suggestLocalSignature False (Just env) (Just tmr) (Just bindings)) -- we still need diagnostics for local bindings
-        Exported -> mapMaybe (generateLensForGlobal gblSigsMp) (filter gbExported gblSigs')
-        Diagnostics -> generateLensFromDiags
-            $ suggestSignature False (Just env) (Just gblSigs) (Just tmr) (Just bindings)
-
-generateLens :: PluginId -> Range -> T.Text -> WorkspaceEdit -> CodeLens
-generateLens pId _range title edit =
-  let cId = mkLspCommand pId (CommandId typeLensCommandId) title (Just [toJSON edit])
-   in CodeLens _range (Just cId) Nothing
-
+    -- We have two ways we can possibly generate code lenses for type lenses.
+    -- Different options are with different "modes" of the type-lenses plugin.
+    -- (Remember here, as the code lens is not resolved yet, we only really need
+    -- the range and any data that will help us resolve it later)
+    let -- The first option is to generate lens from diagnostics about
+        -- top level bindings.
+        generateLensFromGlobalDiags diags =
+          -- We don't actually pass any data to resolve, however we need this
+          -- dummy type to make sure HLS resolves our lens
+          [ CodeLens _range Nothing (Just $ toJSON TypeLensesResolve)
+            | (dFile, _, diag@Diagnostic{_range}) <- diags
+            , dFile == nfp
+            , isGlobalDiagnostic diag]
+        -- The second option is to generate lenses from the GlobalBindingTypeSig
+        -- rule. This is the only type that needs to have the range adjusted
+        -- with PositionMapping.
+        -- PositionMapping for diagnostics doesn't make sense, because we always
+        -- have fresh diagnostics even if current module parsed failed (the
+        -- diagnostic would then be parse failed). See
+        -- https://github.com/haskell/haskell-language-server/pull/3558 for this
+        -- discussion.
+        generateLensFromGlobal sigs mp = do
+          [ CodeLens newRange Nothing (Just $ toJSON TypeLensesResolve)
+            | sig <- sigs
+            , Just range <- [srcSpanToRange (gbSrcSpan sig)]
+            , Just newRange <- [toCurrentRange mp range]]
+    if mode == Always || mode == Exported
+      then do
+        -- In this mode we get the global bindings from the
+        -- GlobalBindingTypeSigs rule.
+        (GlobalBindingTypeSigsResult gblSigs, gblSigsMp) <-
+          runActionE "codeLens.GetGlobalBindingTypeSigs" ideState
+          $ useWithStaleE GetGlobalBindingTypeSigs nfp
+        -- Depending on whether we only want exported or not we filter our list
+        -- of signatures to get what we want
+        let relevantGlobalSigs =
+              if mode == Exported
+                then filter gbExported gblSigs
+                else gblSigs
+        pure $ InL $ generateLensFromGlobal relevantGlobalSigs gblSigsMp
+      else do
+        -- For this mode we exclusively use diagnostics to create the lenses.
+        -- However we will still use the GlobalBindingTypeSigs to resolve them.
+        diags <- liftIO $ atomically $ getDiagnostics ideState
+        hDiags <- liftIO $ atomically $ getHiddenDiagnostics ideState
+        let allDiags = diags <> hDiags
+        pure $ InL $ generateLensFromGlobalDiags allDiags
+
+codeLensResolveProvider :: ResolveFunction IdeState TypeLensesResolve Method_CodeLensResolve
+codeLensResolveProvider ideState pId lens@CodeLens{_range} uri TypeLensesResolve = do
+  nfp <- getNormalizedFilePathE uri
+  (gblSigs@(GlobalBindingTypeSigsResult _), pm) <-
+    runActionE "codeLens.GetGlobalBindingTypeSigs" ideState
+    $ useWithStaleE GetGlobalBindingTypeSigs nfp
+  -- regardless of how the original lens was generated, we want to get the range
+  -- that the global bindings rule would expect here, hence the need to reverse
+  -- position map the range, regardless of whether it was position mapped in the
+  -- beginning or freshly taken from diagnostics.
+  newRange <- handleMaybe PluginStaleResolve (fromCurrentRange pm _range)
+  -- We also pass on the PositionMapping so that the generated text edit can
+  -- have the range adjusted.
+  (title, edit) <-
+        handleMaybe PluginStaleResolve $ suggestGlobalSignature' False (Just gblSigs) (Just pm) newRange
+  pure $ lens & L.command ?~ generateLensCommand pId uri title edit
+
+generateLensCommand :: PluginId -> Uri -> T.Text -> TextEdit -> Command
+generateLensCommand pId uri title edit =
+  let wEdit = WorkspaceEdit (Just $ Map.singleton uri $ [edit]) Nothing Nothing
+  in mkLspCommand pId (CommandId typeLensCommandId) title (Just [toJSON wEdit])
+
+-- Since the lenses are created with diagnostics, and since the globalTypeSig
+-- rule can't be changed as it is also used by the hls-refactor plugin, we can't
+-- rely on actions. Because we can't rely on actions it doesn't make sense to
+-- recompute the edit upon command. Hence the command here just takes a edit
+-- and applies it.
 commandHandler :: CommandFunction IdeState WorkspaceEdit
 commandHandler _ideState wedit = do
   _ <- lift $ LSP.sendRequest SMethod_WorkspaceApplyEdit (ApplyWorkspaceEditParams Nothing wedit) (\_ -> pure ())
   pure $ InR Null
 
 --------------------------------------------------------------------------------
+suggestSignature :: Bool -> Maybe GlobalBindingTypeSigsResult -> Diagnostic -> [(T.Text, TextEdit)]
+suggestSignature isQuickFix mGblSigs diag =
+  maybeToList (suggestGlobalSignature isQuickFix mGblSigs diag)
+
+-- The suggestGlobalSignature is separated into two functions. The main function
+-- works with a diagnostic, which then calls the secondary function with
+-- whatever pieces of the diagnostic it needs. This allows the resolve function,
+-- which no longer has the Diagnostic, to still call the secondary functions.
+suggestGlobalSignature :: Bool -> Maybe GlobalBindingTypeSigsResult -> Diagnostic -> Maybe (T.Text, TextEdit)
+suggestGlobalSignature isQuickFix mGblSigs diag@Diagnostic{_range}
+  | isGlobalDiagnostic diag =
+    suggestGlobalSignature' isQuickFix mGblSigs Nothing _range
+  | otherwise = Nothing
 
-suggestSignature :: Bool -> Maybe HscEnv -> Maybe GlobalBindingTypeSigsResult -> Maybe TcModuleResult -> Maybe Bindings -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestSignature isQuickFix env mGblSigs mTmr mBindings diag =
-  suggestGlobalSignature isQuickFix mGblSigs diag <> suggestLocalSignature isQuickFix env mTmr mBindings diag
+isGlobalDiagnostic :: Diagnostic -> Bool
+isGlobalDiagnostic Diagnostic{_message} = _message =~ ("(Top-level binding|Pattern synonym) with no type signature" :: T.Text)
 
-suggestGlobalSignature :: Bool -> Maybe GlobalBindingTypeSigsResult -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestGlobalSignature isQuickFix mGblSigs Diagnostic{_message, _range}
-  | _message
-      =~ ("(Top-level binding|Pattern synonym) with no type signature" :: T.Text)
-    , Just (GlobalBindingTypeSigsResult sigs) <- mGblSigs
-    , Just sig <- find (\x -> sameThing (gbSrcSpan x) _range) sigs
+-- If a PositionMapping is supplied, this function will call
+-- gblBindingTypeSigToEdit with it to create a TextEdit in the right location.
+suggestGlobalSignature' :: Bool -> Maybe GlobalBindingTypeSigsResult -> Maybe PositionMapping -> Range -> Maybe (T.Text, TextEdit)
+suggestGlobalSignature' isQuickFix mGblSigs pm range
+  |   Just (GlobalBindingTypeSigsResult sigs) <- mGblSigs
+    , Just sig <- find (\x -> sameThing (gbSrcSpan x) range) sigs
     , signature <- T.pack $ gbRendered sig
     , title <- if isQuickFix then "add signature: " <> signature else signature
-    , Just action <- gblBindingTypeSigToEdit sig Nothing =
-    [(title, [action])]
-  | otherwise = []
-
-suggestLocalSignature :: Bool -> Maybe HscEnv -> Maybe TcModuleResult -> Maybe Bindings -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestLocalSignature isQuickFix mEnv mTmr mBindings Diagnostic{_message, _range = _range@Range{..}}
-  | Just (_ :: T.Text, _ :: T.Text, _ :: T.Text, [identifier]) <-
-      (T.unwords . T.words $ _message)
-        =~~ ("Polymorphic local binding with no type signature: (.*) ::" :: T.Text)
-    , Just bindings <- mBindings
-    , Just env <- mEnv
-    , localScope <- getFuzzyScope bindings _start _end
-    , -- we can't use srcspan to lookup scoped bindings, because the error message reported by GHC includes the entire binding, instead of simply the name
-      Just (name, ty) <- find (\(x, _) -> printName x == T.unpack identifier) localScope >>= \(name, mTy) -> (name,) <$> mTy
-    , Just TcModuleResult{tmrTypechecked = TcGblEnv{tcg_rdr_env, tcg_sigs}} <- mTmr
-    , -- not a top-level thing, to avoid duplication
-      not $ name `elemNameSet` tcg_sigs
-    , tyMsg <- printSDocQualifiedUnsafe (mkPrintUnqualifiedDefault env tcg_rdr_env) $ pprSigmaType ty
-    , signature <- T.pack $ printName name <> " :: " <> tyMsg
-    , startCharacter <- _character _start
-    , startOfLine <- Position (_line _start) startCharacter
-    , beforeLine <- Range startOfLine startOfLine
-    , title <- if isQuickFix then "add signature: " <> signature else signature
-    , action <- TextEdit beforeLine $ signature <> "\n" <> T.replicate (fromIntegral startCharacter) " " =
-    [(title, [action])]
-  | otherwise = []
+    , Just action <- gblBindingTypeSigToEdit sig pm =
+    Just (title, action)
+  | otherwise = Nothing
 
 sameThing :: SrcSpan -> Range -> Bool
 sameThing s1 s2 = (_start <$> srcSpanToRange s1) == (_start <$> Just s2)
@@ -209,12 +233,20 @@ gblBindingTypeSigToEdit GlobalBindingTypeSig{..} mmp
   | Just Range{..} <- srcSpanToRange $ getSrcSpan gbName
     , startOfLine <- Position (_line _start) 0
     , beforeLine <- Range startOfLine startOfLine
-    -- If `mmp` is `Nothing`, return the original range, it used by lenses from diagnostic,
+    -- If `mmp` is `Nothing`, return the original range,
     -- otherwise we apply `toCurrentRange`, and the guard should fail if `toCurrentRange` failed.
     , Just range <- maybe (Just beforeLine) (flip toCurrentRange beforeLine) mmp
-    = Just $ TextEdit range $ T.pack gbRendered <> "\n"
+    -- We need to flatten the signature, as otherwise long signatures are
+    -- rendered on multiple lines with invalid formatting.
+    , renderedFlat <- unwords $ lines gbRendered
+    = Just $ TextEdit range $ T.pack renderedFlat <> "\n"
   | otherwise = Nothing
 
+-- |We don't need anything to resolve our lens, but a data field is mandatory
+-- to get types resolved in HLS
+data TypeLensesResolve = TypeLensesResolve
+  deriving (Generic, A.FromJSON, A.ToJSON)
+
 data Mode
   = -- | always displays type lenses of global bindings, no matter what GHC flags are set
     Always

ghcide/test/exe/AsyncTests.hs

@@ -35,7 +35,7 @@ tests = testGroup "async"
               , "foo = id"
               ]
             void waitForDiagnostics
-            codeLenses <- getCodeLenses doc
+            codeLenses <- getAndResolveCodeLenses doc
             liftIO $ [ _title | CodeLens{_command = Just Command{_title}} <- codeLenses] @=?
               [ "foo :: a -> a" ]
     , testSession "request" $ do
@@ -47,7 +47,7 @@ tests = testGroup "async"
               , "foo = id"
               ]
             void waitForDiagnostics
-            codeLenses <- getCodeLenses doc
+            codeLenses <- getAndResolveCodeLenses doc
             liftIO $ [ _title | CodeLens{_command = Just Command{_title}} <- codeLenses] @=?
               [ "foo :: a -> a" ]
     ]