Improved hPutStr lazy for small chunk sizes

src/Data/Text/IO.hs

@@ -1,5 +1,7 @@
-{-# LANGUAGE BangPatterns, CPP, RecordWildCards, ScopedTypeVariables #-}
+{-# LANGUAGE BangPatterns, RecordWildCards, ScopedTypeVariables #-}
 {-# LANGUAGE Trustworthy #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# OPTIONS_GHC -Wno-name-shadowing #-}
 -- |
 -- Module      : Data.Text.IO
 -- Copyright   : (c) 2009, 2010 Bryan O'Sullivan,
@@ -33,6 +35,9 @@ module Data.Text.IO
     , hGetChunk
     , hGetLine
     , hPutStr
+    , hPutStrInit
+    , hPutStr'
+    , commitBuffer
     , hPutStrLn
     -- * Special cases for standard input and output
     , interact
@@ -54,7 +59,7 @@ import qualified Data.Text as T
 import Data.Text.Internal.Fusion (stream)
 import Data.Text.Internal.Fusion.Types (Step(..), Stream(..))
 import Data.Text.Internal.IO (hGetLineWith, readChunk)
-import GHC.IO.Buffer (Buffer(..), BufferState(..), CharBufElem, CharBuffer,
+import GHC.IO.Buffer (Buffer(..), BufferState(..), CharBuffer,
                       RawCharBuffer, emptyBuffer, isEmptyBuffer, newCharBuffer,
                       writeCharBuf)
 import GHC.IO.Exception (IOException(ioe_type), IOErrorType(InappropriateType))
@@ -65,6 +70,7 @@ import GHC.IO.Handle.Types (BufferList(..), BufferMode(..), Handle__(..),
                             HandleType(..), Newline(..))
 import System.IO (hGetBuffering, hFileSize, hSetBuffering, hTell)
 import System.IO.Error (isEOFError)
+import Data.Functor (void)
 
 -- | The 'readFile' function reads a file and returns the contents of
 -- the file as a string.  The entire file is read strictly, as with
@@ -172,27 +178,38 @@ chooseGoodBuffering h = do
 hGetLine :: Handle -> IO Text
 hGetLine = hGetLineWith T.concat
 
+
 -- | Write a string to a handle.
 hPutStr :: Handle -> Text -> IO ()
--- This function is lifted almost verbatim from GHC.IO.Handle.Text.
-hPutStr h t = 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
-     (BlockBuffering _, buf)
-         | nl == CRLF        -> writeBlocksCRLF h buf str
-         | otherwise         -> writeBlocksRaw h buf str
+hPutStr h t = hPutStrInit h $ \mode buf nl -> do
+  (n, b) <- hPutStr' h mode nl buf 0 (stream t)
+  let Buffer{bufRaw,bufSize} = b
+  when (n /= 0) $ void $ commitBuffer h bufRaw bufSize n False True
+
+{-# INLINE hPutStrInit #-}
+hPutStrInit :: Handle -> (BufferMode -> CharBuffer -> Newline -> IO ()) -> IO ()
+hPutStrInit h f = do
+  (mode, buf, nl) <- wantWritableHandle "hPutStr" h $ \h_ -> do
+    (mode, buf) <- getSpareBuffer h_
+    let nl = haOutputNL h_
+    return (mode, buf, nl)
+  f mode buf nl
+
+{-# INLINE hPutStr' #-}
+hPutStr' :: Handle -> BufferMode -> Newline -> CharBuffer -> Int -> Stream Char -> IO (Int, CharBuffer)
+hPutStr' h mode nl = case mode of
+  NoBuffering      -> \_ _ -> hPutChars h 
+  LineBuffering    -> writeLines h nl
+  BlockBuffering _
+    | nl == CRLF   -> writeBlocksCRLF h
+    | otherwise    -> writeBlocksRaw h
 
-hPutChars :: Handle -> Stream Char -> IO ()
+{-# INLINE hPutChars #-}
+hPutChars :: Handle -> Stream Char -> IO (Int, CharBuffer)
 hPutChars h (Stream next0 s0 _len) = loop s0
   where
     loop !s = case next0 s of
-                Done       -> return ()
+                Done       -> return (0, error "no buffer for hPutChars")
                 Skip s'    -> loop s'
                 Yield x s' -> hPutChar h x >> loop s'
 
@@ -206,57 +223,61 @@ hPutChars h (Stream next0 s0 _len) = loop s0
 -- performance improvement.  Lifting out the raw/cooked newline
 -- handling gave a few more percent on top.
 
-writeLines :: Handle -> Newline -> Buffer CharBufElem -> Stream Char -> IO ()
-writeLines h nl buf0 (Stream next0 s0 _len) = outer s0 buf0
+{-# INLINE writeLines #-}
+writeLines :: Handle -> Newline -> CharBuffer -> Int -> Stream Char -> IO (Int, CharBuffer)
+writeLines h nl buf0 n (Stream next0 s0 _len) = outer s0 n buf0
  where
-  outer s1 Buffer{bufRaw=raw, bufSize=len} = inner s1 (0::Int)
+  outer s1 n buf@Buffer{bufRaw=raw, bufSize=len} = inner s1 n
    where
     inner !s !n =
       case next0 s of
-        Done -> commit n False{-no flush-} True{-release-} >> return ()
+        Done -> return (n, buf)
         Skip s' -> inner s' n
         Yield x s'
-          | n + 1 >= len -> commit n True{-needs flush-} False >>= outer s
+          | n + 1 >= len -> commit n True{-needs flush-} False >>= outer s 0
           | x == '\n'    -> do
                    n' <- if nl == CRLF
                          then do n1 <- writeCharBuf raw n '\r'
                                  writeCharBuf raw n1 '\n'
                          else writeCharBuf raw n x
-                   commit n' True{-needs flush-} False >>= outer s'
+                   commit n' True{-needs flush-} False >>= outer s' 0
           | otherwise    -> writeCharBuf raw n x >>= inner s'
     commit = commitBuffer h raw len
 
-writeBlocksCRLF :: Handle -> Buffer CharBufElem -> Stream Char -> IO ()
-writeBlocksCRLF h buf0 (Stream next0 s0 _len) = outer s0 buf0
+{-# INLINE writeBlocksCRLF #-}
+writeBlocksCRLF :: Handle -> CharBuffer -> Int -> Stream Char -> IO (Int, CharBuffer)
+writeBlocksCRLF h buf0 n (Stream next0 s0 _len) = outer s0 n buf0
  where
-  outer s1 Buffer{bufRaw=raw, bufSize=len} = inner s1 (0::Int)
+  outer s1 n buf@Buffer{bufRaw=raw, bufSize=len} = inner s1 n
    where
     inner !s !n =
       case next0 s of
-        Done -> commit n False{-no flush-} True{-release-} >> return ()
+        Done -> return (n, buf)
         Skip s' -> inner s' n
         Yield x s'
-          | n + 1 >= len -> commit n True{-needs flush-} False >>= outer s
+          | n + 1 >= len -> commit n True{-needs flush-} False >>= outer s 0
           | x == '\n'    -> do n1 <- writeCharBuf raw n '\r'
                                writeCharBuf raw n1 '\n' >>= inner s'
           | otherwise    -> writeCharBuf raw n x >>= inner s'
     commit = commitBuffer h raw len
 
-writeBlocksRaw :: Handle -> Buffer CharBufElem -> Stream Char -> IO ()
-writeBlocksRaw h buf0 (Stream next0 s0 _len) = outer s0 buf0
+{-# INLINE writeBlocksRaw #-}
+writeBlocksRaw :: Handle -> CharBuffer -> Int -> Stream Char -> IO (Int, CharBuffer)
+writeBlocksRaw h buf0 n (Stream next0 s0 _len) = outer s0 n buf0
  where
-  outer s1 Buffer{bufRaw=raw, bufSize=len} = inner s1 (0::Int)
+  outer s1 n buf@Buffer{bufRaw=raw, bufSize=len} = inner s1 n
    where
     inner !s !n =
       case next0 s of
-        Done -> commit n False{-no flush-} True{-release-} >> return ()
+        Done -> return (n, buf)
         Skip s' -> inner s' n
         Yield x s'
-          | n + 1 >= len -> commit n True{-needs flush-} False >>= outer s
+          | n + 1 >= len -> commit n True{-needs flush-} False >>= outer s 0
           | otherwise    -> writeCharBuf raw n x >>= inner s'
     commit = commitBuffer h raw len
 
 -- This function is completely lifted from GHC.IO.Handle.Text.
+{-# INLINE getSpareBuffer #-}
 getSpareBuffer :: Handle__ -> IO (BufferMode, CharBuffer)
 getSpareBuffer Handle__{haCharBuffer=ref,
                         haBuffers=spare_ref,

src/Data/Text/Lazy/IO.hs

@@ -1,5 +1,7 @@
-{-# LANGUAGE BangPatterns, CPP, RecordWildCards #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE Trustworthy #-}
+{-# OPTIONS_GHC -Wno-name-shadowing #-}
 -- |
 -- Module      : Data.Text.Lazy.IO
 -- Copyright   : (c) 2009, 2010 Bryan O'Sullivan,
@@ -45,19 +47,23 @@ import System.IO (Handle, IOMode(..), hPutChar, openFile, stdin, stdout,
                   withFile)
 import qualified Data.Text.IO as T
 import qualified Data.Text.Lazy as L
+import qualified Data.Text.Internal.Lazy as L
 import qualified Control.Exception as E
 import Control.Monad (when)
 import Data.IORef (readIORef)
 import Data.Text.Internal.IO (hGetLineWith, readChunk)
 import Data.Text.Internal.Lazy (chunk, empty)
-import GHC.IO.Buffer (isEmptyBuffer)
+import GHC.IO.Buffer (isEmptyBuffer, Buffer(Buffer,bufRaw,bufSize))
 import GHC.IO.Exception (IOException(..), IOErrorType(..), ioException)
 import GHC.IO.Handle.Internals (augmentIOError, hClose_help,
                                 wantReadableHandle, withHandle)
 import GHC.IO.Handle.Types (Handle__(..), HandleType(..))
 import System.IO (BufferMode(..), hGetBuffering, hSetBuffering)
 import System.IO.Error (isEOFError)
 import System.IO.Unsafe (unsafeInterleaveIO)
+import Data.Text.Internal.Fusion (stream)
+import Data.Foldable (foldlM)
+import Data.Functor (void)
 
 -- | Read a file and return its contents as a string.  The file is
 -- read lazily, as with 'getContents'.
@@ -129,7 +135,11 @@ hGetLine = hGetLineWith L.fromChunks
 
 -- | Write a string to a handle.
 hPutStr :: Handle -> Text -> IO ()
-hPutStr h = mapM_ (T.hPutStr h) . L.toChunks
+hPutStr _ L.Empty = pure ()
+hPutStr h t = T.hPutStrInit h $ \mode buf nl -> do
+  (n, b) <- foldlM (\(n, buf) t -> T.hPutStr' h mode nl buf n (stream t)) (0, buf) (L.toChunks t)
+  let Buffer{bufRaw,bufSize} = b
+  when (n /= 0) $ void $ T.commitBuffer h bufRaw bufSize n False True
 
 -- | Write a string to a handle, followed by a newline.
 hPutStrLn :: Handle -> Text -> IO ()