1{-# LANGUAGE BangPatterns #-}
2
3-----------------------------------------------------------------------------
4--
5-- Fast write-buffered Handles
6--
7-- (c) The University of Glasgow 2005-2006
8--
9-- This is a simple abstraction over Handles that offers very fast write
10-- buffering, but without the thread safety that Handles provide.  It's used
11-- to save time in Pretty.printDoc.
12--
13-----------------------------------------------------------------------------
14
15module BufWrite (
16        BufHandle(..),
17        newBufHandle,
18        bPutChar,
19        bPutStr,
20        bPutFS,
21        bPutFZS,
22        bPutPtrString,
23        bPutReplicate,
24        bFlush,
25  ) where
26
27import GhcPrelude
28
29import FastString
30import FastMutInt
31
32import Control.Monad    ( when )
33import Data.ByteString (ByteString)
34import qualified Data.ByteString.Unsafe as BS
35import Data.Char        ( ord )
36import Foreign
37import Foreign.C.String
38import System.IO
39
40-- -----------------------------------------------------------------------------
41
42data BufHandle = BufHandle {-#UNPACK#-}!(Ptr Word8)
43                           {-#UNPACK#-}!FastMutInt
44                           Handle
45
46newBufHandle :: Handle -> IO BufHandle
47newBufHandle hdl = do
48  ptr <- mallocBytes buf_size
49  r <- newFastMutInt
50  writeFastMutInt r 0
51  return (BufHandle ptr r hdl)
52
53buf_size :: Int
54buf_size = 8192
55
56bPutChar :: BufHandle -> Char -> IO ()
57bPutChar b@(BufHandle buf r hdl) !c = do
58  i <- readFastMutInt r
59  if (i >= buf_size)
60        then do hPutBuf hdl buf buf_size
61                writeFastMutInt r 0
62                bPutChar b c
63        else do pokeElemOff buf i (fromIntegral (ord c) :: Word8)
64                writeFastMutInt r (i+1)
65
66bPutStr :: BufHandle -> String -> IO ()
67bPutStr (BufHandle buf r hdl) !str = do
68  i <- readFastMutInt r
69  loop str i
70  where loop "" !i = do writeFastMutInt r i; return ()
71        loop (c:cs) !i
72           | i >= buf_size = do
73                hPutBuf hdl buf buf_size
74                loop (c:cs) 0
75           | otherwise = do
76                pokeElemOff buf i (fromIntegral (ord c))
77                loop cs (i+1)
78
79bPutFS :: BufHandle -> FastString -> IO ()
80bPutFS b fs = bPutBS b $ bytesFS fs
81
82bPutFZS :: BufHandle -> FastZString -> IO ()
83bPutFZS b fs = bPutBS b $ fastZStringToByteString fs
84
85bPutBS :: BufHandle -> ByteString -> IO ()
86bPutBS b bs = BS.unsafeUseAsCStringLen bs $ bPutCStringLen b
87
88bPutCStringLen :: BufHandle -> CStringLen -> IO ()
89bPutCStringLen b@(BufHandle buf r hdl) cstr@(ptr, len) = do
90  i <- readFastMutInt r
91  if (i + len) >= buf_size
92        then do hPutBuf hdl buf i
93                writeFastMutInt r 0
94                if (len >= buf_size)
95                    then hPutBuf hdl ptr len
96                    else bPutCStringLen b cstr
97        else do
98                copyBytes (buf `plusPtr` i) ptr len
99                writeFastMutInt r (i + len)
100
101bPutPtrString :: BufHandle -> PtrString -> IO ()
102bPutPtrString b@(BufHandle buf r hdl) l@(PtrString a len) = l `seq` do
103  i <- readFastMutInt r
104  if (i+len) >= buf_size
105        then do hPutBuf hdl buf i
106                writeFastMutInt r 0
107                if (len >= buf_size)
108                    then hPutBuf hdl a len
109                    else bPutPtrString b l
110        else do
111                copyBytes (buf `plusPtr` i) a len
112                writeFastMutInt r (i+len)
113
114-- | Replicate an 8-bit character
115bPutReplicate :: BufHandle -> Int -> Char -> IO ()
116bPutReplicate (BufHandle buf r hdl) len c = do
117  i <- readFastMutInt r
118  let oc = fromIntegral (ord c)
119  if (i+len) < buf_size
120    then do
121      fillBytes (buf `plusPtr` i) oc len
122      writeFastMutInt r (i+len)
123    else do
124      -- flush the current buffer
125      when (i /= 0) $ hPutBuf hdl buf i
126      if (len < buf_size)
127        then do
128          fillBytes buf oc len
129          writeFastMutInt r len
130        else do
131          -- fill a full buffer
132          fillBytes buf oc buf_size
133          -- flush it as many times as necessary
134          let go n | n >= buf_size = do
135                                       hPutBuf hdl buf buf_size
136                                       go (n-buf_size)
137                   | otherwise     = writeFastMutInt r n
138          go len
139
140bFlush :: BufHandle -> IO ()
141bFlush (BufHandle buf r hdl) = do
142  i <- readFastMutInt r
143  when (i > 0) $ hPutBuf hdl buf i
144  free buf
145  return ()
146