WIP

Retrie/ExactPrint.hs

@@ -29,6 +29,8 @@ module Retrie.ExactPrint
   , transferEntryAnnsT
   , transferEntryDPT
   , transferAnchor
+  , deltaComments
+  , stripCommentsA
     -- * Utils
   , debugDump
   , debugParse
@@ -89,14 +91,28 @@ debug c s = trace s c
 fix :: (Data ast, MonadIO m) => FixityEnv -> ast -> TransformT m ast
 fix env = fixAssociativity >=> fixEntryDP
   where
-    fixAssociativity = everywhereM (mkM (fixOneExpr env) `extM` fixOnePat env)
+    fixAssociativity = everywhereM (mkM (fixOneExpr' env) `extM` fixOnePat env)
     fixEntryDP = everywhereM (mkM fixOneEntryExpr `extM` fixOneEntryPat)
 
 -- Should (x op1 y) op2 z be reassociated as x op1 (y op2 z)?
 associatesRight :: Fixity -> Fixity -> Bool
 associatesRight (Fixity _ p1 a1) (Fixity _ p2 _a2) =
   p2 > p1 || p1 == p2 && a1 == InfixR
 
+
+fixOneExpr'
+  :: MonadIO m
+  => FixityEnv
+  -> LHsExpr GhcPs
+  -> TransformT m (LHsExpr GhcPs)
+fixOneExpr' env ll@(L _l2 (OpApp _x2 (L _l1 (OpApp _x1 _x op1 _y)) op2 _z)) = do
+    ll' <- fixOneExpr env ll
+    -- if (associatesRight (lookupOp op1 env) (lookupOp op2 env))
+    --    then lift $ liftIO $ debugPrint Loud "fixOneExpr':(ll,ll')"  [showAst (ll,ll')]
+    --    else return ()
+    return ll'
+fixOneExpr' env ll = fixOneExpr env ll
+
 -- We know GHC produces left-associated chains, so 'z' is never an
 -- operator application. We also know that this will be applied bottom-up
 -- by 'everywhere', so we can assume the children are already fixed.
@@ -105,38 +121,46 @@ fixOneExpr
   => FixityEnv
   -> LHsExpr GhcPs
   -> TransformT m (LHsExpr GhcPs)
-fixOneExpr env (L l2 (OpApp x2 ap1@(L _l1 (OpApp x1 x op1 y)) op2 z))
+fixOneExpr env (L l2 (OpApp x2 ap1@(L l1 (OpApp x1 x op1 y)) op2 z))
 {-
   pre
   x   is [print]   4:8-12
   op1 is [$]       4:14
   y   is [foo]     4:16-18
   op2 is [`bar`]   4:20-24
   z   is [[1..10]] 4:26-32
+  l1 carries comments that immediately precede or follow op1
+  l2 carries comments that immediately precede or follow op2
 
   (L l2 (OpApp _
           (L l1 (OpApp _ x op1 y))
           op2
           z))
   -- post
+  Comments mut move from l2 to new_loc
   (L l2 (OpApp _
           x
           op1
           (L new_loc (OpApp _ y op2 z)))
 -}
 
   | associatesRight (lookupOp op1 env) (lookupOp op2 env) = do
-    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:(l1,l2)="  [showAst (_l1,l2)]
+    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:(l1,l2)="  [showAst (l1,l2)]
     -- We need a location from start of y to end of z
     -- let ap2' = L (stripComments l2) $ OpApp x2 y op2 z
-    let ap2' :: LHsExpr GhcPs = L (noAnnSrcSpan (combineSrcSpans (locA y) (locA z)) ) $ OpApp x2 y op2 z
+    let ap2' :: LHsExpr GhcPs = L (transferComments l2 op2 $ noAnnSrcSpan (combineSrcSpans (locA y) (locA z)))
+                                  $ OpApp x2 y op2 z
+    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:ap2'"  [showAst ap2']
     -- (_ap1_0, ap2'_0) <- swapEntryDPT ap1 ap2'
     (_ap1_0, ap2'_0) <- return (ap1, ap2')
     -- Even though we process bottom-up, we need to recurse because we
     -- have changed the structure at this level
+    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:recursing rhs"  [showAst (getLoc ap2'_0)]
     rhs <- fixOneExpr env ap2'_0
-    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:returning"  [showAst (L l2 $ OpApp x1 x op1 rhs)]
-    return $ L l2 $ OpApp x1 x op1 rhs
+    let l2' = transferComments l1 op1 (stripComments l2)
+    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:l2'"  [showAst l2']
+    -- lift $ liftIO $ debugPrint Loud "fixOneExpr:returning"  [showAst (L l2' $ OpApp x1 x op1 rhs)]
+    return $ L l2' $ OpApp x1 x op1 rhs
 fixOneExpr _ e = return e
 
 fixOnePat :: Monad m => FixityEnv -> LPat GhcPs -> TransformT m (LPat GhcPs)
@@ -162,6 +186,15 @@ stripComments (SrcSpanAnn EpAnnNotUsed l) = SrcSpanAnn EpAnnNotUsed l
 stripComments (SrcSpanAnn (EpAnn anc an _) l) = SrcSpanAnn (EpAnn anc an emptyComments) l
 #endif
 
+#if __GLASGOW_HASKELL__ >= 910
+transferComments :: EpAnn an -> LocatedAn aa b -> EpAnn an -> EpAnn an
+transferComments (EpAnn _anc _an cs0) (L (EpAnn anc0 _ _) _) (EpAnn anc an cs1) = EpAnn anc an (addCommentOrigDeltas cs')
+    where
+      cs' = case anc0 of
+          EpaSpan (RealSrcSpan r _) -> splitCommentsStart r (cs0 <> cs1)
+          _ -> cs0 <> cs1
+#endif
+
 -- Move leading whitespace from the left child of an operator application
 -- to the application itself. We need this so we have correct offsets when
 -- substituting into patterns and don't end up with extra leading spaces.
@@ -381,6 +414,13 @@ transferAnchor (L (SrcSpanAnn EpAnnNotUsed l)    _) lb = setAnchorAn lb (spanAsA
 transferAnchor (L (SrcSpanAnn (EpAnn anc _ _) _) _) lb = setAnchorAn lb anc              emptyComments
 #endif
 
+#if __GLASGOW_HASKELL__ >= 910
+deltaComments :: LocatedAn a b -> LocatedAn a b
+deltaComments (L (EpAnn anc an cs) a) = L (EpAnn anc an (addCommentOrigDeltas cs)) a
+
+stripCommentsA :: LocatedAn a b -> LocatedAn a b
+stripCommentsA (L (EpAnn anc an _cs) a) = L (EpAnn anc an emptyComments) a
+#endif
 
 isComma :: TrailingAnn -> Bool
 isComma (AddCommaAnn _) = True

Retrie/ExactPrint/Compat.hs

@@ -12,7 +12,10 @@ module Retrie.ExactPrint.Compat
     , E.d0
     , E.uniqueSrcSpanT
 #if __GLASGOW_HASKELL__ >= 910
+    , E.addCommentOrigDeltas
     , transferEntryDP
+    , E.splitCommentsStart
+    , E.splitCommentsEnd
 #else
     , E.transferEntryDP
 #endif

Retrie/Expr.hs

@@ -413,9 +413,9 @@ parenify Context{..} le@(L ll e)
   | needed ctxtParentPrec (precedence ctxtFixityEnv e) && needsParens e = do
     let tokLP = EpTok (EpaDelta (SameLine 0) [])
         tokRP = EpTok (EpaDelta (SameLine 0) [])
-     -- in mkParen' (getEntryDP le) (\_an -> HsPar (tokLP, tokRP) (setEntryDP le (SameLine 0)))
     let le' = setEntryDP le (SameLine 0) :: LHsExpr GhcPs
     let r = L ll (HsPar (tokLP, tokRP) le') :: LHsExpr GhcPs
+    -- lift $ liftIO $ debugPrint Loud "parenify:r="  [showAst  (getLoc r)]
     return r
 #endif
   | otherwise = return le

Retrie/Replace.hs

@@ -90,12 +90,14 @@ replaceImpl c e = do
       -- substitute for quantifiers in grafted template
       r <- subst sub c t'
       -- copy appropriate annotations from old expression to template
-      r0 <- addAllAnnsT e r
+      -- r0 <- addAllAnnsT e r
+      r0 <- pure $ transferAnchor e r
       -- add parens to template if needed
       res' <- (mkM (parenify c) `extM` parenifyT c `extM` parenifyP c) r0
       -- Make sure the replacement has the same anchor as the thing
       -- being replaced
       res <- transferEntryDP e res'
+      let res2 = stripCommentsA res
 
       -- prune the resulting expression and log it with location
       orig <- printNoLeadingSpaces <$> pruneA e
@@ -105,17 +107,21 @@ replaceImpl c e = do
       lift $ liftIO $ debugPrint Loud "replaceImpl:orig="  [orig]
       lift $ liftIO $ debugPrint Loud "replaceImpl:repl="  [repl]
 
-      -- lift $ liftIO $ debugPrint Loud "replaceImpl:e="  [showAst e]
-      -- lift $ liftIO $ debugPrint Loud "replaceImpl:r="  [showAst r]
+      lift $ liftIO $ debugPrint Loud "replaceImpl:e="  [showAst e]
+      lift $ liftIO $ debugPrint Loud "replaceImpl:r="  [showAst r]
       -- lift $ liftIO $ debugPrint Loud "replaceImpl:r0="  [showAst r0]
       -- lift $ liftIO $ debugPrint Loud "replaceImpl:t'=" [showAst t']
       -- lift $ liftIO $ debugPrint Loud "replaceImpl:res=" [showAst res]
+      lift $ liftIO $ debugPrint Loud "replaceImpl:res2=" [showAst res2]
       -- lift $ liftIO $ debugPrint Loud "replaceImpl:res=" [showAst res']
+      -- lift $ liftIO $ debugPrint Loud "replaceImpl:(e)=" [showAst (getLoc e)]
+      -- lift $ liftIO $ debugPrint Loud "replaceImpl:(r0,r)=" [showAst (getLoc r0, getLoc r)]
+      -- lift $ liftIO $ debugPrint Loud "replaceImpl:(res,res')=" [showAst (getLoc res, getLoc res')]
 
       let replacement = Replacement (getLocA e) orig repl
       TransformT $ lift $ tell $ Change [replacement] [tImports]
       -- make the actual replacement
-      return res
+      return res2
 
 
 -- | Records a replacement made. In cases where we cannot use ghc-exactprint

Retrie/Subst.hs

@@ -55,7 +55,7 @@ substExpr ctxt e@(L l1 (HsVar x (L l2 v))) =
     Just (HoleExpr eA') -> do
       let eA = fmap makeDeltaAst eA'
       lift $ liftIO $ debugPrint Loud "substExpr:HoleExpr:e" [showAst e]
-      lift $ liftIO $ debugPrint Loud "substExpr:HoleExpr:eA" [showAst eA]
+      lift $ liftIO $ debugPrint Loud "substExpr:HoleExpr:eA" [showAst eA']
       e0 <- graftA (unparen <$> eA)
       let hasCs = hasComments e0
       e1 <- if hasCs

retrie.cabal

@@ -47,6 +47,7 @@ library
     Retrie.Elaborate,
     Retrie.ExactPrint,
     Retrie.ExactPrint.Annotated,
+    Retrie.ExactPrint.Compat
     Retrie.Expr,
     Retrie.Fixity,
     Retrie.FreeVars,

tests/inputs/Operator1.test

@@ -0,0 +1,17 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+-l Operator1.print
+===
+ module Operator1 where
+
+ main :: IO ()
+ main = do
+-  putStrLn $ show $ {- c1 -} foo || bar
++  print $  {- c1 -} foo || bar
+
+
+
+ {-# RULES "print" forall x. putStrLn $ show $ x = print $ x #-}

tests/inputs/Recursion.test

@@ -11,10 +11,10 @@
 ===
  module Recursion where
 
- -- Goal of these tests is to make sure rewriting doesn't introduce unintended 
+ -- Goal of these tests is to make sure rewriting doesn't introduce unintended
  -- recursion. It doesn't protect against infinite mutual recursion, but we can
  -- at least spot self-recursion.
- 
+
  -- foo should not be rewritten
  foo :: Int -> Int
  foo = foldr bar baz