Make 'Terminal.renderLazy' lazy (#176)

* Make 'Terminal.renderLazy' lazy

* Format extension pragmas

* Add module header

* Reorganise imports

* De-wibbling

Fixes #175.

prettyprinter-ansi-terminal/bench/LargeOutput.hs

@@ -0,0 +1,159 @@
+{-# LANGUAGE DeriveGeneric     #-}
+{-# LANGUAGE LambdaCase        #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+
+-- | This benchmark is derived from the large-output benchmark in prettyprinter, but contains additional annotations.
+module Main (main) where
+
+import Prelude        ()
+import Prelude.Compat
+
+import           Control.DeepSeq
+import           Control.Monad.Compat
+import           Data.Char
+import           Data.Map                              (Map)
+import qualified Data.Map                              as M
+import           Data.Text                             (Text)
+import qualified Data.Text                             as T
+import qualified Data.Text.IO                          as T
+import qualified Data.Text.Lazy                        as TL
+import           Gauge
+import           GHC.Generics
+import           Prettyprinter
+import           Prettyprinter.Render.Terminal         as Terminal
+import qualified Prettyprinter.Render.Text             as Text
+import           Test.QuickCheck
+import           Test.QuickCheck.Gen
+import           Test.QuickCheck.Random
+
+
+
+newtype Program = Program Binds deriving (Show, Generic)
+newtype Binds = Binds (Map Text LambdaForm) deriving (Show, Generic)
+data LambdaForm = LambdaForm ![Text] ![Text] !Expr deriving (Show, Generic)
+data Expr
+    = Let Binds Expr
+    | Case Expr [Alt]
+    | AppF Text [Text]
+    | AppC Text [Text]
+    | AppP Text Text Text
+    | LitE Int
+    deriving (Show, Generic)
+data Alt = Alt Text [Text] Expr deriving (Show, Generic)
+
+instance NFData Program
+instance NFData Binds
+instance NFData LambdaForm
+instance NFData Expr
+instance NFData Alt
+
+instance Arbitrary Program where arbitrary = fmap Program arbitrary
+instance Arbitrary Binds where
+    arbitrary = do
+        NonEmpty xs <- arbitrary
+        pure (Binds (M.fromList xs))
+instance Arbitrary LambdaForm where
+    arbitrary = LambdaForm <$> fromTo 0 2 arbitrary <*> fromTo 0 2 arbitrary <*> arbitrary
+
+instance Arbitrary Expr where
+    arbitrary = (oneof . map scaled)
+        [ Let <$> arbitrary <*> arbitrary
+        , Case <$> arbitrary <*> (do NonEmpty xs <- arbitrary; pure xs)
+        , AppF <$> arbitrary <*> fromTo 0 3 arbitrary
+        , AppC <$> ucFirst arbitrary <*> fromTo 0 3 arbitrary
+        , AppP <$> arbitrary <*> arbitrary <*> arbitrary
+        , LitE <$> arbitrary ]
+instance Arbitrary Alt where arbitrary = Alt <$> ucFirst arbitrary <*> fromTo 0 3 arbitrary <*> arbitrary
+instance Arbitrary Text where
+    arbitrary = do
+        n <- choose (3,6)
+        str <- replicateM n (elements ['a'..'z'])
+        if str `elem` ["let", "in", "case", "of"]
+            then arbitrary
+            else pure (T.pack str)
+
+ucFirst :: Gen Text -> Gen Text
+ucFirst gen = do
+    x <- gen
+    case T.uncons x of
+        Nothing -> pure x
+        Just (t,ext) -> pure (T.cons (toUpper t) ext)
+
+anCol :: Color -> Doc AnsiStyle -> Doc AnsiStyle
+anCol = annotate . color
+
+prettyProgram :: Program -> Doc AnsiStyle
+prettyProgram (Program binds) = annotate italicized $ prettyBinds binds
+
+prettyBinds :: Binds -> Doc AnsiStyle
+prettyBinds (Binds bs) = align (vsep (map prettyBinding (M.assocs bs)))
+  where
+    prettyBinding (var, lambda) = pretty var <+> anCol Red "=" <+> prettyLambdaForm lambda
+
+prettyLambdaForm :: LambdaForm -> Doc AnsiStyle
+prettyLambdaForm (LambdaForm free bound body) = prettyExp . (<+> anCol Blue "->") . prettyBound . prettyFree $ anCol Blue "\\"
+  where
+    prettyFree | null free = id
+               | otherwise = (<> anCol Blue lparen <> hsep (map pretty free) <> anCol Blue rparen)
+    prettyBound | null bound = id
+                | null free = (<> hsep (map pretty bound))
+                | otherwise = (<+> hsep (map pretty bound))
+    prettyExp = (<+> prettyExpr body)
+
+prettyExpr :: Expr -> Doc AnsiStyle
+prettyExpr = \case
+    Let binds body ->
+        align (vsep [ anCol Red "let" <+> align (prettyBinds binds)
+                    , anCol Red "in" <+> prettyExpr body ])
+
+    Case scrutinee alts -> vsep
+        [ anCol Yellow "case" <+> prettyExpr scrutinee <+> anCol Yellow "of"
+        , indent 4 (align (vsep (map prettyAlt alts))) ]
+
+    AppF f [] -> annotate bold . anCol Green $ pretty f
+    AppF f args -> annotate bold . anCol Green $ pretty f <+> hsep (map pretty args)
+
+    AppC c [] -> annotate bold . anCol Green $ pretty c
+    AppC c args -> annotate bold . anCol Green $ pretty c <+> hsep (map pretty args)
+
+    AppP op x y -> annotate bold . anCol Green $ pretty op <+> pretty x <+> pretty y
+
+    LitE lit -> annotate bold . anCol Green $ pretty lit
+
+prettyAlt :: Alt -> Doc AnsiStyle
+prettyAlt (Alt con [] body) = pretty con <+> anCol Yellow "->" <+> prettyExpr body
+prettyAlt (Alt con args body) = pretty con <+> hsep (map pretty args) <+> anCol Yellow "->" <+> prettyExpr body
+
+scaled :: Gen a -> Gen a
+scaled = scale (\n -> n * 2 `quot` 3)
+
+fromTo :: Int -> Int -> Gen b -> Gen b
+fromTo a b gen = do
+    n <- choose (min a b, max a b)
+    resize n gen
+
+randomProgram
+    :: Int -- ^ Seed
+    -> Int -- ^ Generator size
+    -> Program
+randomProgram seed size = let MkGen gen = arbitrary in gen (mkQCGen seed) size
+
+main :: IO ()
+main = do
+    let prog = randomProgram 1 60
+        layoutOpts = defaultLayoutOptions { layoutPageWidth = Unbounded }
+        renderedProg = (renderLazy . layoutPretty layoutOpts . prettyProgram) prog
+        (progLines, progWidth) = let l = TL.lines renderedProg in (length l, maximum (map TL.length l))
+    putDoc ("Program size:" <+> pretty progLines <+> "lines, maximum width:" <+> pretty progWidth)
+
+    let render :: (SimpleDocStream AnsiStyle -> TL.Text) -> Program -> TL.Text
+        render r = r . layoutPretty layoutOpts . prettyProgram
+
+    rnf prog `seq` T.putStrLn "Starting benchmark…"
+
+    defaultMain
+        [ bench "prettyprinter-ansi-terminal" $ nf (render Terminal.renderLazy) prog
+        , bench "prettyprinter" $ nf (render Text.renderLazy) prog
+        ]

prettyprinter-ansi-terminal/prettyprinter-ansi-terminal.cabal

@@ -62,3 +62,21 @@ test-suite doctest
     if impl (ghc < 7.10)
         buildable: False
         -- Doctest cannot search folders in old versions it seems :-(
+
+benchmark large-output
+    build-depends:
+          base >= 4.5 && < 5
+        , base-compat >=0.9.3 && <0.12
+        , containers
+        , deepseq
+        , gauge >= 0.2
+        , prettyprinter
+        , prettyprinter-ansi-terminal
+        , QuickCheck >= 2.7
+        , text
+
+    hs-source-dirs:      bench
+    main-is:             LargeOutput.hs
+    ghc-options:         -O2 -threaded -rtsopts -with-rtsopts=-N -Wall
+    default-language:    Haskell2010
+    type:                exitcode-stdio-1.0

prettyprinter-ansi-terminal/src/Prettyprinter/Render/Terminal/Internal.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP               #-}
+{-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE OverloadedStrings #-}
 
 {-# OPTIONS_HADDOCK not-home #-}
@@ -39,10 +40,8 @@ module Prettyprinter.Render.Terminal.Internal (
 
 
 import           Control.Applicative
-import           Control.Monad.ST
 import           Data.IORef
 import           Data.Maybe
-import           Data.STRef
 import           Data.Text              (Text)
 import qualified Data.Text              as T
 import qualified Data.Text.IO           as T
@@ -149,47 +148,33 @@ underlined = mempty { ansiUnderlining = Just Underlined }
 --
 -- Run the above via @echo -e '...'@ in your terminal to see the coloring.
 renderLazy :: SimpleDocStream AnsiStyle -> TL.Text
-renderLazy sdoc = runST (do
-    styleStackRef <- newSTRef [mempty]
-    outputRef <- newSTRef mempty
-
-    let push x = modifySTRef' styleStackRef (x :)
-        unsafePeek = readSTRef styleStackRef >>= \tok -> case tok of
+renderLazy =
+    let push x = (x :)
+        unsafePeek = \case
             []  -> panicPeekedEmpty
-            x:_ -> pure x
-        unsafePop = readSTRef styleStackRef >>= \tok -> case tok of
+            x:_ -> x
+        unsafePop = \case
             []   -> panicPeekedEmpty
-            x:xs -> writeSTRef styleStackRef xs >> pure x
-        writeOutput x = modifySTRef outputRef (<> x)
+            x:xs -> (x, xs)
 
-    let go = \sds -> case sds of
+        go :: [AnsiStyle] -> SimpleDocStream AnsiStyle -> TLB.Builder
+        go s = \case
             SFail -> panicUncaughtFail
-            SEmpty -> pure ()
-            SChar c rest -> do
-                writeOutput (TLB.singleton c)
-                go rest
-            SText _ t rest -> do
-                writeOutput (TLB.fromText t)
-                go rest
-            SLine i rest -> do
-                writeOutput (TLB.singleton '\n' <> TLB.fromText (T.replicate i (T.singleton ' ')))
-                go rest
-            SAnnPush style rest -> do
-                currentStyle <- unsafePeek
-                let newStyle = style <> currentStyle
-                push newStyle
-                writeOutput (TLB.fromText (styleToRawText newStyle))
-                go rest
-            SAnnPop rest -> do
-                _currentStyle <- unsafePop
-                newStyle <- unsafePeek
-                writeOutput (TLB.fromText (styleToRawText newStyle))
-                go rest
-    go sdoc
-    readSTRef styleStackRef >>= \stack -> case stack of
-        []  -> panicStyleStackFullyConsumed
-        [_] -> fmap TLB.toLazyText (readSTRef outputRef)
-        xs  -> panicStyleStackNotFullyConsumed (length xs) )
+            SEmpty -> mempty
+            SChar c rest -> TLB.singleton c <> go s rest
+            SText _ t rest -> TLB.fromText t <> go s rest
+            SLine i rest -> TLB.singleton '\n' <> TLB.fromText (T.replicate i " ") <> go s rest
+            SAnnPush style rest ->
+                let currentStyle = unsafePeek s
+                    newStyle = style <> currentStyle
+                in  TLB.fromText (styleToRawText newStyle) <> go (push style s) rest
+            SAnnPop rest ->
+                let (_currentStyle, s') = unsafePop s
+                    newStyle = unsafePeek s'
+                in  TLB.fromText (styleToRawText newStyle) <> go s' rest
+
+    in  TLB.toLazyText . go [mempty]
+
 
 -- | @('renderIO' h sdoc)@ writes @sdoc@ to the handle @h@.
 --