Fix terminal corruption bug and clean up SDoc interface.

- Fix #13076 by wrapping `printDoc_` so that the terminal color is
  reset even if an exception occurs.

- Add `printSDoc`, `printSDocLn`, and `bufLeftRenderSDoc` to keep `SDoc`
  values abstract (they are wrappers of `printDoc_`, `printDoc`, and
  `bufLeftRender` respectively).

- Remove unused function: `printForAsm`

Test Plan: manual

Reviewers: RyanGlScott, austin, dfeuer, bgamari

Reviewed By: dfeuer, bgamari

Subscribers: dfeuer, mpickering, thomie

Differential Revision: https://phabricator.haskell.org/D2932

GHC Trac Issues: #13076

4 files changed, 51 insertions, 27 deletions

diff --git a/compiler/llvmGen/LlvmCodeGen/Base.hs b/compiler/llvmGen/LlvmCodeGen/Base.hs
index 824a8595fc..eb4a863e5a 100644
--- a/compiler/llvmGen/LlvmCodeGen/Base.hs
+++ b/compiler/llvmGen/LlvmCodeGen/Base.hs
@@ -46,7 +46,6 @@ import DynFlags
 import FastString
 import Cmm              hiding ( succ )
 import Outputable as Outp
-import qualified Pretty as Prt
 import Platform
 import UniqFM
 import Unique
@@ -330,8 +329,8 @@ renderLlvm sdoc = do
     -- Write to output
     dflags <- getDynFlags
     out <- getEnv envOutput
-    let doc = Outp.withPprStyleDoc dflags (Outp.mkCodeStyle Outp.CStyle) sdoc
-    liftIO $ Prt.bufLeftRender out doc
+    liftIO $ Outp.bufLeftRenderSDoc dflags out
+               (Outp.mkCodeStyle Outp.CStyle) sdoc
 
     -- Dump, if requested
     dumpIfSetLlvm Opt_D_dump_llvm "LLVM Code" sdoc
diff --git a/compiler/main/DynFlags.hs b/compiler/main/DynFlags.hs
index 0bc119a783..8d50e01905 100644
--- a/compiler/main/DynFlags.hs
+++ b/compiler/main/DynFlags.hs
@@ -1697,8 +1697,8 @@ defaultFatalMessager = hPutStrLn stderr
 defaultLogAction :: LogAction
 defaultLogAction dflags reason severity srcSpan style msg
     = case severity of
-      SevOutput      -> printSDoc msg style
-      SevDump        -> printSDoc (msg $$ blankLine) style
+      SevOutput      -> printOut msg style
+      SevDump        -> printOut (msg $$ blankLine) style
       SevInteractive -> putStrSDoc msg style
       SevInfo        -> printErrs msg style
       SevFatal       -> printErrs msg style
@@ -1714,7 +1714,7 @@ defaultLogAction dflags reason severity srcSpan style msg
                            -- whereas converting to string first and using
                            -- hPutStr would just emit the low 8 bits of
                            -- each unicode char.
-    where printSDoc  = defaultLogActionHPrintDoc  dflags stdout
+    where printOut   = defaultLogActionHPrintDoc  dflags stdout
           printErrs  = defaultLogActionHPrintDoc  dflags stderr
           putStrSDoc = defaultLogActionHPutStrDoc dflags stdout
           -- Pretty print the warning flag, if any (#10752)
@@ -1731,17 +1731,16 @@ defaultLogAction dflags reason severity srcSpan style msg
                         groups -> " (in " ++ intercalate ", " (map ("-W"++) groups) ++ ")"
               | otherwise = ""
 
+-- | Like 'defaultLogActionHPutStrDoc' but appends an extra newline.
 defaultLogActionHPrintDoc :: DynFlags -> Handle -> SDoc -> PprStyle -> IO ()
 defaultLogActionHPrintDoc dflags h d sty
  = defaultLogActionHPutStrDoc dflags h (d $$ text "") sty
-      -- Adds a newline
 
 defaultLogActionHPutStrDoc :: DynFlags -> Handle -> SDoc -> PprStyle -> IO ()
 defaultLogActionHPutStrDoc dflags h d sty
-  = Pretty.printDoc_ Pretty.PageMode (pprCols dflags) h doc
-  where   -- Don't add a newline at the end, so that successive
-          -- calls to this log-action can output all on the same line
-    doc = runSDoc d (initSDocContext dflags sty)
+  -- Don't add a newline at the end, so that successive
+  -- calls to this log-action can output all on the same line
+  = printSDoc Pretty.PageMode dflags h sty d
 
 newtype FlushOut = FlushOut (IO ())
 
diff --git a/compiler/nativeGen/AsmCodeGen.hs b/compiler/nativeGen/AsmCodeGen.hs
index 0a15638cc4..7cc7a2804d 100644
--- a/compiler/nativeGen/AsmCodeGen.hs
+++ b/compiler/nativeGen/AsmCodeGen.hs
@@ -346,8 +346,7 @@ finishNativeGen dflags modLoc bufh@(BufHandle _ _ h) us ngs
           dump_stats (Linear.pprStats (concat (ngs_natives ngs)) linearStats)
 
         -- write out the imports
-        Pretty.printDoc Pretty.LeftMode (pprCols dflags) h
-                $ withPprStyleDoc dflags (mkCodeStyle AsmStyle)
+        printSDocLn Pretty.LeftMode dflags h (mkCodeStyle AsmStyle)
                 $ makeImportsDoc dflags (concat (ngs_imports ngs))
         return us'
   where
@@ -481,8 +480,8 @@ cmmNativeGens dflags this_mod modLoc ncgImpl h dbgMap = go
 emitNativeCode :: DynFlags -> BufHandle -> SDoc -> IO ()
 emitNativeCode dflags h sdoc = do
 
-        {-# SCC "pprNativeCode" #-} Pretty.bufLeftRender h
-                $ withPprStyleDoc dflags (mkCodeStyle AsmStyle) sdoc
+        {-# SCC "pprNativeCode" #-} bufLeftRenderSDoc dflags h
+                                      (mkCodeStyle AsmStyle) sdoc
 
         -- dump native code
         dumpIfSet_dyn dflags
diff --git a/compiler/utils/Outputable.hs b/compiler/utils/Outputable.hs
index 371856f5ea..93afffefe2 100644
--- a/compiler/utils/Outputable.hs
+++ b/compiler/utils/Outputable.hs
@@ -43,7 +43,8 @@ module Outputable (
         colWhiteFg, colBinder, colCoerc, colDataCon, colType,
 
         -- * Converting 'SDoc' into strings and outputing it
-        printForC, printForAsm, printForUser, printForUserPartWay,
+        printSDoc, printSDocLn, printForUser, printForUserPartWay,
+        printForC, bufLeftRenderSDoc,
         pprCode, mkCodeStyle,
         showSDoc, showSDocUnsafe, showSDocOneLine,
         showSDocForUser, showSDocDebug, showSDocDump, showSDocDumpOneLine,
@@ -94,6 +95,7 @@ import {-# SOURCE #-}   Module( UnitId, Module, ModuleName, moduleName )
 import {-# SOURCE #-}   OccName( OccName )
 import {-# SOURCE #-}   StaticFlags( opt_PprStyle_Debug, opt_NoDebugOutput )
 
+import BufWrite (BufHandle)
 import FastString
 import qualified Pretty
 import Util
@@ -103,6 +105,7 @@ import Panic
 import GHC.Serialized
 import GHC.LanguageExtensions (Extension)
 
+import Control.Exception (finally)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as BS
 import Data.Char
@@ -298,6 +301,11 @@ code (either C or assembly), or generating interface files.
 ************************************************************************
 -}
 
+-- | Represents a pretty-printable document.
+--
+-- To display an 'SDoc', use 'printSDoc', 'printSDocLn', 'bufLeftRenderSDoc',
+-- or 'renderWithStyle'.  Avoid calling 'runSDoc' directly as it breaks the
+-- abstraction layer.
 newtype SDoc = SDoc { runSDoc :: SDocContext -> Doc }
 
 data SDocContext = SDC
@@ -320,6 +328,9 @@ initSDocContext dflags sty = SDC
 withPprStyle :: PprStyle -> SDoc -> SDoc
 withPprStyle sty d = SDoc $ \ctxt -> runSDoc d ctxt{sdocStyle=sty}
 
+-- | This is not a recommended way to render 'SDoc', since it breaks the
+-- abstraction layer of 'SDoc'.  Prefer to use 'printSDoc', 'printSDocLn',
+-- 'bufLeftRenderSDoc', or 'renderWithStyle' instead.
 withPprStyleDoc :: DynFlags -> PprStyle -> SDoc -> Doc
 withPprStyleDoc dflags sty d = runSDoc d (initSDocContext dflags sty)
 
@@ -409,27 +420,43 @@ ifPprDebug d = SDoc $ \ctx ->
         SDC{sdocStyle=PprDebug} -> runSDoc d ctx
         _                       -> Pretty.empty
 
+-- | The analog of 'Pretty.printDoc_' for 'SDoc', which tries to make sure the
+--   terminal doesn't get screwed up by the ANSI color codes if an exception
+--   is thrown during pretty-printing.
+printSDoc :: Mode -> DynFlags -> Handle -> PprStyle -> SDoc -> IO ()
+printSDoc mode dflags handle sty doc =
+  Pretty.printDoc_ mode cols handle (runSDoc doc ctx)
+    `finally`
+      Pretty.printDoc_ mode cols handle (runSDoc (coloured colReset empty) ctx)
+  where
+    cols = pprCols dflags
+    ctx = initSDocContext dflags sty
+
+-- | Like 'printSDoc' but appends an extra newline.
+printSDocLn :: Mode -> DynFlags -> Handle -> PprStyle -> SDoc -> IO ()
+printSDocLn mode dflags handle sty doc =
+  printSDoc mode dflags handle sty (doc $$ text "")
+
 printForUser :: DynFlags -> Handle -> PrintUnqualified -> SDoc -> IO ()
 printForUser dflags handle unqual doc
-  = Pretty.printDoc PageMode (pprCols dflags) handle
-      (runSDoc doc (initSDocContext dflags (mkUserStyle unqual AllTheWay)))
+  = printSDocLn PageMode dflags handle (mkUserStyle unqual AllTheWay) doc
 
 printForUserPartWay :: DynFlags -> Handle -> Int -> PrintUnqualified -> SDoc
                     -> IO ()
 printForUserPartWay dflags handle d unqual doc
-  = Pretty.printDoc PageMode (pprCols dflags) handle
-      (runSDoc doc (initSDocContext dflags (mkUserStyle unqual (PartWay d))))
+  = printSDocLn PageMode dflags handle (mkUserStyle unqual (PartWay d)) doc
 
--- printForC, printForAsm do what they sound like
+-- | Like 'printSDocLn' but specialized with 'LeftMode' and
+-- @'PprCode' 'CStyle'@.  This is typically used to output C-- code.
 printForC :: DynFlags -> Handle -> SDoc -> IO ()
 printForC dflags handle doc =
-  Pretty.printDoc LeftMode (pprCols dflags) handle
-    (runSDoc doc (initSDocContext dflags (PprCode CStyle)))
+  printSDocLn LeftMode dflags handle (PprCode CStyle) doc
 
-printForAsm :: DynFlags -> Handle -> SDoc -> IO ()
-printForAsm dflags handle doc =
-  Pretty.printDoc LeftMode (pprCols dflags) handle
-    (runSDoc doc (initSDocContext dflags (PprCode AsmStyle)))
+-- | An efficient variant of 'printSDoc' specialized for 'LeftMode' that
+-- outputs to a 'BufHandle'.
+bufLeftRenderSDoc :: DynFlags -> BufHandle -> PprStyle -> SDoc -> IO ()
+bufLeftRenderSDoc dflags bufHandle sty doc =
+  Pretty.bufLeftRender bufHandle (runSDoc doc (initSDocContext dflags sty))
 
 pprCode :: CodeStyle -> SDoc -> SDoc
 pprCode cs d = withPprStyle (PprCode cs) d