Add function to trim trailing spaces

prettyprinter/CHANGELOG.md

@@ -1,3 +1,8 @@
+# 1.2.2
+
+- Add function to trim trailing space in layouted `SimpleDocStream`,
+  `removeTrailingWhitespace`
+
 # 1.2.1
 
 - Add `Pretty` instances for `Identity` and `Const`

prettyprinter/prettyprinter.cabal

@@ -117,6 +117,7 @@ test-suite testsuite
     type: exitcode-stdio-1.0
     hs-source-dirs: test/Testsuite
     main-is: Main.hs
+    other-modules: StripTrailingSpace
     build-depends:
           base
         , prettyprinter

prettyprinter/src/Data/Text/Prettyprint/Doc.hs

@@ -278,6 +278,7 @@ module Data.Text.Prettyprint.Doc (
     SimpleDocStream(..),
     PageWidth(..), LayoutOptions(..), defaultLayoutOptions,
     layoutPretty, layoutCompact, layoutSmart,
+    removeTrailingWhitespace,
 
     -- * Migration guide
     --

prettyprinter/src/Data/Text/Prettyprint/Doc/Internal.hs

@@ -1428,6 +1428,64 @@ data SimpleDocStream ann =
     | SAnnPop (SimpleDocStream ann)
     deriving (Eq, Ord, Show, Generic)
 
+-- | Remove all trailing space characters.
+--
+-- This has some performance impact, because it does an entire additional pass
+-- over the 'SimpleDocStream'.
+--
+-- No trimming will be done inside annotations, which are considered to contain
+-- no (trimmable) whitespace, since the annotation might actually be /about/ the
+-- whitespace, for example a renderer that colors the background of trailing
+-- whitespace, as e.g. @git diff@ can be configured to do.
+removeTrailingWhitespace :: SimpleDocStream ann -> SimpleDocStream ann
+removeTrailingWhitespace = go (WssWithheldWhitespace Nothing 0)
+  where
+    commitSpaces :: Maybe Int -> Int -> SimpleDocStream ann -> SimpleDocStream ann
+    commitSpaces mWithheldNewline withheldSpaces = nl . sp
+      where
+        nl = case mWithheldNewline of
+            Just withheldLine -> SLine withheldLine
+            Nothing -> id
+        sp = case withheldSpaces of
+            0 -> id
+            1 -> SChar ' '
+            n -> SText n (T.replicate n " ")
+
+    go :: WhitespaceStrippingState -> SimpleDocStream ann -> SimpleDocStream ann
+    go annLevel@(WssAnnotationLevel annotationLevel) = \case
+        SFail -> SFail
+        SEmpty -> SEmpty
+        SChar c rest -> SChar c (go annLevel rest)
+        SText textLength text rest -> SText textLength text (go annLevel rest)
+        SLine i rest -> SLine i (go annLevel rest)
+        SAnnPush ann rest -> SAnnPush ann (go (WssAnnotationLevel (annotationLevel+1)) rest)
+        SAnnPop rest
+            | annotationLevel > 1 -> SAnnPop (go (WssAnnotationLevel (annotationLevel-1)) rest)
+            | otherwise -> SAnnPop (go (WssWithheldWhitespace Nothing 0) rest)
+    go (WssWithheldWhitespace withheldLine withheldSpaces) = \case
+        SFail -> SFail
+        SEmpty -> SEmpty
+        SChar c rest
+            | c == ' ' -> go (WssWithheldWhitespace withheldLine (withheldSpaces+1)) rest
+            | otherwise -> commitSpaces withheldLine withheldSpaces (SChar c (go (WssWithheldWhitespace Nothing 0) rest))
+        SText textLength text rest ->
+            let stripped = T.dropWhileEnd (== ' ') text
+                strippedLength = T.length stripped
+                trailingLength = textLength - strippedLength
+                isOnlySpace = strippedLength == 0
+            in if isOnlySpace
+                then go (WssWithheldWhitespace withheldLine (withheldSpaces + textLength)) rest
+                else commitSpaces withheldLine withheldSpaces (SText strippedLength stripped (go (WssWithheldWhitespace Nothing trailingLength) rest))
+        SLine i rest -> go (WssWithheldWhitespace (Just i) 0) rest
+        SAnnPush ann rest -> commitSpaces withheldLine withheldSpaces (SAnnPush ann (go (WssAnnotationLevel 1) rest))
+        SAnnPop _ -> error "Tried skipping spaces in unannotated data! Please report this as a bug in 'prettyprinter'."
+
+data WhitespaceStrippingState
+    = WssAnnotationLevel !Int
+    | WssWithheldWhitespace (Maybe Int) !Int
+
+
+
 -- | Alter the document’s annotations.
 --
 -- This instance makes 'SimpleDocStream' more flexible (because it can be used in

prettyprinter/test/Testsuite/Main.hs

@@ -22,6 +22,8 @@ import Test.Tasty
 import Test.Tasty.HUnit
 import Test.Tasty.QuickCheck
 
+import StripTrailingSpace
+
 #if !(APPLICATIVE_MONAD)
 import Control.Applicative
 #endif
@@ -42,6 +44,7 @@ tests = testGroup "Tests"
         , testProperty "Deep fusion does not change rendering"
                        (fusionDoesNotChangeRendering Deep)
         ]
+    , testStripTrailingSpace
     , testGroup "Performance tests"
         [ testCase "Grouping performance"
                    groupingPerformance

prettyprinter/test/Testsuite/StripTrailingSpace.hs

@@ -0,0 +1,88 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module StripTrailingSpace (testStripTrailingSpace) where
+
+
+
+import           Data.Text (Text)
+import qualified Data.Text as T
+
+import Data.Text.Prettyprint.Doc
+import Data.Text.Prettyprint.Doc.Render.Util.StackMachine
+
+import Test.Tasty
+import Test.Tasty.HUnit
+
+box :: Text -> Text
+box singleLine = unlines'
+    [ "┌─" <> T.replicate (T.length singleLine) "─" <> "─┐"
+    , "│ " <> singleLine <> " │"
+    , "└─" <> T.replicate (T.length singleLine) "─" <> "─┘"
+    ]
+
+bbox :: Text -> Text
+bbox singleLine = unlines'
+    [ "╔═" <> T.replicate (T.length singleLine) "═" <> "═╗"
+    , "║ " <> singleLine <> " ║"
+    , "╚═" <> T.replicate (T.length singleLine) "═" <> "═╝"
+    ]
+
+testStripTrailingSpace :: TestTree
+testStripTrailingSpace = testGroup "Stripping trailing space"
+    [ testCase "No trailing space"
+               (testStripping "No trailing space at all")
+    , testCase "Single trailing space character"
+               (testStripping ("Single trailing character" <> " "))
+    , testCase "Space character inside"
+               (testStripping ("Space character" <> " " <> "inside"))
+    , testCase "Obvious trailing spaces"
+               (testStripping ("Obvious trailing space" <> "   "))
+    , testCase "Multiple spaces inside"
+               (testStripping ("Multiple spaces" <> "    " <> "inside"))
+    , testCase "Whitespace inside text"
+               (testStripping ("Whitespace inside text   "))
+    , testCase "Indented blank line"
+               (testStripping (nest 4 (vcat ["Indented blank line", "", "<end>"])))
+    , testCase "Multiple indented blank lines"
+               (testStripping (nest 4 (vcat ["Indented blank lines", "", "", "", "<end>"])))
+    , testCase "Annotation"
+               (testStripping (annotate () "Annotation with trailing space   "))
+    , testCase "Document with annotation"
+               (testStripping ("Here comes an" <> annotate () "annotation   " <> "and some trailing space again  " <> "  "))
+    , testCase "Nested annotations"
+               (testStripping ("A " <> annotate () ("nested   " <> annotate () "annotation ") <> "and some trailing space again  " <> "  "))
+    , testCase "Stress test"
+               (testStripping (nest 4 (vcat ["Stress test", "", "" <> annotate () "hello ", "", "world " <> "   ", annotate () "", "", "end"])))
+    ]
+
+testStripping :: Doc ann -> Assertion
+testStripping doc = case hasTrailingWhitespace (render removeTrailingWhitespace doc) of
+    False -> pure ()
+    True  -> (assertFailure . T.unpack . T.unlines)
+        [ bbox "Input is not stripped correctly!"
+        , ""
+        , box "Rendered/stripped:"
+        , (revealSpaces . render removeTrailingWhitespace) doc
+        , ""
+        , box "Rendered/unstripped:"
+        , (revealSpaces . render id) doc
+        , ""
+        , box "Rendered/unstripped, later stripped via Text API:"
+        , (revealSpaces . removeTrailingSpaceText . render id) doc ]
+  where
+
+    render :: (SimpleDocStream ann -> SimpleDocStream ann) -> Doc ann -> Text
+    render f = renderSimplyDecorated id (const "<ann>") (const "</ann>") . f . layoutPretty defaultLayoutOptions
+
+    removeTrailingSpaceText :: Text -> Text
+    removeTrailingSpaceText = unlines' . map T.stripEnd . T.lines
+
+    hasTrailingWhitespace :: Text -> Bool
+    hasTrailingWhitespace x = removeTrailingSpaceText x /= x
+
+    revealSpaces :: Text -> Text
+    revealSpaces = T.map (\x -> if x == ' ' then '␣' else x)
+
+-- Text.unlines appends a trailing whitespace, so T.unlines . T.lines /= id
+unlines' :: [Text] -> Text
+unlines' = T.intercalate (T.singleton '\n')