Splicer and context reprinting

src/Text/Reprinter.hs

@@ -17,10 +17,13 @@ module Text.Reprinter
   , RefactorType(..)
   , Refactorable(..)
   , Reprinting
+  , ZipperReprinting
   , catchAll
   , genReprinting
   , reprint
   , reprintSort
+  , Line(..)
+  , Col(..)
   ) where
 
 -- Import solely for re-exporting for library clients
@@ -86,32 +89,37 @@ type Span = (Position, Position)
 -- @i@ is the input type (something with a '[Char]'-like interface)
 type Reprinting i m = forall node . (Typeable node) => node -> m (Maybe (RefactorType, i, Span))
 
+-- | Zipper reprinting allows reprinting functions to take into account context
+type ZipperReprinting i ast m = Zipper ast -> m (Maybe (RefactorType, i, Span))
+
 -- | Specify a refactoring type
 data RefactorType = Before | After | Replace
     deriving Show -- for debugging
 
 -- | The reprint algorithm takes a refactoring (parameteric in
 -- | some monad m) and turns an arbitrary pretty-printable type 'ast'
 -- | into a monadic 'StringLike i' transformer.
-reprint :: (Monad m, Data ast, StringLike i) => Reprinting i m -> ast -> i -> m i
-reprint reprinting ast input
+reprint :: (Monad m, Data ast, StringLike i)
+   => Reprinting i m -> (i -> m i) -> ast -> i -> m i
+reprint reprinting splicer ast input
   -- If the input is empty return empty
   | slNull input = return mempty
-
   -- Otherwise proceed with the algorithm
   | otherwise = do
     -- Initial state comprises start cursor and input source
     let state_0 = (initPosition, input)
     -- Enter the top-node of a zipper for `ast'
-    let comp = enter reprinting (toZipper ast)
-    (out, (_, remaining)) <- runStateT comp state_0
-    -- Add to the output source the remaining input source
-    return (out <> remaining)
+    let comp    = enter reprinting splicer (toZipper ast)
+    (output, (_, inp)) <- runStateT comp state_0
+    -- Splice the remaining input and append to output source
+    remaining <- splicer inp
+    return (output <> remaining)
 
 -- | Take a refactoring and a zipper producing a stateful 'StringLike i'
 -- | transformer with Position state.
-enter :: (Monad m, StringLike i) => Reprinting i m -> Zipper ast -> StateT (Position, i) m i
-enter reprinting zipper = do
+enter :: (Monad m, StringLike i) =>
+  Reprinting i m -> (i -> m i) -> Zipper ast -> StateT (Position, i) m i
+enter reprinting splicer zipper = do
     -- Step 1: Apply a refactoring
     refactoringInfo <- lift (query reprinting zipper)
 
@@ -120,7 +128,7 @@ enter reprinting zipper = do
       -- No refactoring; go to children
       Nothing -> go down'
       -- A refactoring was applied
-      Just r -> splice r
+      Just r -> splice splicer r
     -- Step 3: Enter the right sibling of the current context
     outputSib <- go right
 
@@ -132,53 +140,55 @@ enter reprinting zipper = do
     go direction =
         case direction zipper of
           -- Go to next node if there is one
-          Just zipper -> enter reprinting zipper
+          Just zipper -> enter reprinting splicer zipper
           -- Otherwise return the empty string
           Nothing -> return mempty
 
 
 -- | The reprint algorithm takes a refactoring (parameteric in
--- | some monad m) and turns an arbitrary pretty-printable type 'ast'
--- | into a monadic 'StringLike i' transformer.
-reprintSort :: (Monad m, Data ast, StringLike i) => Reprinting i m -> ast -> i -> m i
-reprintSort reprinting ast input
+-- some monad m), and a splicer (also parametric in m) turns an arbitrary
+-- pretty-printable type 'ast' into a monadic 'StringLike i' transformer.
+reprintSort :: (Monad m, Data ast, StringLike i) =>
+  ZipperReprinting i ast m -> (i -> m i) -> ast -> i -> m i
+reprintSort reprinting splicer ast input
   -- If the input is empty return empty
   | slNull input = return mempty
-
   -- Otherwise proceed with the algorithm
   | otherwise = do
     -- Initial state comprises start cursor and input source
     let state_0 = (initPosition, input)
     -- Enter the top-node of a zipper for `ast'
-    let comp = enter' reprinting (toZipper ast)
-    (out, (_, remaining)) <- runStateT comp state_0
-    -- Add to the output source the remaining input source
+    let comp    = enter' reprinting splicer (toZipper ast)
+    (out, (_, inp)) <- runStateT comp state_0
+    remaining       <- splicer inp
     return (out <> remaining)
 
 
 -- | Take a refactoring and a zipper to produce a list of refactorings
-enter' :: (Monad m, StringLike i) => Reprinting i m -> Zipper ast
-      -> StateT (Position, i) m i
-enter' reprinting zipper = do
+enter' :: (Monad m, StringLike i)
+  => ZipperReprinting i ast m -> (i -> m i) -> Zipper ast
+  -> StateT (Position, i) m i
+enter' reprinting splicer zipper = do
     -- Step 1: Get refactorings via AST zipper traversal
     rs <- lift $ getRefactorings reprinting zipper []
     -- Step 2: Do the splicing on the sorted refactorings
-    srcs <- mapM splice (sortBySpan . reverse $ rs)
+    srcs <- mapM (splice splicer) . sortBySpan . reverse $ rs
     return $ mconcat srcs
   where
     sortBySpan = sortOn (\(_,_,sp) -> sp)
 
-getRefactorings :: (Monad m, StringLike i) => Reprinting i m -> Zipper ast -> [(RefactorType, i, Span)]
-                    -> m [(RefactorType, i, Span)]
+getRefactorings :: (Monad m, StringLike i)
+  => ZipperReprinting i ast m -> Zipper ast -> [(RefactorType, i, Span)]
+  -> m [(RefactorType, i, Span)]
 getRefactorings reprinting zipper acc = do
     -- Step 1: Apply a refactoring
-    refactoringInfo <- query reprinting zipper
+    refactoringInfo <- reprinting zipper
     -- Step 2: Deal with refactored code or go to children
-    acc <- case refactoringInfo of
+    acc             <- case refactoringInfo of
       -- No refactoring; go to children
       Nothing -> go down' acc
       -- A refactoring was applied, add it to the accumulator
-      Just r -> return (r : acc)
+      Just r  -> return (r : acc)
     -- Step 3: Enter the left sibling of the current focus
     acc <- go right acc
     -- Finally return the accumulated refactorings
@@ -192,29 +202,37 @@ getRefactorings reprinting zipper acc = do
           -- Otherwise return the empty string
           Nothing -> return acc
 
-splice :: (Monad m, StringLike i) => (RefactorType, i, Span) -> StateT (Position, i) m i
-splice (typ, output, (lb, ub)) = do
+-- | This takes in a splicer function to allow for changes some as line breaks
+-- and indentation.
+splice :: (Monad m, StringLike i) =>
+  (i -> m i) -> (RefactorType, i, Span) -> StateT (Position, i) m i
+splice splicer (typ, output, (lb, ub)) = do
     (cursor, inp) <- get
     case typ of
       Replace -> do
-        -- Get soure up to start of refactored node
-        let (pre, inp') = splitBySpan (cursor, lb) inp
+        -- Get source up to start of refactored node
+        let (pre, inp' ) = splitBySpan (cursor, lb) inp
         -- Remove source covered by refactoring
-        let (_, inp'') = splitBySpan (lb, ub) inp'
+            (_  , inp'') = splitBySpan (lb, ub) inp'
         put (ub, inp'')
-        return (pre <> output)
+        spliceAndTrim [pre, output]
       After -> do
         -- Get source up to end of the refactored node
         let (pre, inp') = splitBySpan (cursor, ub) inp
         put (ub, inp')
-        return (pre <> output)
+        spliceAndTrim [pre, output]
       Before -> do
         -- Get source up to start of refactored node
-        let (pre, inp') = splitBySpan (cursor, lb) inp
+        let (pre, inp')   = splitBySpan (cursor, lb) inp
         -- Discard portion consumed by the refactoring
         let (post, inp'') = splitBySpan (lb, ub) inp'
         put (ub, inp'')
-        return (pre <> output <> post)
+        spliceAndTrim [pre, output, post]
+  where
+    spliceAndTrim = lift . fmap (foldr trimSpaces mempty) . traverse splicer
+    trimSpaces x y = if y /= mempty && (fst <$> slUncons y) == Just '\n'
+      then slDropWhileEnd (== ' ') x <> y
+      else x <> y
 
 -- | Given a lower-bound and upper-bound pair of Positions, split the
 -- | incoming 'StringLike i' based on the distance between the Position pairs.
@@ -236,16 +254,14 @@ splitBySpan (lower, upper) =
       | otherwise = done
       where done = (slReverse acc, input)
 
-
-
 -- | Infrastructure for building the reprinter "plugins"
 class Refactorable t where
   isRefactored :: t -> Maybe RefactorType
   getSpan      :: t -> Span
 
 -- | Essentially wraps the refactorable interface
-genReprinting :: (Monad m, Refactorable t, Typeable t, StringLike i)
-              => (t -> m i) -> t -> m (Maybe (RefactorType, i, Span))
+genReprinting :: (Monad m, Refactorable t, Typeable t, StringLike i) =>
+  (t -> m i) -> t -> m (Maybe (RefactorType, i, Span))
 genReprinting f z = case isRefactored z of
     Nothing -> return Nothing
     Just refactorType -> do
@@ -254,4 +270,4 @@ genReprinting f z = case isRefactored z of
 
 -- | Catch all generic query
 catchAll :: Monad m => a -> m (Maybe b)
-catchAll _ = return Nothing
+catchAll _ = return Nothing

src/Text/Reprinter/Example.lhs

@@ -223,7 +223,7 @@ let's code that reprinting and the required pretty printer:
 
 \begin{code}
 -- See the 2017 paper and SYB documentation for more info on 'extQ' queries.
-exprReprinter :: Reprinting String Identity
+exprReprinter :: Monad m => Reprinting String m
 exprReprinter = catchAll `extQ` reprintExpr
   where
     reprintExpr x = genReprinting (return . prettyExpr) (x :: Expr Bool)
@@ -252,7 +252,7 @@ reprints.
 refactor :: String -> String
 refactor s =
       runIdentity
-    . flip (reprint exprReprinter) s
+    . flip (reprint exprReprinter pure) s
     . refactorZero
     . parse $ s
 
@@ -273,8 +273,8 @@ commentPrinter = catchAll `extQ` decl
     decl (Decl _ s v e) = do
       val <- eval (e :: Expr Bool)
       case val of
-        Nothing -> return $ Nothing -- declaration expression referenced a
-                                    -- variable before assignment: no annotation
+        Nothing -> return Nothing -- declaration expression referenced a
+                                  -- variable before assignment: no annotation
         Just val -> do
           modify ((v,val) :)    -- add mapping to environment
           let msg = " // " <> v <> " = " <> show val
@@ -291,7 +291,7 @@ eval (Var _ _ s) = get >>= return . lookup s
 refactorComment :: String -> String
 refactorComment input =
       flip evalState []
-    . flip (reprint commentPrinter) input
+    . flip (reprint commentPrinter pure) input
     . parse $ input
 \end{code}
 
@@ -357,7 +357,7 @@ parseExpr = do
     else do
        isVar <- peekChar isAlpha
        if isVar then do
-           name <- many isAlpha
+           name <- many (\x -> isAlpha x || x == '_')
            p2 <- getPos
            return $ Var False (p1, p2) name
        else do

src/Text/Reprinter/StringLike.hs

@@ -5,7 +5,7 @@ module Text.Reprinter.StringLike
   , IsString(..)
   ) where
 
-import           Data.List   (uncons)
+import           Data.List   (uncons, dropWhileEnd)
 import           Data.String (IsString(..))
 
 import qualified Data.Text                  as TextStrict
@@ -20,13 +20,14 @@ import qualified Data.ByteString.Lazy.Char8 as BSCLazy
 -- type. Only operations required by the reprinting algorithm are included.
 -- Where possible, operations are prefilled using presumed-existing instances
 -- (any @[Char]@-like should be a monoid and have a @String -> a@).
-class (Monoid a, IsString a) => StringLike a where
+class (Eq a, Monoid a, IsString a) => StringLike a where
     slCons :: Char -> a -> a
     slUncons :: a -> Maybe (Char, a)
     slNull :: a -> Bool
     slReverse :: a -> a
     -- | like @unpack@
     slToString :: a -> String
+    slDropWhileEnd :: (Char -> Bool) -> a -> a
 
 -- same trick as used in IsString, to avoid possible ambiguity issues
 instance (a ~ Char) => StringLike [a] where
@@ -35,31 +36,36 @@ instance (a ~ Char) => StringLike [a] where
     slNull = null
     slReverse = reverse
     slToString = id
+    slDropWhileEnd = dropWhileEnd
 
 instance StringLike TextStrict.Text where
     slCons = TextStrict.cons
     slUncons = TextStrict.uncons
     slNull = TextStrict.null
     slReverse = TextStrict.reverse
     slToString = TextStrict.unpack
+    slDropWhileEnd = TextStrict.dropWhileEnd
 
 instance StringLike TextLazy.Text where
     slCons = TextLazy.cons
     slUncons = TextLazy.uncons
     slNull = TextLazy.null
     slReverse = TextLazy.reverse
     slToString = TextLazy.unpack
+    slDropWhileEnd = TextLazy.dropWhileEnd
 
 instance StringLike BSCStrict.ByteString where
     slCons = BSCStrict.cons
     slUncons = BSCStrict.uncons
     slNull = BSCStrict.null
     slReverse = BSCStrict.reverse
     slToString = BSCStrict.unpack
+    slDropWhileEnd p = snd . BSCStrict.spanEnd p
 
 instance StringLike BSCLazy.ByteString where
     slCons = BSCLazy.cons
     slUncons = BSCLazy.uncons
     slNull = BSCLazy.null
     slReverse = BSCLazy.reverse
     slToString = BSCLazy.unpack
+    slDropWhileEnd p = BSCLazy.fromStrict . snd . BSCStrict.spanEnd p . BSCLazy.toStrict