Retrieve Type from typecheck result for type lenses

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

@@ -51,6 +51,7 @@ import           Development.IDE.Plugin.CodeAction.ExactPrint
 import           Development.IDE.Plugin.CodeAction.PositionIndexed
 import           Development.IDE.Plugin.TypeLenses                 (suggestSignature)
 import           Development.IDE.Spans.Common
+import           Development.IDE.Spans.LocalBindings               (Bindings)
 import           Development.IDE.Types.Exports
 import           Development.IDE.Types.HscEnvEq
 import           Development.IDE.Types.Location
@@ -97,13 +98,14 @@ codeAction state _ (CodeActionParams _ _ (TextDocumentIdentifier uri) _range Cod
     let text = Rope.toText . (_text :: VirtualFile -> Rope.Rope) <$> contents
         mbFile = toNormalizedFilePath' <$> uriToFilePath uri
     diag <- fmap (\(_, _, d) -> d) . filter (\(p, _, _) -> mbFile == Just p) <$> getDiagnostics state
-    (ideOptions, join -> parsedModule, join -> env, join -> annotatedPS, join -> tcM, join -> har) <- runAction "CodeAction" state $
-      (,,,,,) <$> getIdeOptions
+    (ideOptions, join -> parsedModule, join -> env, join -> annotatedPS, join -> tcM, join -> har, join -> bindings) <- runAction "CodeAction" state $
+      (,,,,,,) <$> getIdeOptions
             <*> getParsedModule `traverse` mbFile
             <*> use GhcSession `traverse` mbFile
             <*> use GetAnnotatedParsedSource `traverse` mbFile
             <*> use TypeCheck `traverse` mbFile
             <*> use GetHieAst `traverse` mbFile
+            <*> use GetBindings `traverse` mbFile
     -- This is quite expensive 0.6-0.7s on GHC
     pkgExports   <- maybe mempty envPackageExports env
     localExports <- readVar (exportsMap $ shakeExtras state)
@@ -112,7 +114,7 @@ codeAction state _ (CodeActionParams _ _ (TextDocumentIdentifier uri) _range Cod
       df = ms_hspp_opts . pm_mod_summary <$> parsedModule
       actions =
         [ mkCA title  [x] edit
-        | x <- xs, (title, tedit) <- suggestAction exportsMap ideOptions parsedModule text df annotatedPS tcM har x
+        | x <- xs, (title, tedit) <- suggestAction exportsMap ideOptions parsedModule text df annotatedPS tcM har bindings x
         , let edit = WorkspaceEdit (Just $ Map.singleton uri $ List tedit) Nothing
         ]
       actions' = caRemoveRedundantImports parsedModule text diag xs uri
@@ -144,12 +146,13 @@ suggestAction
   -> Maybe (Annotated ParsedSource)
   -> Maybe TcModuleResult
   -> Maybe HieAstResult
+  -> Maybe Bindings
   -> Diagnostic
   -> [(T.Text, [TextEdit])]
-suggestAction packageExports ideOptions parsedModule text df annSource tcM har diag =
+suggestAction packageExports ideOptions parsedModule text df annSource tcM har bindings diag =
     concat
    -- Order these suggestions by priority
-    [ suggestSignature True diag
+    [ suggestSignature True tcM bindings diag
     , rewrite df annSource $ \_ ps -> suggestExtendImport packageExports ps diag
     , rewrite df annSource $ \df ps ->
         suggestImportDisambiguation df text ps diag

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

@@ -6,36 +6,53 @@ module Development.IDE.Plugin.TypeLenses
   )
 where
 
+import           ConLike                             (ConLike (PatSynCon))
+import           Control.Applicative                 ((<|>))
 import           Control.Monad.IO.Class
-import           Data.Aeson.Types               (Value (..), toJSON)
-import qualified Data.HashMap.Strict            as Map
-import qualified Data.Text                      as T
-import           Development.IDE.Core.RuleTypes (TypeCheck (TypeCheck))
-import           Development.IDE.Core.Rules     (IdeState, runAction)
-import           Development.IDE.Core.Service   (getDiagnostics)
-import           Development.IDE.Core.Shake     (getHiddenDiagnostics, use)
-import           Development.IDE.Types.Location (Position (Position, _character, _line),
-                                                 Range (Range, _end, _start),
-                                                 toNormalizedFilePath',
-                                                 uriToFilePath')
-import           Ide.PluginUtils                (mkLspCommand)
-import           Ide.Types                      (CommandFunction,
-                                                 CommandId (CommandId),
-                                                 PluginCommand (PluginCommand),
-                                                 PluginDescriptor (..),
-                                                 PluginId,
-                                                 defaultPluginDescriptor,
-                                                 mkPluginHandler)
-import qualified Language.LSP.Server            as LSP
-import           Language.LSP.Types             (ApplyWorkspaceEditParams (ApplyWorkspaceEditParams),
-                                                 CodeLens (CodeLens),
-                                                 CodeLensParams (CodeLensParams, _textDocument),
-                                                 Diagnostic (..), List (..),
-                                                 ResponseError, SMethod (..),
-                                                 TextDocumentIdentifier (TextDocumentIdentifier),
-                                                 TextEdit (TextEdit),
-                                                 WorkspaceEdit (WorkspaceEdit))
-import           Text.Regex.TDFA                ((=~))
+import           Data.Aeson.Types                    (Value (..), toJSON)
+import           Data.Char                           (isAlpha)
+import           Data.Generics                       (mkQ, something)
+import qualified Data.HashMap.Strict                 as Map
+import           Data.List                           (find)
+import qualified Data.Text                           as T
+import           Development.IDE.Core.Compile        (TcModuleResult (..))
+import           Development.IDE.Core.RuleTypes      (GetBindings (GetBindings),
+                                                      TypeCheck (TypeCheck))
+import           Development.IDE.Core.Rules          (IdeState, runAction)
+import           Development.IDE.Core.Service        (getDiagnostics)
+import           Development.IDE.Core.Shake          (getHiddenDiagnostics, use)
+import           Development.IDE.GHC.Compat
+import           Development.IDE.GHC.Util            (printName)
+import           Development.IDE.Spans.Common        (safeTyThingType)
+import           Development.IDE.Spans.LocalBindings (Bindings, getFuzzyScope)
+import           Development.IDE.Types.Location      (Position (Position, _character, _line),
+                                                      Range (Range, _end, _start),
+                                                      toNormalizedFilePath',
+                                                      uriToFilePath')
+import           HscTypes                            (lookupTypeEnv,
+                                                      mkPrintUnqualified)
+import           Ide.PluginUtils                     (mkLspCommand)
+import           Ide.Types                           (CommandFunction,
+                                                      CommandId (CommandId),
+                                                      PluginCommand (PluginCommand),
+                                                      PluginDescriptor (..),
+                                                      PluginId,
+                                                      defaultPluginDescriptor,
+                                                      mkPluginHandler)
+import qualified Language.LSP.Server                 as LSP
+import           Language.LSP.Types                  (ApplyWorkspaceEditParams (ApplyWorkspaceEditParams),
+                                                      CodeLens (CodeLens),
+                                                      CodeLensParams (CodeLensParams, _textDocument),
+                                                      Diagnostic (..),
+                                                      List (..), ResponseError,
+                                                      SMethod (..),
+                                                      TextDocumentIdentifier (TextDocumentIdentifier),
+                                                      TextEdit (TextEdit),
+                                                      WorkspaceEdit (WorkspaceEdit))
+import           Outputable                          (showSDocForUser)
+import           TcRnTypes                           (TcGblEnv (TcGblEnv, tcg_rdr_env, tcg_rn_decls, tcg_type_env))
+import           TcType                              (pprSigmaType)
+import           Text.Regex.TDFA                     ((=~))
 
 typeLensCommandId :: T.Text
 typeLensCommandId = "typesignature.add"
@@ -55,14 +72,15 @@ codeLensProvider ::
 codeLensProvider ideState pId CodeLensParams {_textDocument = TextDocumentIdentifier uri} = do
   fmap (Right . List) $ case uriToFilePath' uri of
     Just (toNormalizedFilePath' -> filePath) -> liftIO $ do
-      _ <- runAction "codeLens" ideState (use TypeCheck filePath)
+      tmr <- runAction "codeLens.TypeCheck" ideState (use TypeCheck filePath)
+      bindings <- runAction "bindings.GetBindings" ideState (use GetBindings filePath)
       diag <- getDiagnostics ideState
       hDiag <- getHiddenDiagnostics ideState
       sequence
         [ generateLens pId _range title edit
           | (dFile, _, dDiag@Diagnostic {_range = _range}) <- diag ++ hDiag,
             dFile == filePath,
-            (title, tedit) <- suggestSignature False dDiag,
+            (title, tedit) <- suggestSignature False tmr bindings dDiag,
             let edit = WorkspaceEdit (Just $ Map.singleton uri $ List tedit) Nothing
         ]
     Nothing -> pure []
@@ -77,33 +95,38 @@ commandHandler _ideState wedit = do
   _ <- LSP.sendRequest SWorkspaceApplyEdit (ApplyWorkspaceEditParams Nothing wedit) (\_ -> pure ())
   return $ Right Null
 
-suggestSignature :: Bool -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestSignature isQuickFix Diagnostic {_range = _range@Range {..}, ..}
+suggestSignature :: Bool -> Maybe TcModuleResult -> Maybe Bindings -> Diagnostic -> [(T.Text, [TextEdit])]
+suggestSignature isQuickFix mTmr mBindings Diagnostic {_message, _range = Range {..}}
   | _message
-      =~ ("(Top-level binding|Polymorphic local binding|Pattern synonym) with no type signature" :: T.Text) =
-    let signature =
-          removeInitialForAll $
-            T.takeWhile (\x -> x /= '*' && x /= '•') $
-              T.strip $ unifySpaces $ last $ T.splitOn "type signature: " $ filterNewlines _message
-        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 startCharacter " "
-     in [(title, [action])]
-  where
-    removeInitialForAll :: T.Text -> T.Text
-    removeInitialForAll (T.breakOnEnd " :: " -> (nm, ty))
-      | "forall" `T.isPrefixOf` ty = nm <> T.drop 2 (snd (T.breakOn "." ty))
-      | otherwise = nm <> ty
-    startCharacter
-      | "Polymorphic local binding" `T.isPrefixOf` _message =
-        _character _start
-      | otherwise =
-        0
-suggestSignature _ _ = []
+      =~ ("(Top-level binding|Polymorphic local binding|Pattern synonym) with no type signature" :: T.Text),
+    Just bindings <- mBindings,
+    Just TcModuleResult {tmrTypechecked = TcGblEnv {tcg_type_env, tcg_rn_decls, tcg_rdr_env}} <- mTmr,
+    localScope <- getFuzzyScope bindings _start _end,
+    Just group <- tcg_rn_decls,
+    Just name <- getFirstIdAtLine (succ $ _line _start) group,
+    Just (isPatSyn, ty) <-
+      (lookupTypeEnv tcg_type_env name >>= \x -> (isTyThingPatSyn x,) <$> safeTyThingType x)
+        <|> ((False,) <$> (find (\(x, _) -> x == name) localScope >>= snd)),
+    tyMsg <- showSDocForUser unsafeGlobalDynFlags (mkPrintUnqualified unsafeGlobalDynFlags tcg_rdr_env) $ pprSigmaType ty,
+    signature <- T.pack $ (if isPatSyn then "pattern " else "") <> printName name <> " :: " <> tyMsg,
+    startCharacter <- if "Polymorphic local" `T.isPrefixOf` T.dropWhile (not . isAlpha) _message then _character _start else 0,
+    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 startCharacter " " =
+    [(title, [action])]
+  | otherwise = []
 
-unifySpaces :: T.Text -> T.Text
-unifySpaces = T.unwords . T.words
+isTyThingPatSyn :: TyThing -> Bool
+isTyThingPatSyn (AConLike (PatSynCon _)) = True
+isTyThingPatSyn _                        = False
 
-filterNewlines :: T.Text -> T.Text
-filterNewlines = T.concat . T.lines
+getFirstIdAtLine :: Int -> HsGroup GhcRn -> Maybe Name
+getFirstIdAtLine line = something (mkQ Nothing f)
+  where
+    f :: Located Name -> Maybe Name
+    f (L l name)
+      | RealSrcSpan s <- l,
+        srcSpanStartLine s == line =
+        Just name
+      | otherwise = Nothing

ghcide/src/Development/IDE/Spans/Common.hs

@@ -25,7 +25,6 @@ import qualified Data.Text                    as T
 import           GHC.Generics
 
 import           ConLike
-import           DataCon
 import           DynFlags
 import           GHC
 import           NameEnv
@@ -66,9 +65,9 @@ safeTyThingType (ATyCon tycon)    = Just (tyConKind tycon)
 safeTyThingType _                 = Nothing
 
 safeTyThingId :: TyThing -> Maybe Id
-safeTyThingId (AnId i)                    = Just i
-safeTyThingId (AConLike (RealDataCon dc)) = Just $ dataConWrapId dc
-safeTyThingId _                           = Nothing
+safeTyThingId (AnId i)           = Just i
+safeTyThingId (AConLike conLike) = conLikeWrapId_maybe conLike
+safeTyThingId _                  = Nothing
 
 -- Possible documentation for an element in the code
 data SpanDoc

ghcide/test/exe/Main.hs

@@ -3157,7 +3157,13 @@ addSigLensesTests :: TestTree
 addSigLensesTests = let
   missing = "{-# OPTIONS_GHC -Wmissing-signatures -Wmissing-pattern-synonym-signatures -Wunused-matches #-}"
   notMissing = "{-# OPTIONS_GHC -Wunused-matches #-}"
-  moduleH = "{-# LANGUAGE PatternSynonyms #-}\nmodule Sigs where\nimport qualified Data.Complex as C"
+  moduleH = T.unlines
+    [
+    "{-# LANGUAGE PatternSynonyms,TypeApplications,DataKinds,RankNTypes,ScopedTypeVariables,TypeOperators #-}"
+    , "module Sigs where"
+    , "import qualified Data.Complex as C"
+    , "import Data.Data (Proxy (..), type (:~:) (..), mkCharType)"
+    ]
   other = T.unlines ["f :: Integer -> Integer", "f x = 3"]
   before  withMissing def
     = T.unlines $ (if withMissing then (missing :) else (notMissing :)) [moduleH, def, other]
@@ -3183,6 +3189,13 @@ addSigLensesTests = let
       , sigSession enableWarnings "a `haha` b = a b"        "haha :: (t1 -> t2) -> t1 -> t2"
       , sigSession enableWarnings "pattern Some a = Just a" "pattern Some :: a -> Maybe a"
       , sigSession enableWarnings "qualifiedSigTest= C.realPart" "qualifiedSigTest :: C.Complex a -> a"
+      , sigSession enableWarnings "head = 233"              "head :: Integer"
+      , sigSession enableWarnings "rank2Test (k :: forall a . a -> a) = (k 233 :: Int, k \"QAQ\")"
+         "rank2Test :: (forall a. a -> a) -> (Int, [Char])"
+      , sigSession enableWarnings "symbolKindTest = Proxy @\"qwq\""   "symbolKindTest :: Proxy \"qwq\""
+      , sigSession enableWarnings "promotedKindTest = Proxy @Nothing" "promotedKindTest :: Proxy 'Nothing"
+      , sigSession enableWarnings "typeOperatorTest = Refl"           "typeOperatorTest :: a :~: a"
+      , sigSession enableWarnings "notInScopeTest = mkCharType"       "notInScopeTest :: String -> Data.Data.DataType"
       ]
       | (title, enableWarnings) <-
         [("with warnings enabled", True)