Use exact print for suggest missing constraint code actions

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

@@ -32,6 +32,7 @@ import Development.IDE.Core.Shake
 import Development.IDE.GHC.Error
 import Development.IDE.GHC.ExactPrint
 import Development.IDE.LSP.Server
+import Development.IDE.Plugin.CodeAction.ExactPrint
 import Development.IDE.Plugin.CodeAction.PositionIndexed
 import Development.IDE.Plugin.CodeAction.RuleTypes
 import Development.IDE.Plugin.CodeAction.Rules
@@ -53,7 +54,7 @@ import Data.List.NonEmpty (NonEmpty((:|)))
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Text as T
 import Text.Regex.TDFA (mrAfter, (=~), (=~~))
-import Outputable (ppr, showSDocUnsafe)
+import Outputable (Outputable, ppr, showSDoc, showSDocUnsafe)
 import Data.Function
 import Control.Arrow ((>>>))
 import Data.Functor
@@ -91,20 +92,38 @@ codeAction lsp state (TextDocumentIdentifier uri) _range CodeActionContext{_diag
     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) <- runAction "CodeAction" state $
-      (,,) <$> getIdeOptions
+    (ideOptions, join -> parsedModule, join -> env, join -> annotatedPS) <- runAction "CodeAction" state $
+      (,,,) <$> getIdeOptions
             <*> getParsedModule `traverse` mbFile
             <*> use GhcSession `traverse` mbFile
+            <*> use GetAnnotatedParsedSource `traverse` mbFile
     -- This is quite expensive 0.6-0.7s on GHC
     pkgExports <- runAction "CodeAction:PackageExports" state $ (useNoFile_ . PackageExports) `traverse` env
     localExports <- readVar (exportsMap $ shakeExtras state)
-    let exportsMap = localExports <> fromMaybe mempty pkgExports
-    pure . Right $
-        [ CACodeAction $ CodeAction title (Just CodeActionQuickFix) (Just $ List [x]) (Just edit) Nothing
+    let
+      exportsMap = localExports <> fromMaybe mempty pkgExports
+      df = ms_hspp_opts . pm_mod_summary <$> parsedModule
+      actions =
+        [ mkCA title  [x] edit
         | x <- xs, (title, tedit) <- suggestAction exportsMap ideOptions parsedModule text x
         , let edit = WorkspaceEdit (Just $ Map.singleton uri $ List tedit) Nothing
         ] <> caRemoveRedundantImports parsedModule text diag xs uri
 
+      actions' =
+          [mkCA title [x] edit
+          | x <- xs
+          , Just ps <- [annotatedPS]
+          , Just dynflags <- [df]
+          , (title, graft) <- suggestExactAction dynflags ps x
+          , let edit = either error id $
+                        rewriteToEdit dynflags uri (annsA ps) graft
+          ]
+    pure $ Right $ actions' <> actions
+
+mkCA :: T.Text -> [Diagnostic] -> WorkspaceEdit -> CAResult
+mkCA title diags edit =
+  CACodeAction $ CodeAction title (Just CodeActionQuickFix) (Just $ List diags) (Just edit) Nothing
+
 -- | Generate code lenses.
 codeLens
     :: LSP.LspFuncs c
@@ -151,6 +170,16 @@ commandHandler lsp _ideState ExecuteCommandParams{..}
     | otherwise
     = return (Right Null, Nothing)
 
+suggestExactAction ::
+  DynFlags ->
+  Annotated ParsedSource ->
+  Diagnostic ->
+  [(T.Text, Rewrite)]
+suggestExactAction df ps x =
+  concat
+    [ suggestConstraint df (astA ps) x
+    ]
+
 suggestAction
   :: ExportsMap
   -> IdeOptions
@@ -169,15 +198,32 @@ suggestAction packageExports ideOptions parsedModule text diag = concat
     , removeRedundantConstraints text diag
     , suggestAddTypeAnnotationToSatisfyContraints text diag
     ] ++ concat
-    [  suggestConstraint pm text diag
-    ++ suggestNewDefinition ideOptions pm text diag
+    [  suggestNewDefinition ideOptions pm text diag
     ++ suggestNewImport packageExports pm diag
     ++ suggestDeleteUnusedBinding pm text diag
     ++ suggestExportUnusedTopBinding text pm diag
     | Just pm <- [parsedModule]
     ] ++
     suggestFillHole diag                   -- Lowest priority
 
+findSigOfDecl :: (IdP p -> Bool) -> [LHsDecl p] -> Maybe (Sig p)
+findSigOfDecl pred decls =
+  listToMaybe
+    [ sig
+      | L _ (SigD _ sig@(TypeSig _ idsSig _)) <- decls,
+        any (pred . unLoc) idsSig
+    ]
+
+findInstanceHead :: (Outputable (HsType p)) => DynFlags -> String -> [LHsDecl p] -> Maybe (LHsType p)
+findInstanceHead df instanceHead decls =
+  listToMaybe
+    [ hsib_body
+      | L _ (InstD _ (ClsInstD _ ClsInstDecl {cid_poly_ty = HsIB {hsib_body}})) <- decls,
+        showSDoc df (ppr hsib_body) == instanceHead
+    ]
+
+findDeclContainingLoc :: Position -> [Located a] -> Maybe (Located a)
+findDeclContainingLoc loc = find (\(L l _) -> loc `isInsideSrcSpan` l)
 
 suggestRemoveRedundantImport :: ParsedModule -> Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
 suggestRemoveRedundantImport ParsedModule{pm_parsed_source = L _  HsModule{hsmodImports}} contents Diagnostic{_range=_range,..}
@@ -210,14 +256,9 @@ caRemoveRedundantImports m contents digs ctxDigs uri
       = caRemoveCtx ++ [caRemoveAll]
   | otherwise = []
   where
-    removeSingle title tedit diagnostic = CACodeAction CodeAction{..} where
+    removeSingle title tedit diagnostic = mkCA title [diagnostic] WorkspaceEdit{..} where
         _changes = Just $ Map.singleton uri $ List tedit
-        _title = title
-        _kind = Just CodeActionQuickFix
-        _diagnostics = Just $ List [diagnostic]
         _documentChanges = Nothing
-        _edit = Just WorkspaceEdit{..}
-        _command = Nothing
     removeAll tedit = CACodeAction CodeAction {..} where
         _changes = Just $ Map.singleton uri $ List tedit
         _title = "Remove all redundant imports"
@@ -687,13 +728,12 @@ suggestSignature isQuickFix Diagnostic{_range=_range@Range{..},..}
 suggestSignature _ _ = []
 
 -- | Suggests a constraint for a declaration for which a constraint is missing.
-suggestConstraint :: ParsedModule -> Maybe T.Text -> Diagnostic -> [(T.Text, [TextEdit])]
-suggestConstraint parsedModule mContents diag@Diagnostic {..}
-  | Just contents <- mContents
-  , Just missingConstraint <- findMissingConstraint _message
+suggestConstraint :: DynFlags -> ParsedSource -> Diagnostic -> [(T.Text, Rewrite)]
+suggestConstraint df parsedModule diag@Diagnostic {..}
+  | Just missingConstraint <- findMissingConstraint _message
   = let codeAction = if _message =~ ("the type signature for:" :: String)
-                        then suggestFunctionConstraint parsedModule
-                        else suggestInstanceConstraint contents
+                        then suggestFunctionConstraint df parsedModule
+                        else suggestInstanceConstraint df parsedModule
      in codeAction diag missingConstraint
   | otherwise = []
     where
@@ -702,59 +742,43 @@ suggestConstraint parsedModule mContents diag@Diagnostic {..}
         let regex = "(No instance for|Could not deduce) \\((.+)\\) arising from a use of"
          in matchRegexUnifySpaces t regex <&> last
 
-normalizeConstraints :: T.Text -> T.Text -> T.Text
-normalizeConstraints existingConstraints constraint =
-  let constraintsInit = if "(" `T.isPrefixOf` existingConstraints
-                           then T.dropEnd 1 existingConstraints
-                           else "(" <> existingConstraints
-   in constraintsInit <> ", " <> constraint <> ")"
-
 -- | Suggests a constraint for an instance declaration for which a constraint is missing.
-suggestInstanceConstraint :: T.Text -> Diagnostic -> T.Text -> [(T.Text, [TextEdit])]
-suggestInstanceConstraint contents Diagnostic {..} missingConstraint
--- Suggests a constraint for an instance declaration with no existing constraints.
--- • No instance for (Eq a) arising from a use of ‘==’
---   Possible fix: add (Eq a) to the context of the instance declaration
--- • In the expression: x == y
---   In an equation for ‘==’: (Wrap x) == (Wrap y) = x == y
---   In the instance declaration for ‘Eq (Wrap a)’
-  | Just [instanceDeclaration] <- matchRegexUnifySpaces _message "In the instance declaration for ‘([^`]*)’"
-  = let instanceLine = contents
-          & T.splitOn ("instance " <> instanceDeclaration)
-          & head & T.lines & length
-        startOfConstraint = Position instanceLine (length ("instance " :: String))
-        range = Range startOfConstraint startOfConstraint
-        newConstraint = missingConstraint <> " => "
-     in [(actionTitle missingConstraint, [TextEdit range newConstraint])]
-
--- Suggests a constraint for an instance declaration with one or more existing constraints.
--- • Could not deduce (Eq b) arising from a use of ‘==’
---   from the context: Eq a
---     bound by the instance declaration at /path/to/Main.hs:7:10-32
---   Possible fix: add (Eq b) to the context of the instance declaration
--- • In the second argument of ‘(&&)’, namely ‘x' == y'’
---   In the expression: x == y && x' == y'
---   In an equation for ‘==’:
---       (Pair x x') == (Pair y y') = x == y && x' == y'
-  | Just [instanceLineStr, constraintFirstCharStr]
-    <- matchRegexUnifySpaces _message "bound by the instance declaration at .+:([0-9]+):([0-9]+)"
-  = let existingConstraints = findExistingConstraints _message
-        newConstraints = normalizeConstraints existingConstraints missingConstraint
-        instanceLine = readPositionNumber instanceLineStr
-        constraintFirstChar = readPositionNumber constraintFirstCharStr
-        startOfConstraint = Position instanceLine constraintFirstChar
-        endOfConstraint = Position instanceLine $
-          constraintFirstChar + T.length existingConstraints
-        range = Range startOfConstraint endOfConstraint
-     in [(actionTitle missingConstraint, [TextEdit range newConstraints])]
+suggestInstanceConstraint :: DynFlags -> ParsedSource -> Diagnostic -> T.Text -> [(T.Text, Rewrite)]
+
+suggestInstanceConstraint df (L _ HsModule {hsmodDecls}) Diagnostic {..} missingConstraint
+  | Just instHead <- instanceHead
+  = [(actionTitle missingConstraint , appendConstraint (T.unpack missingConstraint) instHead)]
   | otherwise = []
     where
-      findExistingConstraints :: T.Text -> T.Text
-      findExistingConstraints t =
-        T.replace "from the context: " "" . T.strip $ T.lines t !! 1
+      instanceHead
+        -- Suggests a constraint for an instance declaration with no existing constraints.
+        -- • No instance for (Eq a) arising from a use of ‘==’
+        --   Possible fix: add (Eq a) to the context of the instance declaration
+        -- • In the expression: x == y
+        --   In an equation for ‘==’: (Wrap x) == (Wrap y) = x == y
+        --   In the instance declaration for ‘Eq (Wrap a)’
+        | Just [instanceDeclaration] <- matchRegexUnifySpaces _message "In the instance declaration for ‘([^`]*)’"
+        , Just instHead <- findInstanceHead df (T.unpack instanceDeclaration) hsmodDecls
+        = Just instHead
+        -- Suggests a constraint for an instance declaration with one or more existing constraints.
+        -- • Could not deduce (Eq b) arising from a use of ‘==’
+        --   from the context: Eq a
+        --     bound by the instance declaration at /path/to/Main.hs:7:10-32
+        --   Possible fix: add (Eq b) to the context of the instance declaration
+        -- • In the second argument of ‘(&&)’, namely ‘x' == y'’
+        --   In the expression: x == y && x' == y'
+        --   In an equation for ‘==’:
+        --       (Pair x x') == (Pair y y') = x == y && x' == y'
+        | Just [instanceLineStr, constraintFirstCharStr]
+            <- matchRegexUnifySpaces _message "bound by the instance declaration at .+:([0-9]+):([0-9]+)"
+        , Just (L _ (InstD _ (ClsInstD _ ClsInstDecl {cid_poly_ty = HsIB{hsib_body}})))
+            <- findDeclContainingLoc (Position (readPositionNumber instanceLineStr) (readPositionNumber constraintFirstCharStr)) hsmodDecls
+        = Just hsib_body
+        | otherwise
+        = Nothing
 
       readPositionNumber :: T.Text -> Int
-      readPositionNumber = T.unpack >>> read >>> pred
+      readPositionNumber = T.unpack >>> read
 
       actionTitle :: T.Text -> T.Text
       actionTitle constraint = "Add `" <> constraint
@@ -768,8 +792,9 @@ findTypeSignatureLine contents typeSignatureName =
   T.splitOn (typeSignatureName <> " :: ") contents & head & T.lines & length
 
 -- | Suggests a constraint for a type signature with any number of existing constraints.
-suggestFunctionConstraint :: ParsedModule -> Diagnostic -> T.Text -> [(T.Text, [TextEdit])]
-suggestFunctionConstraint ParsedModule{pm_parsed_source = L _ HsModule{hsmodDecls}} Diagnostic{..} missingConstraint
+suggestFunctionConstraint :: DynFlags -> ParsedSource -> Diagnostic -> T.Text -> [(T.Text, Rewrite)]
+
+suggestFunctionConstraint df (L _ HsModule {hsmodDecls}) Diagnostic {..} missingConstraint
 -- • No instance for (Eq a) arising from a use of ‘==’
 --   Possible fix:
 --     add (Eq a) to the context of
@@ -792,43 +817,13 @@ suggestFunctionConstraint ParsedModule{pm_parsed_source = L _ HsModule{hsmodDecl
 --   In an equation for ‘eq’:
 --       eq (Pair x y) (Pair x' y') = x == x' && y == y'
   | Just typeSignatureName <- findTypeSignatureName _message
-  = let mExistingConstraints = findExistingConstraints _message
-        newConstraint = buildNewConstraints missingConstraint mExistingConstraints
-     in case findRangeOfContextForFunctionNamed typeSignatureName of
-       Just range -> [(actionTitle missingConstraint typeSignatureName, [TextEdit range newConstraint])]
-       Nothing -> []
-  | otherwise = []
+  , Just (TypeSig _ _ HsWC{hswc_body = HsIB {hsib_body = sig}})
+    <- findSigOfDecl ((T.unpack typeSignatureName ==) . showSDoc df . ppr) hsmodDecls
+  , title <- actionTitle missingConstraint typeSignatureName
+  = [(title, appendConstraint (T.unpack $ missingConstraint) sig)]
+  | otherwise
+  = []
     where
-      findRangeOfContextForFunctionNamed :: T.Text -> Maybe Range
-      findRangeOfContextForFunctionNamed typeSignatureName = do
-          locatedType <- listToMaybe
-              [ locatedType
-              | L _ (SigD _ (TypeSig _ identifiers (HsWC _ (HsIB _ locatedType)))) <- hsmodDecls
-              , any (`isSameName` T.unpack typeSignatureName) $ fmap unLoc identifiers
-              ]
-          let typeBody = dropForAll locatedType
-          srcSpanToRange $ case splitLHsQualTy typeBody of
-            (L contextSrcSpan _ , _) ->
-              if isGoodSrcSpan contextSrcSpan
-                then contextSrcSpan -- The type signature has explicit context
-                else -- No explicit context, return SrcSpan at the start of type (after a potential `forall`)
-                    let start = srcSpanStart $ getLoc typeBody in mkSrcSpan start start
-
-      isSameName :: IdP GhcPs -> String -> Bool
-      isSameName x name = showSDocUnsafe (ppr x) == name
-
-      findExistingConstraints :: T.Text -> Maybe T.Text
-      findExistingConstraints message =
-        if message =~ ("from the context:" :: String)
-           then fmap (T.strip . head) $ matchRegexUnifySpaces message "\\. ([^=]+)"
-           else Nothing
-
-      buildNewConstraints :: T.Text -> Maybe T.Text -> T.Text
-      buildNewConstraints constraint mExistingConstraints =
-        case mExistingConstraints of
-          Just existingConstraints -> normalizeConstraints existingConstraints constraint
-          Nothing -> constraint <> " => "
-
       actionTitle :: T.Text -> T.Text -> T.Text
       actionTitle constraint typeSignatureName = "Add `" <> constraint
         <> "` to the context of the type signature for `" <> typeSignatureName <> "`"

ghcide/test/exe/Main.hs

@@ -2028,7 +2028,7 @@ addFunctionConstraintTests = let
     , ""
     , "data Pair a b = Pair a b"
     , ""
-    , "eq :: " <> constraint <> " => Pair a b -> Pair a b -> Bool"
+    , "eq :: ( " <> constraint <> " ) => Pair a b -> Pair a b -> Bool"
     , "eq (Pair x y) (Pair x' y') = x == x' && y == y'"
     ]
 
@@ -2038,7 +2038,7 @@ addFunctionConstraintTests = let
     [ "module Testing where"
     , "data Pair a b = Pair a b"
     , "eq "
-    , "    :: " <> constraint
+    , "    :: (" <> constraint <> ")"
     , "    => Pair a b -> Pair a b -> Bool"
     , "eq (Pair x y) (Pair x' y') = x == x' && y == y'"
     ]
@@ -2082,13 +2082,13 @@ addFunctionConstraintTests = let
   , check
     "preexisting constraint, with extra spaces in context"
     "Add `Eq b` to the context of the type signature for `eq`"
-    (incompleteConstraintSourceCodeWithExtraCharsInContext "( Eq a )")
-    (incompleteConstraintSourceCodeWithExtraCharsInContext "(Eq a, Eq b)")
+    (incompleteConstraintSourceCodeWithExtraCharsInContext "Eq a")
+    (incompleteConstraintSourceCodeWithExtraCharsInContext "Eq a, Eq b")
   , check
     "preexisting constraint, with newlines in type signature"
     "Add `Eq b` to the context of the type signature for `eq`"
-    (incompleteConstraintSourceCodeWithNewlinesInTypeSignature "(Eq a)")
-    (incompleteConstraintSourceCodeWithNewlinesInTypeSignature "(Eq a, Eq b)")
+    (incompleteConstraintSourceCodeWithNewlinesInTypeSignature "Eq a")
+    (incompleteConstraintSourceCodeWithNewlinesInTypeSignature "Eq a, Eq b")
   ]
 
 removeRedundantConstraintsTests :: TestTree