Skip to content

Commit

Permalink
Make putStrLn more atomic with line or block buffering
Browse files Browse the repository at this point in the history
  • Loading branch information
Lysxia authored and Bodigrim committed Jun 2, 2024
1 parent 4fba353 commit 0bcaa72
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 7 deletions.
10 changes: 6 additions & 4 deletions src/Data/Text/IO.hs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
32 changes: 31 additions & 1 deletion src/Data/Text/Internal/Fusion.hs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ module Data.Text.Internal.Fusion

-- * Creation and elimination
, stream
, streamLn
, unstream
, reverseStream

Expand All @@ -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)
Expand Down Expand Up @@ -98,6 +99,35 @@ 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\')@
--
-- @since 2.1.2
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.
--
Expand Down
6 changes: 4 additions & 2 deletions src/Data/Text/Lazy/IO.hs
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ 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 Data.Text.Internal.Lazy (Text(..), chunk, empty)
import GHC.IO.Buffer (isEmptyBuffer)
import GHC.IO.Exception (IOException(..), IOErrorType(..), ioException)
import GHC.IO.Handle.Internals (augmentIOError, hClose_help,
Expand Down Expand Up @@ -133,7 +133,9 @@ hPutStr h = mapM_ (T.hPutStr h) . L.toChunks

-- | 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 Empty = hPutChar h '\n'
hPutStrLn h (Chunk t Empty) = T.hPutStrLn h t -- print the newline after the last chunk atomically
hPutStrLn h (Chunk t ts) = T.hPutStr h t >> hPutStrLn h ts

-- | The 'interact' function takes a function of type @Text -> Text@
-- as its argument. The entire input from the standard input device is
Expand Down

0 comments on commit 0bcaa72

Please sign in to comment.