Be much more intelligent about splitting matches

ghcide/src/Development/IDE/GHC/ExactPrint.hs

@@ -1,19 +1,19 @@
-{-# LANGUAGE CPP                #-}
-{-# LANGUAGE DerivingStrategies #-}
-{-# LANGUAGE GADTs              #-}
-{-# LANGUAGE OverloadedStrings  #-}
-{-# LANGUAGE RankNTypes         #-}
-{-# LANGUAGE TypeFamilies       #-}
+{-# LANGUAGE CPP          #-}
+{-# LANGUAGE DerivingVia  #-}
+{-# LANGUAGE GADTs        #-}
+{-# LANGUAGE RankNTypes   #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module Development.IDE.GHC.ExactPrint
     ( Graft(..),
       graftDecls,
       graftDeclsWithM,
       annotate,
+      annotateDecl,
       hoistGraft,
       graftWithM,
-      graftWithSmallestM,
-      graftSmallestDecls,
+      genericGraftWithSmallestM,
+      genericGraftWithLargestM,
       graftSmallestDeclsWithM,
       transform,
       transformM,
@@ -27,6 +27,9 @@ module Development.IDE.GHC.ExactPrint
       TransformT,
       Anns,
       Annotate,
+      mkBindListT,
+      setPrecedingLinesT,
+      everywhereM',
     )
 where
 
@@ -63,7 +66,8 @@ import Parser (parseIdentifier)
 import Data.Traversable (for)
 import Data.Foldable (Foldable(fold))
 import Data.Bool (bool)
-import Data.Monoid (All(All))
+import Data.Monoid (All(All), Any(Any))
+import Data.Functor.Compose (Compose(Compose))
 #if __GLASGOW_HASKELL__ == 808
 import Control.Arrow
 #endif
@@ -195,7 +199,7 @@ needsParensSpace ExplicitSum{}   = (All False, All False)
 needsParensSpace HsCase{}        = (All False, All False)
 needsParensSpace HsIf{}          = (All False, All False)
 needsParensSpace HsMultiIf{}     = (All False, All False)
-needsParensSpace HsLet{}         = (All False, All False)
+needsParensSpace HsLet{}         = (All False, All True)
 needsParensSpace HsDo{}          = (All False, All False)
 needsParensSpace ExplicitList{}  = (All False, All False)
 needsParensSpace RecordCon{}     = (All False, All False)
@@ -286,30 +290,56 @@ graftWithM dst trans = Graft $ \dflags a -> do
         )
         a
 
-graftWithSmallestM ::
-    forall ast m a.
-    (Fail.MonadFail m, Data a, ASTElement ast) =>
+-- | A generic query intended to be used for calling 'smallestM' and
+-- 'largestM'. If the current node is a 'Located', returns whether or not the
+-- given 'SrcSpan' is a subspan. For all other nodes, returns 'Nothing', which
+-- indicates uncertainty. The search strategy in 'smallestM' et al. will
+-- continue searching uncertain nodes.
+genericIsSubspan ::
+    forall ast.
+    Typeable ast =>
+    -- | The type of nodes we'd like to consider.
+    Proxy (Located ast) ->
     SrcSpan ->
-    (Located ast -> TransformT m (Maybe (Located ast))) ->
+    GenericQ (Maybe Bool)
+genericIsSubspan _ dst = mkQ Nothing $ \case
+  (L span _ :: Located ast) -> Just $ dst `isSubspanOf` span
+
+-- | Run the given transformation only on the smallest node in the tree that
+-- contains the 'SrcSpan'.
+genericGraftWithSmallestM ::
+    forall m a ast.
+    (Monad m, Data a, Typeable ast) =>
+    -- | The type of nodes we'd like to consider when finding the smallest.
+    Proxy (Located ast) ->
+    SrcSpan ->
+    (DynFlags -> GenericM (TransformT m)) ->
     Graft m a
-graftWithSmallestM dst trans = Graft $ \dflags a -> do
-    everywhereM'
-        ( mkM $
-            \case
-                val@(L src _ :: Located ast)
-                    | dst `isSubspanOf` src -> do
-                        mval <- trans val
-                        case mval of
-                            Just val' -> do
-                                (anns, val'') <-
-                                    hoistTransform (either Fail.fail pure) $
-                                        annotate dflags True $ maybeParensAST val'
-                                modifyAnnsT $ mappend anns
-                                pure val''
-                            Nothing -> pure val
-                l -> pure l
-        )
-        a
+genericGraftWithSmallestM proxy dst trans = Graft $ \dflags ->
+    smallestM (genericIsSubspan proxy dst) (trans dflags)
+
+-- | Run the given transformation only on the largest node in the tree that
+-- contains the 'SrcSpan'.
+genericGraftWithLargestM ::
+    forall m a ast.
+    (Monad m, Data a, Typeable ast) =>
+    -- | The type of nodes we'd like to consider when finding the largest.
+    Proxy (Located ast) ->
+    SrcSpan ->
+    (DynFlags -> GenericM (TransformT m)) ->
+    Graft m a
+genericGraftWithLargestM proxy dst trans = Graft $ \dflags ->
+    largestM (genericIsSubspan proxy dst) (trans dflags)
+
+
+-- | Lift a function that replaces a value with several values into a generic
+-- function. The result doesn't perform any searching, so should be driven via
+-- 'everywhereM' or friends.
+--
+-- The 'Int' argument is the index in the list being bound.
+mkBindListT :: forall b m. (Typeable b, Data b, Monad m) => (Int -> b -> m [b]) -> GenericM m
+mkBindListT f = mkM $ fmap join . traverse (uncurry f) . zip [0..]
+
 
 graftDecls ::
     forall a.
@@ -319,32 +349,13 @@ graftDecls ::
     Graft (Either String) a
 graftDecls dst decs0 = Graft $ \dflags a -> do
     decs <- forM decs0 $ \decl -> do
-        (anns, decl') <- annotateDecl dflags decl
-        modifyAnnsT $ mappend anns
-        pure decl'
+        annotateDecl dflags decl
     let go [] = DL.empty
         go (L src e : rest)
             | src == dst = DL.fromList decs <> DL.fromList rest
             | otherwise = DL.singleton (L src e) <> go rest
     modifyDeclsT (pure . DL.toList . go) a
 
-graftSmallestDecls ::
-    forall a.
-    (HasDecls a) =>
-    SrcSpan ->
-    [LHsDecl GhcPs] ->
-    Graft (Either String) a
-graftSmallestDecls dst decs0 = Graft $ \dflags a -> do
-    decs <- forM decs0 $ \decl -> do
-        (anns, decl') <- annotateDecl dflags decl
-        modifyAnnsT $ mappend anns
-        pure decl'
-    let go [] = DL.empty
-        go (L src e : rest)
-            | dst `isSubspanOf` src = DL.fromList decs <> DL.fromList rest
-            | otherwise = DL.singleton (L src e) <> go rest
-    modifyDeclsT (pure . DL.toList . go) a
-
 graftSmallestDeclsWithM ::
     forall a.
     (HasDecls a) =>
@@ -356,11 +367,8 @@ graftSmallestDeclsWithM dst toDecls = Graft $ \dflags a -> do
         go (e@(L src _) : rest)
             | dst `isSubspanOf` src = toDecls e >>= \case
                 Just decs0 -> do
-                    decs <- forM decs0 $ \decl -> do
-                        (anns, decl') <-
-                            annotateDecl dflags decl
-                        modifyAnnsT $ mappend anns
-                        pure decl'
+                    decs <- forM decs0 $ \decl ->
+                        annotateDecl dflags decl
                     pure $ DL.fromList decs <> DL.fromList rest
                 Nothing -> (DL.singleton e <>) <$> go rest
             | otherwise = (DL.singleton e <>) <$> go rest
@@ -377,12 +385,9 @@ graftDeclsWithM dst toDecls = Graft $ \dflags a -> do
         go (e@(L src _) : rest)
             | src == dst = toDecls e >>= \case
                 Just decs0 -> do
-                    decs <- forM decs0 $ \decl -> do
-                        (anns, decl') <-
-                            hoistTransform (either Fail.fail pure) $
-                            annotateDecl dflags decl
-                        modifyAnnsT $ mappend anns
-                        pure decl'
+                    decs <- forM decs0 $ \decl ->
+                        hoistTransform (either Fail.fail pure) $
+                          annotateDecl dflags decl
                     pure $ DL.fromList decs <> DL.fromList rest
                 Nothing -> (DL.singleton e <>) <$> go rest
             | otherwise = (DL.singleton e <>) <$> go rest
@@ -461,7 +466,7 @@ annotate dflags needs_space ast = do
     pure (anns', expr')
 
 -- | Given an 'LHsDecl', compute its exactprint annotations.
-annotateDecl :: DynFlags -> LHsDecl GhcPs -> TransformT (Either String) (Anns, LHsDecl GhcPs)
+annotateDecl :: DynFlags -> LHsDecl GhcPs -> TransformT (Either String) (LHsDecl GhcPs)
 -- The 'parseDecl' function fails to parse 'FunBind' 'ValD's which contain
 -- multiple matches. To work around this, we split the single
 -- 'FunBind'-of-multiple-'Match'es into multiple 'FunBind's-of-one-'Match',
@@ -474,24 +479,23 @@ annotateDecl dflags
     let set_matches matches =
           ValD ext fb { fun_matches = mg { mg_alts = L alt_src matches }}
 
-    (anns', alts') <- fmap unzip $ for (zip [0..] alts) $ \(ix :: Int, alt) -> do
+    (anns', alts') <- fmap unzip $ for alts $ \alt -> do
       uniq <- show <$> uniqueSrcSpanT
       let rendered = render dflags $ set_matches [alt]
       lift (mapLeft show $ parseDecl dflags uniq rendered) >>= \case
         (ann, L _ (ValD _ FunBind { fun_matches = MG { mg_alts = L _ [alt']}}))
-           -> pure (bool id (setPrecedingLines alt' 1 0) (ix /= 0) ann, alt')
+           -> pure (setPrecedingLines alt' 1 0 ann, alt')
         _ ->  lift $ Left "annotateDecl: didn't parse a single FunBind match"
 
-    let expr' = L src $ set_matches alts'
-        anns'' = setPrecedingLines expr' 1 0 $ fold anns'
-
-    pure (anns'', expr')
+    modifyAnnsT $ mappend $ fold anns'
+    pure $ L src $ set_matches alts'
 annotateDecl dflags ast = do
     uniq <- show <$> uniqueSrcSpanT
     let rendered = render dflags ast
     (anns, expr') <- lift $ mapLeft show $ parseDecl dflags uniq rendered
     let anns' = setPrecedingLines expr' 1 0 anns
-    pure (anns', expr')
+    modifyAnnsT $ mappend anns'
+    pure expr'
 
 ------------------------------------------------------------------------------
 
@@ -504,3 +508,77 @@ render dflags = showSDoc dflags . ppr
 -- | Put parentheses around an expression if required.
 parenthesize :: LHsExpr GhcPs -> LHsExpr GhcPs
 parenthesize = parenthesizeHsExpr appPrec
+
+
+------------------------------------------------------------------------------
+-- Custom SYB machinery
+------------------------------------------------------------------------------
+
+-- | Generic monadic transformations that return side-channel data.
+type GenericMQ r m = forall a. Data a => a -> m (r, a)
+
+------------------------------------------------------------------------------
+-- | Apply the given 'GenericM' at all every node whose children fail the
+-- 'GenericQ', but which passes the query itself.
+--
+-- The query must be a monotonic function when it returns 'Just'. That is, if
+-- @s@ is a subtree of @t@, @q t@ should return @Just True@ if @q s@ does. It
+-- is the True-to-false edge of the query that triggers the transformation.
+--
+-- Why is the query a @Maybe Bool@? The GHC AST intersperses 'Located' nodes
+-- with data nodes, so for any given node we can only definitely return an
+-- answer if it's a 'Located'. See 'genericIsSubspan' for how this parameter is
+-- used.
+smallestM :: forall m. Monad m => GenericQ (Maybe Bool) -> GenericM m -> GenericM m
+smallestM q f = fmap snd . go
+  where
+    go :: GenericMQ Any m
+    go x = do
+      case q x of
+        Nothing -> gmapMQ go x
+        Just True -> do
+          it@(r, x') <- gmapMQ go x
+          case r of
+            Any True -> pure it
+            Any False -> fmap (Any True,) $ f x'
+        Just False -> pure (mempty, x)
+
+------------------------------------------------------------------------------
+-- | Apply the given 'GenericM' at every node that passes the 'GenericQ', but
+-- don't descend into children if the query matches. Because this traversal is
+-- root-first, this policy will find the largest subtrees for which the query
+-- holds true.
+--
+-- Why is the query a @Maybe Bool@? The GHC AST intersperses 'Located' nodes
+-- with data nodes, so for any given node we can only definitely return an
+-- answer if it's a 'Located'. See 'genericIsSubspan' for how this parameter is
+-- used.
+largestM :: forall m. Monad m => GenericQ (Maybe Bool) -> GenericM m -> GenericM m
+largestM q f = go
+  where
+    go :: GenericM m
+    go x = do
+      case q x of
+        Just True -> f x
+        Just False -> pure x
+        Nothing -> gmapM go x
+
+newtype MonadicQuery r m a = MonadicQuery
+  { runMonadicQuery :: m (r, a)
+  }
+  deriving stock (Functor)
+  deriving Applicative via Compose m ((,) r)
+
+
+------------------------------------------------------------------------------
+-- | Like 'gmapM', but also returns side-channel data.
+gmapMQ ::
+    forall f r a. (Monoid r, Data a, Applicative f) =>
+    (forall d. Data d => d -> f (r, d)) ->
+    a ->
+    f (r, a)
+gmapMQ f = runMonadicQuery . gfoldl k pure
+  where
+    k :: Data d => MonadicQuery r f (d -> b) -> d -> MonadicQuery r f b
+    k c x = c <*> MonadicQuery (f x)
+

plugins/hls-tactics-plugin/src/Wingman/GHC.hs

@@ -176,17 +176,22 @@ allOccNames = everything (<>) $ mkQ mempty $ \case
     a -> S.singleton a
 
 
+------------------------------------------------------------------------------
+-- | Unpack the relevant parts of a 'Match'
+pattern AMatch :: HsMatchContext (NameOrRdrName (IdP GhcPs)) -> [Pat GhcPs] -> HsExpr GhcPs -> Match GhcPs (LHsExpr GhcPs)
+pattern AMatch ctx pats body <-
+  Match { m_ctxt = ctx
+        , m_pats = fmap fromPatCompatPs -> pats
+        , m_grhss = UnguardedRHSs body
+        }
 
 
 ------------------------------------------------------------------------------
 -- | A pattern over the otherwise (extremely) messy AST for lambdas.
 pattern Lambda :: [Pat GhcPs] -> HsExpr GhcPs -> HsExpr GhcPs
 pattern Lambda pats body <-
   HsLam _
-    (MG {mg_alts = L _ [L _
-      (Match { m_pats = fmap fromPatCompatPs -> pats
-             , m_grhss = UnguardedRHSs body
-             })]})
+    (MG {mg_alts = L _ [L _ (AMatch _ pats body) ]})
   where
     -- If there are no patterns to bind, just stick in the body
     Lambda [] body   = body