feat: update type signature during add argument action

plugins/hls-refactor-plugin/src/Development/IDE/GHC/ExactPrint.hs

@@ -23,6 +23,9 @@ module Development.IDE.GHC.ExactPrint
 #if MIN_VERSION_ghc(9,2,1)
       modifySmallestDeclWithM,
       modifyMgMatchesT,
+      modifyMgMatchesT',
+      modifySigWithM,
+      genAnchor1,
 #endif
 #if !MIN_VERSION_ghc(9,2,0)
       Anns,
@@ -111,11 +114,18 @@ import           GHC.Parser.Annotation                   (AnnContext (..),
                                                           deltaPos)
 #endif
 
+#if MIN_VERSION_ghc(9,2,1)
+import Data.List (partition)
+import GHC (Anchor(..), realSrcSpan, AnchorOperation, DeltaPos(..), SrcSpanAnnN)
+import GHC.Types.SrcLoc (generatedSrcSpan)
+import Control.Lens ((&), _last)
+import Control.Lens.Operators ((%~))
+#endif
+
 #if MIN_VERSION_ghc(9,2,0)
 setPrecedingLines :: Default t => LocatedAn t a -> Int -> Int -> LocatedAn t a
 setPrecedingLines ast n c = setEntryDP ast (deltaPos n c)
 #endif
-
 ------------------------------------------------------------------------------
 
 data Log = LogShake Shake.Log deriving Show
@@ -449,32 +459,129 @@ graftDecls dst decs0 = Graft $ \dflags a -> do
 --
 -- For example, if you would like to move a where-clause-defined variable to the same
 -- level as its parent HsDecl, you could use this function.
+--
+-- When matching declaration is found in the sub-declarations of `a`, `Just r` is also returned with the new `a`. If
+-- not declaration matched, then `Nothing` is returned.
 modifySmallestDeclWithM ::
-  forall a m.
+  forall a m r.
   (HasDecls a, Monad m) =>
   (SrcSpan -> m Bool) ->
-  (LHsDecl GhcPs -> TransformT m [LHsDecl GhcPs]) ->
+  (LHsDecl GhcPs -> TransformT m ([LHsDecl GhcPs], r)) ->
   a ->
-  TransformT m a
+  TransformT m (a, Maybe r)
 modifySmallestDeclWithM validSpan f a = do
-  let modifyMatchingDecl [] = pure DL.empty
-      modifyMatchingDecl (e@(L src _) : rest) =
+  let modifyMatchingDecl [] = pure (DL.empty, Nothing)
+      modifyMatchingDecl (ldecl@(L src _) : rest) =
         lift (validSpan $ locA src) >>= \case
             True -> do
-              decs' <- f e
-              pure $ DL.fromList decs' <> DL.fromList rest
-            False -> (DL.singleton e <>) <$> modifyMatchingDecl rest
-  modifyDeclsT (fmap DL.toList . modifyMatchingDecl) a
-
--- | Modify the each LMatch in a MatchGroup
+              (decs', r) <- f ldecl
+              pure $ (DL.fromList decs' <> DL.fromList rest, Just r)
+            False -> first (DL.singleton ldecl <>) <$> modifyMatchingDecl rest
+  modifyDeclsT' (fmap (first DL.toList) . modifyMatchingDecl) a
+
+generatedAnchor :: AnchorOperation -> Anchor
+generatedAnchor anchorOp = GHC.Anchor (GHC.realSrcSpan generatedSrcSpan) anchorOp
+
+setAnchor :: Anchor -> SrcSpanAnnN -> SrcSpanAnnN
+setAnchor anc (SrcSpanAnn (EpAnn _ nameAnn comments) span) =
+  SrcSpanAnn (EpAnn anc nameAnn comments) span
+setAnchor _ spanAnnN = spanAnnN
+
+removeTrailingAnns :: SrcSpanAnnN -> SrcSpanAnnN
+removeTrailingAnns (SrcSpanAnn (EpAnn anc nameAnn comments) span) =
+  let nameAnnSansTrailings = nameAnn {nann_trailing = []}
+  in SrcSpanAnn (EpAnn anc nameAnnSansTrailings comments) span
+removeTrailingAnns spanAnnN = spanAnnN
+
+-- | Modify the type signature for the given IdP. This function handles splitting a multi-sig
+-- SigD into multiple SigD if the type signature is changed.
+--
+-- For example, update the type signature for `foo` from `Int` to `Bool`:
+--
+-- - foo :: Int
+-- + foo :: Bool
+--
+-- - foo, bar :: Int
+-- + bar :: Int
+-- + foo :: Bool
+--
+-- - foo, bar, baz :: Int
+-- + bar, baz :: Int
+-- + foo :: Bool
+modifySigWithM ::
+  forall a m.
+  (HasDecls a, Monad m) =>
+  IdP GhcPs ->
+  (LHsSigType GhcPs -> LHsSigType GhcPs) ->
+  a ->
+  TransformT m a
+modifySigWithM queryId f a = do
+  let modifyMatchingSigD :: [LHsDecl GhcPs] -> TransformT m (DL.DList (LHsDecl GhcPs))
+      modifyMatchingSigD [] = pure (DL.empty)
+      modifyMatchingSigD (ldecl@(L annSigD (SigD xsig (TypeSig xTypeSig ids (HsWC xHsWc lHsSig)))) : rest)
+        | queryId `elem` (unLoc <$> ids) = do
+            let newSig = f lHsSig
+            -- If this signature update caused no change, then we don't need to split up multi-signatures
+            if newSig `geq` lHsSig
+              then pure $ DL.singleton ldecl <> DL.fromList rest
+              else case partition ((== queryId) . unLoc) ids of
+                ([L annMatchedId matchedId], otherIds) ->
+                  let matchedId' = L (setAnchor genAnchor0 $ removeTrailingAnns annMatchedId) matchedId
+                      matchedIdSig =
+                        let sig' = SigD xsig (TypeSig xTypeSig [matchedId'] (HsWC xHsWc newSig))
+                            epAnn = bool (noAnnSrcSpanDP generatedSrcSpan (DifferentLine 1 0)) annSigD (null otherIds)
+                        in L epAnn sig'
+                      otherSig = case otherIds of
+                        [] -> []
+                        (L (SrcSpanAnn epAnn span) id1:ids) -> [
+                          let epAnn' = case epAnn of
+                                EpAnn _ nameAnn commentsId1 -> EpAnn genAnchor0 nameAnn commentsId1
+                                EpAnnNotUsed -> EpAnn genAnchor0 mempty emptyComments
+                              ids' = L (SrcSpanAnn epAnn' span) id1:ids
+                              ids'' = ids' & _last %~ first removeTrailingAnns
+                            in L annSigD (SigD xsig (TypeSig xTypeSig ids'' (HsWC xHsWc lHsSig)))
+                            ]
+                  in pure $ DL.fromList otherSig <> DL.singleton matchedIdSig <> DL.fromList rest
+                _ -> error "multiple ids matched"
+      modifyMatchingSigD (ldecl : rest) = (DL.singleton ldecl <>) <$> modifyMatchingSigD rest
+  modifyDeclsT (fmap DL.toList . modifyMatchingSigD) a
+
+genAnchor0 :: Anchor
+genAnchor0 = generatedAnchor m0
+
+genAnchor1 :: Anchor
+genAnchor1 = generatedAnchor m1
+
+-- | Apply a transformation to the decls contained in @t@
+modifyDeclsT' :: (HasDecls t, HasTransform m)
+             => ([LHsDecl GhcPs] -> m ([LHsDecl GhcPs], r))
+             -> t -> m (t, r)
+modifyDeclsT' action t = do
+  decls <- liftT $ hsDecls t
+  (decls', r) <- action decls
+  t' <- liftT $ replaceDecls t decls'
+  pure (t', r)
+
+-- | Modify each LMatch in a MatchGroup
 modifyMgMatchesT ::
   Monad m =>
   MatchGroup GhcPs (LHsExpr GhcPs) ->
   (LMatch GhcPs (LHsExpr GhcPs) -> TransformT m (LMatch GhcPs (LHsExpr GhcPs))) ->
   TransformT m (MatchGroup GhcPs (LHsExpr GhcPs))
-modifyMgMatchesT (MG xMg (L locMatches matches) originMg) f = do
-  matches' <- mapM f matches
-  pure $ MG xMg (L locMatches matches') originMg
+modifyMgMatchesT mg f = fst <$> modifyMgMatchesT' mg (fmap (, ()) . f) () ((.) pure . const)
+
+-- | Modify the each LMatch in a MatchGroup
+modifyMgMatchesT' ::
+  Monad m =>
+  MatchGroup GhcPs (LHsExpr GhcPs) ->
+  (LMatch GhcPs (LHsExpr GhcPs) -> TransformT m (LMatch GhcPs (LHsExpr GhcPs), r)) ->
+  r ->
+  (r -> r -> m r) ->
+  TransformT m (MatchGroup GhcPs (LHsExpr GhcPs), r)
+modifyMgMatchesT' (MG xMg (L locMatches matches) originMg) f def combineResults = do
+  (unzip -> (matches', rs)) <- mapM f matches
+  r' <- lift $ foldM combineResults def rs
+  pure $ (MG xMg (L locMatches matches') originMg, r')
 #endif
 
 graftSmallestDeclsWithM ::

plugins/hls-refactor-plugin/src/Development/IDE/Plugin/CodeAction.hs

@@ -93,12 +93,14 @@ import qualified Text.Fuzzy.Parallel                               as TFP
 import           Text.Regex.TDFA                                   (mrAfter,
                                                                     (=~), (=~~))
 #if MIN_VERSION_ghc(9,2,1)
+import           Data.Either.Extra                                 (maybeToEither)
 import           GHC.Types.SrcLoc                                  (generatedSrcSpan)
 import           Language.Haskell.GHC.ExactPrint                   (noAnnSrcSpanDP1,
                                                                     runTransformT)
 #endif
 #if MIN_VERSION_ghc(9,2,0)
-import           Extra                                             (maybeToEither)
+import           Control.Monad.Except                              (lift)
+import           Debug.Trace
 import           GHC                                               (AddEpAnn (AddEpAnn),
                                                                     Anchor (anchor_op),
                                                                     AnchorOperation (..),
@@ -107,7 +109,17 @@ import           GHC                                               (AddEpAnn (Ad
                                                                     EpAnn (..),
                                                                     EpaLocation (..),
                                                                     LEpaComment,
-                                                                    LocatedA)
+                                                                    LocatedA,
+                                                                    SrcSpanAnn' (SrcSpanAnn),
+                                                                    SrcSpanAnnA,
+                                                                    SrcSpanAnnN,
+                                                                    TrailingAnn (..),
+                                                                    addTrailingAnnToA,
+                                                                    emptyComments,
+                                                                    noAnn)
+import           GHC.Hs                                            (IsUnicodeSyntax (..))
+import           Language.Haskell.GHC.ExactPrint.Transform         (d1)
+
 #else
 import           Language.Haskell.GHC.ExactPrint.Types             (Annotation (annsDP),
                                                                     DeltaPos,
@@ -958,8 +970,6 @@ newDefinitionAction IdeOptions {..} parsedModule Range {_start} name typ
 -- When we receive either of these errors, we produce a text edit that will add a new argument (as a new pattern in the
 -- last position of each LHS of the top-level bindings for this HsDecl).
 --
--- TODO Include logic to also update the type signature of a binding
---
 -- NOTE When adding a new argument to a declaration, the corresponding argument's type in declaration's signature might
 --      not be the last type in the signature, such as:
 --         foo :: a -> b -> c -> d
@@ -973,31 +983,100 @@ suggestAddArgument parsedModule Diagnostic {_message, _range}
   where
     message = unifySpaces _message
 
--- TODO use typ to modify type signature
+-- Given a name for the new binding, add a new pattern to the match in the last position,
+-- returning how many patterns there were in this match prior to the transformation:
+--      addArgToMatch "foo" `bar arg1 arg2 = ...`
+--   => (`bar arg1 arg2 foo = ...`, 2)
+addArgToMatch :: T.Text -> GenLocated l (Match GhcPs body) -> (GenLocated l (Match GhcPs body), Int)
+addArgToMatch name (L locMatch (Match xMatch ctxMatch pats rhs)) =
+  let unqualName = mkRdrUnqual $ mkVarOcc $ T.unpack name
+      newPat = L (noAnnSrcSpanDP1 generatedSrcSpan) $ VarPat NoExtField (noLocA unqualName)
+  in (L locMatch (Match xMatch ctxMatch (pats <> [newPat]) rhs), length pats)
+
+-- Attempt to insert a binding pattern into each match for the given LHsDecl; succeeds only if the function is a FunBind.
+-- Also return:
+--   - the declaration's name
+--   - the number of bound patterns in the declaration's matches prior to the transformation
+--
+-- For example:
+--    insertArg "new_pat" `foo bar baz = 1`
+-- => (`foo bar baz new_pat = 1`, Just ("foo", 2))
+appendFinalPatToMatches :: T.Text -> LHsDecl GhcPs -> TransformT (Either ResponseError) (LHsDecl GhcPs, Maybe (GenLocated SrcSpanAnnN RdrName, Int))
+appendFinalPatToMatches name = \case
+  (L locDecl (ValD xVal (FunBind xFunBind idFunBind mg coreFunBind))) -> do
+    (mg', numPatsMay) <- modifyMgMatchesT' mg (pure . second Just . addArgToMatch name) Nothing combineMatchNumPats
+    numPats <- lift $ maybeToEither (responseError "Unexpected empty match group in HsDecl") numPatsMay
+    let decl' = L locDecl (ValD xVal (FunBind xFunBind idFunBind mg' coreFunBind))
+    pure (decl', Just (idFunBind, numPats))
+  decl -> pure (decl, Nothing)
+  where
+    combineMatchNumPats  Nothing other = pure other
+    combineMatchNumPats  other Nothing = pure other
+    combineMatchNumPats  (Just l) (Just r)
+      | l == r = pure (Just l)
+      | otherwise = Left $ responseError "Unexpected different numbers of patterns in HsDecl MatchGroup"
+
+-- The add argument works as follows:
+--  1. Attempt to add the given name as the last pattern of the declaration that contains `range`.
+--  2. If such a declaration exists, use that declaration's name to modify the signature of said declaration, if it
+--     has a type signature.
+--
+-- NOTE For the following situation, the type signature is not updated (it's unclear what should happen):
+--   type FunctionTySyn = () -> Int
+--   foo :: FunctionTySyn
+--   foo () = new_def
+--
+-- TODO instead of inserting a typed hole; use GHC's suggested type from the error
 addArgumentAction :: ParsedModule -> Range -> T.Text -> Maybe T.Text -> Either ResponseError [(T.Text, [TextEdit])]
-addArgumentAction (ParsedModule _ parsedSource _ _) range name _typ =
-  do
-    let addArgToMatch (L locMatch (Match xMatch ctxMatch pats rhs)) = do
-          let unqualName = mkRdrUnqual $ mkVarOcc $ T.unpack name
-          let newPat = L (noAnnSrcSpanDP1 generatedSrcSpan) $ VarPat NoExtField (noLocA unqualName)
-          pure $ L locMatch (Match xMatch ctxMatch (pats <> [newPat]) rhs)
-        insertArg = \case
-          (L locDecl (ValD xVal (FunBind xFunBind idFunBind mg coreFunBind))) -> do
-            mg' <- modifyMgMatchesT mg addArgToMatch
-            let decl' = L locDecl (ValD xVal (FunBind xFunBind idFunBind mg' coreFunBind))
-            pure [decl']
-          decl -> pure [decl]
-    case runTransformT $ modifySmallestDeclWithM spanContainsRangeOrErr insertArg (makeDeltaAst parsedSource) of
-      Left err -> Left err
-      Right (newSource, _, _) ->
-        let diff = makeDiffTextEdit (T.pack $ exactPrint parsedSource) (T.pack $ exactPrint newSource)
-         in pure [("Add argument ‘" <> name <> "’ to function", fromLspList diff)]
-      where
-        spanContainsRangeOrErr = maybeToEither (responseError "SrcSpan was not valid range") . (`spanContainsRange` range)
-#endif
+addArgumentAction (ParsedModule _ moduleSrc _ _) range name _typ = do
+    (newSource, _, _) <- runTransformT $ do
+      (moduleSrc', join -> matchedDeclNameMay) <- addNameAsLastArgOfMatchingDecl (makeDeltaAst moduleSrc)
+      case matchedDeclNameMay of
+          Just (matchedDeclName, numPats) -> modifySigWithM (unLoc matchedDeclName) (addTyHoleToTySigArg numPats) moduleSrc'
+          Nothing -> pure moduleSrc'
+    let diff = makeDiffTextEdit (T.pack $ exactPrint moduleSrc) (T.pack $ exactPrint newSource)
+    pure [("Add argument ‘" <> name <> "’ to function", fromLspList diff)]
+  where
+    addNameAsLastArgOfMatchingDecl = modifySmallestDeclWithM spanContainsRangeOrErr addNameAsLastArg
+    addNameAsLastArg = fmap (first (:[])) . appendFinalPatToMatches name
+
+    spanContainsRangeOrErr = maybeToEither (responseError "SrcSpan was not valid range") . (`spanContainsRange` range)
+
+-- Transform an LHsType into a list of arguments and return type, to make transformations easier.
+hsTypeToFunTypeAsList :: LHsType GhcPs -> ([(SrcSpanAnnA, XFunTy GhcPs, HsArrow GhcPs, LHsType GhcPs)], LHsType GhcPs)
+hsTypeToFunTypeAsList = \case
+  L spanAnnA (HsFunTy xFunTy arrow lhs rhs) ->
+    let (rhsArgs, rhsRes) = hsTypeToFunTypeAsList rhs
+    in ((spanAnnA, xFunTy, arrow, lhs):rhsArgs, rhsRes)
+  ty -> ([], ty)
+
+-- The inverse of `hsTypeToFunTypeAsList`
+hsTypeFromFunTypeAsList :: ([(SrcSpanAnnA, XFunTy GhcPs, HsArrow GhcPs, LHsType GhcPs)], LHsType GhcPs) -> LHsType GhcPs
+hsTypeFromFunTypeAsList (args, res) =
+  foldr (\(spanAnnA, xFunTy, arrow, argTy) res -> L spanAnnA $ HsFunTy xFunTy arrow argTy res) res args
+
+-- Add a typed hole to a type signature in the given argument position:
+--   0 `foo :: ()` => foo :: _ -> ()
+--   2 `foo :: FunctionTySyn` => foo :: FunctionTySyn
+--   1 `foo :: () -> () -> Int` => foo :: () -> _ -> () -> Int
+addTyHoleToTySigArg :: Int -> LHsSigType GhcPs -> (LHsSigType GhcPs)
+addTyHoleToTySigArg loc (L annHsSig (HsSig xHsSig tyVarBndrs lsigTy)) =
+    let (args, res) = hsTypeToFunTypeAsList lsigTy
+        wildCardAnn = SrcSpanAnn (EpAnn genAnchor1 (AnnListItem [AddRarrowAnn d1]) emptyComments) generatedSrcSpan
+        newArg = (SrcSpanAnn mempty generatedSrcSpan, noAnn, HsUnrestrictedArrow NormalSyntax, L wildCardAnn $ HsWildCardTy noExtField)
+        -- NOTE if the location that the argument wants to be placed at is not one more than the number of arguments
+        --      in the signature, then we return the original type signature.
+        --      This situation most likely occurs due to a function type synonym in the signature
+        insertArg n _ | n < 0 = error "Not possible"
+        insertArg 0 as = newArg:as
+        insertArg _ [] = []
+        insertArg n (a:as) = a : insertArg (n - 1) as
+        lsigTy' = hsTypeFromFunTypeAsList (insertArg loc args, res)
+    in L annHsSig (HsSig xHsSig tyVarBndrs lsigTy')
 
 fromLspList :: List a -> [a]
 fromLspList (List a) = a
+#endif
 
 suggestFillTypeWildcard :: Diagnostic -> [(T.Text, TextEdit)]
 suggestFillTypeWildcard Diagnostic{_range=_range,..}

plugins/hls-refactor-plugin/test/Main.hs

@@ -54,7 +54,8 @@ import           Test.Tasty.HUnit
 import           Text.Regex.TDFA                          ((=~))
 
 
-import           Development.IDE.Plugin.CodeAction        (matchRegExMultipleImports)
+import           Development.IDE.Plugin.CodeAction        (bindingsPluginDescriptor,
+                                                           matchRegExMultipleImports)
 import           Test.Hls
 
 import           Control.Applicative                      (liftA2)
@@ -2371,10 +2372,38 @@ addFunctionArgumentTests =
         liftIO $ actionTitle @?= "Add argument ‘select’ to function"
         executeCodeAction action
         contentAfterAction <- documentContents docB
-        liftIO $ contentAfterAction @?= T.unlines foo'
+        liftIO $ contentAfterAction @?= T.unlines foo',
+        mkGoldenAddArgTest "AddArgWithSig" (R 1 0 1 50),
+        mkGoldenAddArgTest "AddArgWithSigAndDocs" (R 8 0 8 50),
+        mkGoldenAddArgTest "AddArgFromLet" (R 2 0 2 50),
+        mkGoldenAddArgTest "AddArgFromWhere" (R 3 0 3 50),
+        mkGoldenAddArgTest "AddArgWithTypeSynSig" (R 2 0 2 50),
+        mkGoldenAddArgTest "AddArgWithTypeSynSigContravariant" (R 2 0 2 50),
+        mkGoldenAddArgTest "AddArgWithLambda" (R 1 0 1 50),
+        mkGoldenAddArgTest "MultiSigFirst" (R 2 0 2 50),
+        mkGoldenAddArgTest "MultiSigLast" (R 2 0 2 50),
+        mkGoldenAddArgTest "MultiSigMiddle" (R 2 0 2 50)
     ]
 #endif
 
+mkGoldenAddArgTest :: FilePath -> Range -> TestTree
+mkGoldenAddArgTest testFileName range = do
+    let action docB = do
+          _ <- waitForDiagnostics
+          InR action@CodeAction {_title = actionTitle} : _ <-
+            filter (\(InR CodeAction {_title = x}) -> "Add" `isPrefixOf` T.unpack x)
+              <$> getCodeActions docB range
+          liftIO $ actionTitle @?= "Add argument ‘new_def’ to function"
+          executeCodeAction action
+    goldenWithHaskellDoc
+      (Refactor.bindingsPluginDescriptor mempty "ghcide-code-actions-bindings")
+      (testFileName <> " (golden)")
+      "test/data/golden/add-arg"
+      testFileName
+      "expected"
+      "hs"
+      action
+
 deleteUnusedDefinitionTests :: TestTree
 deleteUnusedDefinitionTests = testGroup "delete unused definition action"
   [ testSession "delete unused top level binding" $

plugins/hls-refactor-plugin/test/data/golden/add-arg/AddArgFromLet.expected.hs

@@ -0,0 +1,6 @@
+foo :: Bool -> _ -> Int
+foo True new_def =
+  let bar = new_def
+  in bar
+
+foo False new_def = 1