Make putStrLn more atomic with line or block buffering

src/Data/Text/IO.hs

@@ -51,7 +51,7 @@ import qualified Control.Exception as E
 import Control.Monad (liftM2, when)
 import Data.IORef (readIORef, writeIORef)
 import qualified Data.Text as T
-import Data.Text.Internal.Fusion (stream)
+import Data.Text.Internal.Fusion (stream, streamLn)
 import Data.Text.Internal.Fusion.Types (Step(..), Stream(..))
 import Data.Text.Internal.IO (hGetLineWith, readChunk)
 import GHC.IO.Buffer (Buffer(..), BufferState(..), RawCharBuffer, CharBuffer,
@@ -174,13 +174,15 @@ hGetLine = hGetLineWith T.concat
 
 -- | Write a string to a handle.
 hPutStr :: Handle -> Text -> IO ()
+hPutStr h = hPutStr' h . stream
+
 -- This function is lifted almost verbatim from GHC.IO.Handle.Text.
-hPutStr h t = do
+hPutStr' :: Handle -> Stream Char -> IO ()
+hPutStr' h str = do
   (buffer_mode, nl) <-
        wantWritableHandle "hPutStr" h $ \h_ -> do
                      bmode <- getSpareBuffer h_
                      return (bmode, haOutputNL h_)
-  let str = stream t
   case buffer_mode of
      (NoBuffering, _)        -> hPutChars h str
      (LineBuffering, buf)    -> writeLines h nl buf str
@@ -276,7 +278,7 @@ commitBuffer hdl !raw !sz !count flush release =
 
 -- | Write a string to a handle, followed by a newline.
 hPutStrLn :: Handle -> Text -> IO ()
-hPutStrLn h t = hPutStr h t >> hPutChar h '\n'
+hPutStrLn h = hPutStr' h . streamLn
 
 -- | The 'interact' function takes a function of type @Text -> Text@
 -- as its argument. The entire input from the standard input device is

src/Data/Text/Internal/Fusion.hs

@@ -25,6 +25,7 @@ module Data.Text.Internal.Fusion
 
     -- * Creation and elimination
     , stream
+    , streamLn
     , unstream
     , reverseStream
 
@@ -49,7 +50,7 @@ module Data.Text.Internal.Fusion
     , countChar
     ) where
 
-import Prelude (Bool(..), Char, Maybe(..), Monad(..), Int,
+import Prelude (Bool(..), Char, Eq(..), Maybe(..), Monad(..), Int,
                 Num(..), Ord(..), ($),
                 otherwise)
 import Data.Bits (shiftL, shiftR)
@@ -98,6 +99,33 @@ stream (Text arr off len) = Stream next off (betweenSize (len `shiftR` 2) len)
               _ -> U8.chr4 n0 n1 n2 n3
 {-# INLINE [0] stream #-}
 
+-- | /O(n)/ @'streamLn' t = 'stream' (t <> \'\\n\')@
+streamLn ::
+#if defined(ASSERTS)
+  HasCallStack =>
+#endif
+  Text -> Stream Char
+streamLn (Text arr off len) = Stream next off (betweenSize (len `shiftR` 2) (len + 1))
+    where
+      !end = off+len
+      next !i
+          | i > end = Done
+          | i == end  = Yield '\n' (i + 1)
+          | otherwise = Yield chr (i + l)
+          where
+            n0 = A.unsafeIndex arr i
+            n1 = A.unsafeIndex arr (i + 1)
+            n2 = A.unsafeIndex arr (i + 2)
+            n3 = A.unsafeIndex arr (i + 3)
+
+            l  = U8.utf8LengthByLeader n0
+            chr = case l of
+              1 -> unsafeChr8 n0
+              2 -> U8.chr2 n0 n1
+              3 -> U8.chr3 n0 n1 n2
+              _ -> U8.chr4 n0 n1 n2 n3
+{-# INLINE [0] streamLn #-}
+
 -- | /O(n)/ Converts 'Text' into a 'Stream' 'Char', but iterates
 -- backwards through the text.
 --

src/Data/Text/Internal/Lazy/Fusion.hs

@@ -17,6 +17,7 @@
 module Data.Text.Internal.Lazy.Fusion
     (
       stream
+    , streamLn
     , unstream
     , unstreamChunks
     , length
@@ -56,6 +57,22 @@ stream text = Stream next (text :*: 0) unknownSize
         where Iter c d = iter t i
 {-# INLINE [0] stream #-}
 
+-- | /O(n)/ @'streamLn' t = 'stream' (t <> \'\\n\')@
+streamLn ::
+#if defined(ASSERTS)
+  HasCallStack =>
+#endif
+  Text -> Stream Char
+streamLn text = Stream next (text :*: 0) unknownSize
+  where
+    next (Empty :*: 0) = Yield '\n' (Empty :*: 1)
+    next (Empty :*: _) = Done
+    next (txt@(Chunk t@(I.Text _ _ len) ts) :*: i)
+        | i >= len  = next (ts :*: 0)
+        | otherwise = Yield c (txt :*: i+d)
+        where Iter c d = iter t i
+{-# INLINE [0] streamLn #-}
+
 -- | /O(n)/ Convert a 'Stream Char' into a 'Text', using the given
 -- chunk size.
 unstreamChunks ::

src/Data/Text/Lazy/IO.hs

@@ -50,6 +50,7 @@ import Control.Monad (when)
 import Data.IORef (readIORef)
 import Data.Text.Internal.IO (hGetLineWith, readChunk)
 import Data.Text.Internal.Lazy (chunk, empty)
+import Data.Text.Internal.Lazy.Fusion (stream, streamLn)
 import GHC.IO.Buffer (isEmptyBuffer)
 import GHC.IO.Exception (IOException(..), IOErrorType(..), ioException)
 import GHC.IO.Handle.Internals (augmentIOError, hClose_help,
@@ -129,11 +130,11 @@ hGetLine = hGetLineWith L.fromChunks
 
 -- | Write a string to a handle.
 hPutStr :: Handle -> Text -> IO ()
-hPutStr h = mapM_ (T.hPutStr h) . L.toChunks
+hPutStr h = T.hPutStr h . stream
 
 -- | Write a string to a handle, followed by a newline.
 hPutStrLn :: Handle -> Text -> IO ()
-hPutStrLn h t = hPutStr h t >> hPutChar h '\n'
+hPutStrLn h = hPutStr h . streamLn
 
 -- | The 'interact' function takes a function of type @Text -> Text@
 -- as its argument. The entire input from the standard input device is