Draft: Support add to where command

plugins/hls-refactor-plugin/hls-refactor-plugin.cabal

@@ -97,6 +97,7 @@ test-suite tests
   default-language: Haskell2010
   hs-source-dirs:   test
   main-is:          Main.hs
+  other-modules:    Test.AddToWhere
   ghc-options:      -O0 -threaded -rtsopts -with-rtsopts=-N -Wunused-imports
   build-depends:
     , base

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

@@ -47,6 +47,8 @@ module Development.IDE.GHC.ExactPrint
       ExceptStringT (..),
       TransformT,
       Log(..),
+      prependDecl,
+      prependDeclToWhereDecls,
     )
 where
 
@@ -109,6 +111,13 @@ import           GHC.Parser.Annotation                   (AnnContext (..),
                                                           DeltaPos (SameLine),
                                                           EpaLocation (EpaDelta),
                                                           deltaPos)
+import GHC (DeltaPos(..))
+import GHC.Types.SrcLoc (generatedSrcSpan)
+import GHC (Anchor(..))
+import GHC (AnchorOperation(..))
+import GHC (realSrcSpan)
+import GHC (EpAnnComments(..))
+import Debug.Trace (traceShowM, trace, traceM)
 #endif
 
 #if MIN_VERSION_ghc(9,2,0)
@@ -475,6 +484,88 @@ modifyMgMatchesT ::
 modifyMgMatchesT (MG xMg (L locMatches matches) originMg) f = do
   matches' <- mapM f matches
   pure $ MG xMg (L locMatches matches') originMg
+
+insertAtStart' :: (Monad m, HasDecls ast) => ast -> HsDecl GhcPs -> TransformT m ast
+insertAtStart' old newDecl = do
+    srcs <- replicateM 5 uniqueSrcSpanT
+    liftT $ insertAt (insertDeclAtStart srcs) old . noLocA $ newDecl
+  where
+    insertDeclAtStart ssps (L _ newDecl) [] = [L (noAnnSrcSpanDP (ssps !! 0) (DifferentLine 1 4)) newDecl]
+    insertDeclAtStart _ (L _ newDecl) [L (SrcSpanAnn (EpAnn (Anchor dRealSpan (MovedAnchor (SameLine _))) dAnn dComments) dSpan) d] =
+      [ L (noAnnSrcSpanDP generatedSrcSpan (DifferentLine 1 2)) newDecl,
+        L (SrcSpanAnn (EpAnn (Anchor dRealSpan (MovedAnchor $ DifferentLine 1 0)) dAnn dComments) dSpan) d
+      ]
+    insertDeclAtStart ssps (L _ newDecl) (L (SrcSpanAnn (EpAnn (Anchor dRealSpan ancOp) dAnn dComments) dSpan) d:ds) =
+      let newDeclSpan = ssps !! 0
+          newDecl' = L
+            (SrcSpanAnn
+              (EpAnn (Anchor dRealSpan ancOp) mempty emptyComments)
+              newDeclSpan)
+            newDecl
+          secondDecl' = L
+            (SrcSpanAnn
+              (EpAnn
+                (Anchor (realSrcSpan $ newDeclSpan) (MovedAnchor $ DifferentLine 1 0))
+                dAnn
+                dComments)
+              dSpan)
+            d
+      in newDecl' : secondDecl' : ds
+    insertDeclAtStart _ d ds = d : ds
+
+prependDeclToWhereDecls :: (Monad m, HasDecls b) => b -> LHsBindLR GhcPs GhcPs -> TransformT m b
+prependDeclToWhereDecls decl newWhereDecl = do
+  ds <- balanceCommentsList =<< hsDecls decl
+  let ds' = prependDecl (wrapDecl newWhereDecl) ds
+  replaceDecls decl ds'
+
+prependDecl :: LHsDecl GhcPs -> [LHsDecl GhcPs] -> [LHsDecl GhcPs]
+prependDecl ldecl = \case
+   [] -> [setEntryDP ldecl (DifferentLine 1 2)]
+   ld1:lds -> ldecl':ld1'':lds
+     where
+       (ancOp, ld1'') = case ld1 of
+         L (SrcSpanAnn (EpAnn d1Anc d1Ann (EpaComments (L (Anchor c1Rss cAnc) c1:restCs))) ss) d1 ->
+           -- NOTE cannot use setEntryDP to simply assign `DL 1 0` here because when there is no prior decl, the
+           --      DeltaPos on the declaration is absolute instead of relative, and so we must manually update the
+           --      DeltaPos to be relative (since there is about to be a prior declaration).
+           let ld1' = L
+                 (SrcSpanAnn
+                   (EpAnn
+                     (setAnchorDp d1Anc $ DifferentLine 1 0)
+                     d1Ann
+                     (EpaCommentsBalanced (L (Anchor c1Rss $ MovedAnchor $ DifferentLine 1 0) c1:restCs) []))
+                   ss)
+                 d1
+           in (cAnc, ld1')
+         (L (SrcSpanAnn (EpAnn d1Anc d1Ann (EpaCommentsBalanced (L (Anchor c1Rss cAnc) c1:restCs) d1AfterCs)) ss) d1) ->
+           -- NOTE cannot use setEntryDP to simply assign `DL 1 0` here because when there is no prior decl, the
+           --      DeltaPos on the declaration is absolute instead of relative, and so we must manually update the
+           --      DeltaPos to be relative (since there is about to be a prior declaration).
+           let ld1' = L
+                 (SrcSpanAnn
+                   (EpAnn
+                     (setAnchorDp d1Anc $ DifferentLine 1 0)
+                     d1Ann
+                     (EpaCommentsBalanced (L (Anchor c1Rss $ MovedAnchor $ DifferentLine 1 0) c1:restCs) d1AfterCs))
+                   ss)
+                 d1
+           in (cAnc, ld1')
+         (L (SrcSpanAnn (EpAnn (Anchor d1Rss d1AncOp) d1Ann epaCs@(EpaComments [])) ss) d1) ->
+           let ld1' = L (SrcSpanAnn (EpAnn (Anchor d1Rss $ MovedAnchor $ DifferentLine 1 0) d1Ann epaCs) ss) d1
+           in (d1AncOp, ld1')
+         (L (SrcSpanAnn (EpAnn (Anchor d1Rss d1AncOp) d1Ann epaCs@(EpaCommentsBalanced [] _)) ss) d1) ->
+           let ld1' = L (SrcSpanAnn (EpAnn (Anchor d1Rss $ MovedAnchor $ DifferentLine 1 0) d1Ann epaCs) ss) d1
+           in (d1AncOp, ld1')
+         L (SrcSpanAnn EpAnnNotUsed _) _ -> error "Unexpected EpAnnNotUsed"
+       ldecl' = setEntryDP ldecl (maybe (error "what to do with UnchangedAnchor?") id $ getAnchorOpDp ancOp)
+
+setAnchorDp :: Anchor -> DeltaPos -> Anchor
+setAnchorDp (Anchor rss _) dp = Anchor rss (MovedAnchor dp)
+
+getAnchorOpDp :: AnchorOperation -> Maybe DeltaPos
+getAnchorOpDp (MovedAnchor dp) = Just dp
+getAnchorOpDp UnchangedAnchor = Nothing
 #endif
 
 graftSmallestDeclsWithM ::

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

@@ -1,6 +1,6 @@
 -- Copyright (c) 2019 The DAML Authors. All rights reserved.
 -- SPDX-License-Identifier: Apache-2.0
-{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GADTs     #-}
 
 module Development.IDE.Plugin.CodeAction
     (
@@ -92,21 +92,27 @@ import qualified Text.Fuzzy.Parallel                               as TFP
 import           Text.Regex.TDFA                                   (mrAfter,
                                                                     (=~), (=~~))
 #if MIN_VERSION_ghc(9,2,1)
+import           GHC                                               (realSrcSpan)
+import           GHC.Parser.Annotation                             (emptyComments)
 import           GHC.Types.SrcLoc                                  (generatedSrcSpan)
-import           Language.Haskell.GHC.ExactPrint                   (noAnnSrcSpanDP1,
-                                                                    runTransformT)
+import           Language.Haskell.GHC.ExactPrint.Transform
 #endif
 #if MIN_VERSION_ghc(9,2,0)
+import           Control.Monad.Except                              (lift)
+import           Control.Monad.Identity                            (Identity (..))
 import           Extra                                             (maybeToEither)
 import           GHC                                               (AddEpAnn (AddEpAnn),
-                                                                    Anchor (anchor_op),
+                                                                    Anchor (..),
                                                                     AnchorOperation (..),
                                                                     AnnsModule (am_main),
                                                                     DeltaPos (..),
                                                                     EpAnn (..),
                                                                     EpaLocation (..),
                                                                     LEpaComment,
                                                                     LocatedA)
+import Debug.Trace
+import Control.Lens (bimap)
+
 #else
 import           Language.Haskell.GHC.ExactPrint.Types             (Annotation (annsDP),
                                                                     DeltaPos,
@@ -177,7 +183,8 @@ bindingsPluginDescriptor recorder plId = mkExactprintPluginDescriptor recorder $
 #endif
     , wrap suggestNewDefinition
 #if MIN_VERSION_ghc(9,2,1)
-    , wrap suggestAddArgument
+    , wrap (undefinedVariableCodeAction addAsLastArgument)
+    , wrap (undefinedVariableCodeAction addToWhere)
 #endif
     , wrap suggestDeleteUnusedBinding
     ]
@@ -950,12 +957,30 @@ newDefinitionAction IdeOptions {..} parsedModule Range {_start} name typ
     ParsedModule {pm_parsed_source = L _ HsModule {hsmodDecls}} = parsedModule
 
 #if MIN_VERSION_ghc(9,2,1)
+
+type UndefinedVariableCodeAction = ParsedModule -> Diagnostic -> Either ResponseError [(T.Text, [TextEdit])]
+
+type UndefinedVariableHandler = ParsedModule -> Range -> T.Text -> Maybe T.Text -> Either ResponseError [(T.Text, [TextEdit])]
+
 -- When GHC tells us that a variable is not bound, it will tell us either:
 --  - there is an unbound variable with a given type
 --  - there is an unbound variable (GHC provides no type suggestion)
+--  - there is a typed hole `_`
 --
--- 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).
+-- When we receive either of these errors, we can produce a text edit that will put this variable in scope, such as:
+--  - Adding an argument binding
+--  - Adding a binding in a where clause
+--  - ... up to your creativity :)
+undefinedVariableCodeAction :: UndefinedVariableHandler -> UndefinedVariableCodeAction
+undefinedVariableCodeAction handler parsedModule Diagnostic {_message, _range}
+  | Just (name, typ) <- matchVariableNotInScope message = handler parsedModule _range name typ
+  | Just (name, typ) <- matchFoundHoleIncludeUnderscore message = handler parsedModule _range name (Just typ)
+  | otherwise = pure []
+  where
+    message = unifySpaces _message
+
+-- Handle unbound variables by adding 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
 --
@@ -964,17 +989,8 @@ newDefinitionAction IdeOptions {..} parsedModule Range {_start} name typ
 --         foo :: a -> b -> c -> d
 --         foo a b = \c -> ...
 --      In this case a new argument would have to add its type between b and c in the signature.
-suggestAddArgument :: ParsedModule -> Diagnostic -> Either ResponseError [(T.Text, [TextEdit])]
-suggestAddArgument parsedModule Diagnostic {_message, _range}
-  | Just (name, typ) <- matchVariableNotInScope message = addArgumentAction parsedModule _range name typ
-  | Just (name, typ) <- matchFoundHoleIncludeUnderscore message = addArgumentAction parsedModule _range name (Just typ)
-  | otherwise = pure []
-  where
-    message = unifySpaces _message
-
--- TODO use typ to modify type signature
-addArgumentAction :: ParsedModule -> Range -> T.Text -> Maybe T.Text -> Either ResponseError [(T.Text, [TextEdit])]
-addArgumentAction (ParsedModule _ parsedSource _ _) range name _typ =
+addAsLastArgument :: UndefinedVariableHandler
+addAsLastArgument (ParsedModule _ parsedSource _ _) range name _typ =
   do
     let addArgToMatch (L locMatch (Match xMatch ctxMatch pats rhs)) = do
           let unqualName = mkRdrUnqual $ mkVarOcc $ T.unpack name
@@ -986,17 +1002,55 @@ addArgumentAction (ParsedModule _ parsedSource _ _) range name _typ =
             let decl' = L locDecl (ValD xVal (FunBind xFunBind idFunBind mg' coreFunBind))
             pure [decl']
           decl -> pure [decl]
-    case runTransformT $ modifySmallestDeclWithM spanContainsRangeOrErr insertArg (makeDeltaAst parsedSource) of
+    case runTransformT $ modifySmallestDeclWithM (flip spanContainsRangeOrErr range) 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
+
+-- TODO use typ to initialise type signature
+addToWhere :: ParsedModule -> Range -> T.Text -> Maybe T.Text -> Either ResponseError [(T.Text, [TextEdit])]
+addToWhere (ParsedModule _ parsedSource _ _) range name _typ = bimap traceShowId id $ do
+  let mkUnqual name = noLocA $ mkRdrUnqual $ mkVarOcc $ T.unpack name
+      equalAnn dp = AddEpAnn AnnEqual (EpaDelta dp [])
+      addArg = modifySmallestDeclWithM (flip spanContainsRangeOrErr range) $ \case
+        (L locDecl (ValD xVal (FunBind xFunBind idFunBind mg coreFunBind))) -> do
+          declMatchSrcSpan <- uniqueSrcSpanT
+          grhsSrcSpan <- uniqueSrcSpanT
+          mg' <- modifyMgMatchesT mg $ \match -> do
+            spanInRange <- lift $ getLoc match `spanContainsRangeOrErr` range
+            if spanInRange
+              then do
+                let grhs_ann = GrhsAnn Nothing $ equalAnn (SameLine 0)
+                    rhs_hole = L (noAnnSrcSpanDP1 generatedSrcSpan) $ HsVar NoExtField (mkUnqual "_")
+                    grhs = GRHS (EpAnn (generatedAnchor m1) grhs_ann emptyComments) [] rhs_hole
+                    grhss = GRHSs emptyComments [L grhsSrcSpan grhs] (EmptyLocalBinds NoExtField)
+                    newDeclMatchAnn = emptyEpAnnAnchor (generatedAnchor m0)
+                    newDeclMatch =
+                      noLocA (Match newDeclMatchAnn (FunRhs (mkUnqual name) Prefix SrcStrict) [] grhss)
+                    newDeclMg = MG NoExtField (L (noAnnSrcSpanDP0 declMatchSrcSpan) [newDeclMatch]) Generated
+                    newDecl = (FunBind NoExtField (mkUnqual name) newDeclMg [])
+                prependDeclToWhereDecls match (noLocA newDecl)
+              else pure match
+          let decl' = L locDecl (ValD xVal (FunBind xFunBind idFunBind mg' coreFunBind))
+          pure [decl']
+        _ -> pure []
+  (newSource, _, _) <- runTransformT $ addArg (makeDeltaAst parsedSource)
+  let diffText = makeDiffTextEdit (T.pack $ exactPrint parsedSource) (T.pack $ exactPrint newSource)
+  pure [("Add to where ‘" <> name <> "’", fromLspList diffText)]
+
+spanContainsRangeOrErr :: SrcSpan -> Range -> Either ResponseError Bool
+spanContainsRangeOrErr srcSpan range = maybeToEither (responseError "SrcSpan was not valid range") . (`spanContainsRange` range) $ srcSpan
+
+generatedAnchor :: AnchorOperation -> Anchor
+generatedAnchor anchorOp = GHC.Anchor (GHC.realSrcSpan generatedSrcSpan) anchorOp
+
+emptyEpAnnAnchor :: Monoid a => Anchor -> EpAnn a
+emptyEpAnnAnchor anchor = EpAnn anchor mempty emptyComments
 
 fromLspList :: List a -> [a]
 fromLspList (List a) = a
+#endif
 
 suggestFillTypeWildcard :: Diagnostic -> [(T.Text, TextEdit)]
 suggestFillTypeWildcard Diagnostic{_range=_range,..}