Fix it less wrong

sjakobi · Nov 5, 2019 · b825fbf · b825fbf
1 parent 6ffbc89
commit b825fbf
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/prettyprinter/src/Data/Text/Prettyprint/Doc/Internal.hs b/prettyprinter/src/Data/Text/Prettyprint/Doc/Internal.hs
@@ -1468,7 +1468,7 @@ removeTrailingWhitespace = go (RecordedWhitespace [] 0)
     commitSpaces [] 0 = id
     commitSpaces [] 1 = SChar ' '
     commitSpaces [] n = SText n (T.replicate n " ")
-    commitSpaces [_] n = SLine 0 . commitSpaces [] n
+    commitSpaces [i] n = SLine i . commitSpaces [] n
     commitSpaces (_:is) n = SLine 0 . commitSpaces is n
 
     go :: WhitespaceStrippingState -> SimpleDocStream ann -> SimpleDocStream ann
@@ -1494,7 +1494,7 @@ removeTrailingWhitespace = go (RecordedWhitespace [] 0)
         SChar c rest
             | c == ' ' -> go (RecordedWhitespace withheldLines (withheldSpaces+1)) rest
             | otherwise -> commitSpaces
-                               withheldLines
+                               (reverse withheldLines)
                                withheldSpaces
                                (SChar c (go (RecordedWhitespace [] 0) rest))
         SText textLength text rest ->

diff --git a/prettyprinter/test/Testsuite/Main.hs b/prettyprinter/test/Testsuite/Main.hs
@@ -74,8 +74,12 @@ tests = testGroup "Tests"
             , testCase "Reduce to single trailing newline"
                        removeTrailingWhitespaceInTrailingNewlines
             ]
-        , testCase "removeTrailingWhitespace restores indentation in the wrong spot (#93)"
-                   removeTrailingWhitespaceDontRestoreIndentation
+        , testGroup "removeTrailingWhitespace restores indentation in the wrong spot (#93)"
+            [ testCase "Don't restore indentation in the wrong spot"
+                       removeTrailingWhitespaceDontRestoreIndentationInTheWrongSpot
+            , testCase "Preserve leading indentation"
+                       removeTrailingWhitespacePreserveIndentation
+            ]
         ]
     ]
 
@@ -280,9 +284,15 @@ badFallbackAlign
         expected = "/Fallback\n Fallback\n Too wide!!!!!"
     in assertEqual "" expected actual
 
-removeTrailingWhitespaceDontRestoreIndentation :: Assertion
-removeTrailingWhitespaceDontRestoreIndentation
+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)