Support removing brackets from templates/substs

ndmitchell · Mar 19, 2021 · 990a4f3 · 990a4f3
1 parent ca21afb
commit 990a4f3
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 32 deletions.
diff --git a/src/GHC/Util/HsExpr.hs b/src/GHC/Util/HsExpr.hs
@@ -27,7 +27,9 @@ import GHC.Util.FreeVars
 import GHC.Util.View
 
 import Control.Applicative
+import Control.Monad.Trans.Class
 import Control.Monad.Trans.State
+import Control.Monad.Trans.Writer.CPS
 
 import Data.Data
 import Data.Generics.Uniplate.DataOnly
@@ -91,10 +93,13 @@ transformAppsM :: Monad m => (LHsExpr GhcPs -> m (LHsExpr GhcPs)) -> LHsExpr Ghc
 transformAppsM f x = f =<< descendAppsM (transformAppsM f) x
 
 descendIndex :: Data a => (Int -> a -> a) -> a -> a
-descendIndex f x = flip evalState 0 $ flip descendM x $ \y -> do
+descendIndex f = fst . descendIndex' (\x a -> writer (f x a, ()))
+
+descendIndex' :: (Data a, Monoid w) => (Int -> a -> Writer w a) -> a -> (a, w)
+descendIndex' f x = runWriter $ flip evalStateT 0 $ flip descendM x $ \y -> do
     i <- get
     modify (+1)
-    pure $ f i y
+    lift $ f i y
 
 --  There are differences in pretty-printing between GHC and HSE. This
 --  version never removes brackets.
@@ -114,7 +119,6 @@ appsBracket :: [LHsExpr GhcPs] -> LHsExpr GhcPs
 appsBracket = foldl1 mkApp
   where mkApp x y = rebracket1 (noLoc $ HsApp noExtField x y)
 
-
 simplifyExp :: LHsExpr GhcPs -> LHsExpr GhcPs
 -- Replace appliciations 'f $ x' with 'f (x)'.
 simplifyExp (L l (OpApp _ x op y)) | isDol op = L l (HsApp noExtField x (noLoc (HsPar noExtField y)))
@@ -136,7 +140,6 @@ niceDotApp :: LHsExpr GhcPs -> LHsExpr GhcPs -> LHsExpr GhcPs
 niceDotApp (L _ (HsVar _ (L _ r))) b | occNameStr r == "$" = b
 niceDotApp a b = dotApp a b
 
-
 -- Generate a lambda expression but prettier if possible.
 niceLambda :: [String] -> LHsExpr GhcPs -> LHsExpr GhcPs
 niceLambda ss e = fst (niceLambdaR ss e)-- We don't support refactorings yet.
@@ -271,37 +274,45 @@ needBracketOld i parent child
   | isDotApp parent, isDotApp child, i == 2 = False
   | otherwise = needBracket i parent child
 
-transformBracketOld :: (LHsExpr GhcPs -> Maybe (LHsExpr GhcPs)) -> LHsExpr GhcPs -> (LHsExpr GhcPs, LHsExpr GhcPs)
+transformBracketOld :: (LHsExpr GhcPs -> Maybe (LHsExpr GhcPs))
+                    -> LHsExpr GhcPs
+                    -> (LHsExpr GhcPs, (LHsExpr GhcPs, [String]))
 transformBracketOld op = first snd . g
   where
     g = first f . descendBracketOld g
     f x = maybe (False, x) (True, ) (op x)
 
 -- Descend, and if something changes then add/remove brackets
--- appropriately. Returns (suggested replacement, refactor template).
--- Whenever a bracket is added to the suggested replacement, a
--- corresponding bracket is added to the refactor template.
-descendBracketOld :: (LHsExpr GhcPs -> ((Bool, LHsExpr GhcPs), LHsExpr GhcPs))
-                   -> LHsExpr GhcPs
-                   -> (LHsExpr GhcPs, LHsExpr GhcPs)
-descendBracketOld op x = (descendIndex g1 x, descendIndex g2 x)
+-- appropriately. Returns (suggested replacement, (refactor template, no bracket vars)),
+-- where "no bracket vars" is a list of substitution variables which, when expanded,
+-- should have the brackets stripped.
+descendBracketOld :: (LHsExpr GhcPs -> ((Bool, LHsExpr GhcPs), (LHsExpr GhcPs, [String])))
+                  -> LHsExpr GhcPs
+                  -> (LHsExpr GhcPs, (LHsExpr GhcPs, [String]))
+descendBracketOld op x = (descendIndex g1 x, descendIndex' g2 x)
   where
-    g i y = if a then (f1 i b z, f2 i b z) else (b, z)
-      where ((a, b), z) = op y
-
-    g1 = (fst .) . g
-    g2 = (snd .) . g
-
-    f i (L _ (HsPar _ y)) z
-      | not $ needBracketOld i x y = (y, z)
-    f i y z
-      | needBracketOld i x y = (addParen y, addParen z)
+    g i y = if a then (f1 i b z w, f2 i b z w) else (b, (z, w))
+      where ((a, b), (z, w)) = op y
+
+    g1 a b = fst (g a b)
+    g2 a b = writer $ snd (g a b)
+
+    f i (L _ (HsPar _ y)) z w
+      | not $ needBracketOld i x y = (y, removeBracket z)
+      where
+        -- If the template expr is a Var, record it so that we can remove the brackets
+        -- later when expanding it. Otherwise, remove the enclosing brackets (if any).
+        removeBracket = \case
+          var@(L _ HsVar{}) -> (z, varToStr var : w)
+          other -> (fromParen z, w)
+    f i y z w
+      | needBracketOld i x y = (addParen y, (addParen z, w))
       -- https://github.com/mpickering/apply-refact/issues/7
-      | isOp y = (y, addParen z)
-    f _ y z                 = (y, z)
+      | isOp y = (y, (addParen z, w))
+    f _ y z w = (y, (z, w))
 
-    f1 = ((fst .) .) . f
-    f2 = ((snd .) .) . f
+    f1 a b c d = fst (f a b c d)
+    f2 a b c d = snd (f a b c d)
 
     isOp = \case
       L _ (HsVar _ (L _ name)) -> isSymbolRdrName name

diff --git a/src/GHC/Util/Unify.hs b/src/GHC/Util/Unify.hs
@@ -62,10 +62,14 @@ removeParens noParens (Subst xs) = Subst $
   map (\(x, y) -> if x `elem` noParens then (x, fromParen y) else (x, y)) xs
 
 -- Peform a substition.
--- Returns (suggested replacement, refactor template), both with brackets added
--- as needed.
--- Example: (traverse foo (bar baz), traverse f (x))
-substitute :: Subst (LHsExpr GhcPs) -> LHsExpr GhcPs -> (LHsExpr GhcPs, LHsExpr GhcPs)
+-- Returns (suggested replacement, (refactor template, no bracket vars)). It adds/removes brackets
+-- for both the suggested replacement and the refactor template appropriately. The "no bracket vars"
+-- is a list of substituation variables which, when expanded, should have the brackets stripped.
+--
+-- Examples:
+--   (traverse foo (bar baz), (traverse f (x), []))
+--   (zipWith foo bar baz, (f a b, [f]))
+substitute :: Subst (LHsExpr GhcPs) -> LHsExpr GhcPs -> (LHsExpr GhcPs, (LHsExpr GhcPs, [String]))
 substitute (Subst bind) = transformBracketOld exp . transformBi pat . transformBi typ
   where
     exp :: LHsExpr GhcPs -> Maybe (LHsExpr GhcPs)

diff --git a/src/Hint/Match.hs b/src/Hint/Match.hs
@@ -146,7 +146,7 @@ matchIdea sb declName HintRule{..} parent x = do
   -- (with 'e') need to unqualify before substitution (with 'res').
   let rhs' | Just fun <- extra = rebracket1 $ noLoc (HsApp noExtField fun rhs)
            | otherwise = rhs
-      (e, tpl) = substitute u rhs'
+      (e, (tpl, substNoParens)) = substitute u rhs'
       noParens = [varToStr $ fromParen x | L _ (HsApp _ (varToStr -> "_noParen_") x) <- universe tpl]
 
   u <- pure (removeParens noParens u)
@@ -167,7 +167,11 @@ matchIdea sb declName HintRule{..} parent x = do
   (u, tpl) <- pure $ if any ((== noSrcSpan) . getLoc . snd) (fromSubst u) then (mempty, res) else (u, tpl)
   tpl <- pure $ unqualify sa sb (performSpecial tpl)
 
-  pure (res, tpl, hintRuleNotes, [(s, toSS pos) | (s, pos) <- fromSubst u, getLoc pos /= noSrcSpan])
+  pure ( res, tpl, hintRuleNotes,
+         [ (s, toSS pos') | (s, pos) <- fromSubst u, getLoc pos /= noSrcSpan
+                          , let pos' = if s `elem` substNoParens then fromParen pos else pos
+         ]
+       )
 
 ---------------------------------------------------------------------
 -- SIDE CONDITIONS