Merge branch 'sjakobi-93-rtw-indentation'

sjakobi · Nov 5, 2019 · 36b92c5 · 36b92c5
2 parents e660360 + cce4098
commit 36b92c5
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 9 deletions.
diff --git a/prettyprinter/src/Data/Text/Prettyprint/Doc/Internal.hs b/prettyprinter/src/Data/Text/Prettyprint/Doc/Internal.hs
@@ -1460,16 +1460,19 @@ data SimpleDocStream ann =
 removeTrailingWhitespace :: SimpleDocStream ann -> SimpleDocStream ann
 removeTrailingWhitespace = go (RecordedWhitespace [] 0)
   where
-    commitSpaces
+    commitWhitespace
         :: [Int] -- Withheld lines
         -> Int -- Withheld spaces
         -> SimpleDocStream ann
         -> SimpleDocStream ann
-    commitSpaces [] 0 = id
-    commitSpaces [] 1 = SChar ' '
-    commitSpaces [] n = SText n (T.replicate n " ")
-    commitSpaces [i] n = SLine i . commitSpaces [] n
-    commitSpaces (_:is) n = SLine 0 . commitSpaces is n
+    commitWhitespace is0 n0 = commitLines is0 . commitSpaces n0
+      where
+        commitLines [] = id
+        commitLines (i:is) = foldr (\_ f -> SLine 0 . f) (SLine i) is
+
+        commitSpaces 0 = id
+        commitSpaces 1 = SChar ' '
+        commitSpaces n = SText n (T.replicate n " ")
 
     go :: WhitespaceStrippingState -> SimpleDocStream ann -> SimpleDocStream ann
     -- We do not strip whitespace inside annotated documents, since it might
@@ -1493,7 +1496,7 @@ removeTrailingWhitespace = go (RecordedWhitespace [] 0)
         SEmpty -> foldr (\_i sds' -> SLine 0 sds') SEmpty withheldLines
         SChar c rest
             | c == ' ' -> go (RecordedWhitespace withheldLines (withheldSpaces+1)) rest
-            | otherwise -> commitSpaces
+            | otherwise -> commitWhitespace
                                withheldLines
                                withheldSpaces
                                (SChar c (go (RecordedWhitespace [] 0) rest))
@@ -1504,14 +1507,14 @@ removeTrailingWhitespace = go (RecordedWhitespace [] 0)
                 isOnlySpace = strippedLength == 0
             in if isOnlySpace
                 then go (RecordedWhitespace withheldLines (withheldSpaces + textLength)) rest
-                else commitSpaces
+                else commitWhitespace
                         withheldLines
                         withheldSpaces
                         (SText strippedLength
                                stripped
                                (go (RecordedWhitespace [] trailingLength) rest))
         SLine i rest -> go (RecordedWhitespace (i:withheldLines) 0) rest
-        SAnnPush ann rest -> commitSpaces
+        SAnnPush ann rest -> commitWhitespace
                                  withheldLines
                                  withheldSpaces
                                  (SAnnPush ann (go (AnnotationLevel 1) rest))

diff --git a/prettyprinter/test/Testsuite/Main.hs b/prettyprinter/test/Testsuite/Main.hs
@@ -74,6 +74,12 @@ tests = testGroup "Tests"
             , testCase "Reduce to single trailing newline"
                        removeTrailingWhitespaceInTrailingNewlines
             ]
+        , testGroup "removeTrailingWhitespace restores indentation in the wrong spot (#93)"
+            [ testCase "Don't restore indentation in the wrong spot"
+                       removeTrailingWhitespaceDontRestoreIndentationInTheWrongSpot
+            , testCase "Preserve leading indentation"
+                       removeTrailingWhitespacePreserveIndentation
+            ]
         ]
     ]
 
@@ -277,3 +283,16 @@ badFallbackAlign
         actual = renderStrict (layoutSmart (LayoutOptions (AvailablePerLine 12 1)) doc)
         expected = "/Fallback\n Fallback\n Too wide!!!!!"
     in assertEqual "" expected actual
+
+removeTrailingWhitespaceDontRestoreIndentationInTheWrongSpot :: Assertion
+removeTrailingWhitespaceDontRestoreIndentationInTheWrongSpot
+  = let sdoc :: SimpleDocStream ()
+        sdoc = SLine 2 (SLine 0 (SChar 'x' SEmpty))
+        sdoc' = SLine 0 (SLine 0 (SChar 'x' SEmpty))
+    in assertEqual "" sdoc' (removeTrailingWhitespace sdoc)
+
+removeTrailingWhitespacePreserveIndentation :: Assertion
+removeTrailingWhitespacePreserveIndentation
+  = let sdoc :: SimpleDocStream ()
+        sdoc = SLine 2 (SChar 'x' SEmpty)
+    in assertEqual "" sdoc (removeTrailingWhitespace sdoc)