{-# LANGUAGE MagicHash #-}

-----------------------------------------------------------------------------
--
-- Fast write-buffered Handles
--
-- (c) The University of Glasgow 2005-2006
--
-- This is a simple abstraction over Handles that offers very fast write
-- buffering, but without the thread safety that Handles provide.  It's used
-- to save time in GHC.Utils.Ppr.printDoc.
--
-----------------------------------------------------------------------------

module GHC.Utils.BufHandle (
        BufHandle(..),
        newBufHandle,
        bPutChar,
        bPutStr,
        bPutFS,
        bPutFZS,
        bPutPtrString,
        bPutReplicate,
        bFlush,
  ) where

import GHC.Prelude.Basic

import GHC.Data.FastString
import GHC.Data.FastMutInt

import Control.Monad    ( when )
import Data.ByteString (ByteString)
import qualified Data.ByteString.Unsafe as BS
import Data.Char        ( ord )
import Foreign
import Foreign.C.String
import System.IO

-- for RULES
import GHC.Exts (unpackCString#, unpackNBytes#, Int(..))
import GHC.Ptr (Ptr(..))

-- -----------------------------------------------------------------------------

data BufHandle = BufHandle {-#UNPACK#-}!(Ptr Word8)
                           {-#UNPACK#-}!FastMutInt
                           Handle

newBufHandle :: Handle -> IO BufHandle
newBufHandle :: Handle -> IO BufHandle
newBufHandle Handle
hdl = do
  ptr <- Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
mallocBytes Int
buf_size
  r <- newFastMutInt 0
  return (BufHandle ptr r hdl)

buf_size :: Int
buf_size :: Int
buf_size = Int
8192

bPutChar :: BufHandle -> Char -> IO ()
bPutChar :: BufHandle -> Char -> IO ()
bPutChar b :: BufHandle
b@(BufHandle Ptr Word8
buf FastMutInt
r Handle
hdl) !Char
c = do
  i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  if (i >= buf_size)
        then do hPutBuf hdl buf buf_size
                writeFastMutInt r 0
                bPutChar b c
        else do pokeElemOff buf i (fromIntegral (ord c) :: Word8)
                writeFastMutInt r (i+1)

-- Equivalent of the text/str, text/unpackNBytes#, text/[] rules
-- in GHC.Utils.Ppr.
{-# RULES "hdoc/str"
    forall a h. bPutStr h (unpackCString# a) = bPutPtrString h (mkPtrString# a)
  #-}
{-# RULES "hdoc/unpackNBytes#"
    forall p n h. bPutStr h (unpackNBytes# p n) = bPutPtrString h (PtrString (Ptr p) (I# n))
  #-}
{-# RULES "hdoc/[]#"
    forall h. bPutStr h [] = return ()
  #-}

{-# NOINLINE [0] bPutStr #-}  -- Give the RULE a chance to fire
                              -- It must wait till after phase 1 when
                              -- the unpackCString first is manifested

bPutStr :: BufHandle -> String -> IO ()
bPutStr :: BufHandle -> String -> IO ()
bPutStr (BufHandle Ptr Word8
buf FastMutInt
r Handle
hdl) !String
str = do
  i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  loop str i
  where loop :: String -> Int -> IO ()
loop String
"" !Int
i = do FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r Int
i; () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        loop (Char
c:String
cs) !Int
i
           | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
buf_size = do
                Handle -> Ptr Word8 -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr Word8
buf Int
buf_size
                String -> Int -> IO ()
loop (Char
cChar -> String -> String
forall a. a -> [a] -> [a]
:String
cs) Int
0
           | Bool
otherwise = do
                Ptr Word8 -> Int -> Word8 -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr Word8
buf Int
i (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
ord Char
c))
                String -> Int -> IO ()
loop String
cs (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

bPutFS :: BufHandle -> FastString -> IO ()
bPutFS :: BufHandle -> FastString -> IO ()
bPutFS BufHandle
b FastString
fs = BufHandle -> ByteString -> IO ()
bPutBS BufHandle
b (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ FastString -> ByteString
bytesFS FastString
fs

bPutFZS :: BufHandle -> FastZString -> IO ()
bPutFZS :: BufHandle -> FastZString -> IO ()
bPutFZS BufHandle
b FastZString
fs = BufHandle -> ByteString -> IO ()
bPutBS BufHandle
b (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ FastZString -> ByteString
fastZStringToByteString FastZString
fs

bPutBS :: BufHandle -> ByteString -> IO ()
bPutBS :: BufHandle -> ByteString -> IO ()
bPutBS BufHandle
b ByteString
bs = ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BS.unsafeUseAsCStringLen ByteString
bs ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ BufHandle -> CStringLen -> IO ()
bPutCStringLen BufHandle
b

bPutCStringLen :: BufHandle -> CStringLen -> IO ()
bPutCStringLen :: BufHandle -> CStringLen -> IO ()
bPutCStringLen b :: BufHandle
b@(BufHandle Ptr Word8
buf FastMutInt
r Handle
hdl) cstr :: CStringLen
cstr@(Ptr CChar
ptr, Int
len) = do
  i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  if (i + len) >= buf_size
        then do hPutBuf hdl buf i
                writeFastMutInt r 0
                if (len >= buf_size)
                    then hPutBuf hdl ptr len
                    else bPutCStringLen b cstr
        else do
                copyBytes (buf `plusPtr` i) ptr len
                writeFastMutInt r (i + len)

bPutPtrString :: BufHandle -> PtrString -> IO ()
bPutPtrString :: BufHandle -> PtrString -> IO ()
bPutPtrString b :: BufHandle
b@(BufHandle Ptr Word8
buf FastMutInt
r Handle
hdl) l :: PtrString
l@(PtrString Ptr Word8
a Int
len) = PtrString
l PtrString -> IO () -> IO ()
forall a b. a -> b -> b
`seq` do
  i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  if (i+len) >= buf_size
        then do hPutBuf hdl buf i
                writeFastMutInt r 0
                if (len >= buf_size)
                    then hPutBuf hdl a len
                    else bPutPtrString b l
        else do
                copyBytes (buf `plusPtr` i) a len
                writeFastMutInt r (i+len)

-- | Replicate an 8-bit character
bPutReplicate :: BufHandle -> Int -> Char -> IO ()
bPutReplicate :: BufHandle -> Int -> Char -> IO ()
bPutReplicate (BufHandle Ptr Word8
buf FastMutInt
r Handle
hdl) Int
len Char
c = do
  i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  let oc = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Char -> Int
ord Char
c)
  if (i+len) < buf_size
    then do
      fillBytes (buf `plusPtr` i) oc len
      writeFastMutInt r (i+len)
    else do
      -- flush the current buffer
      when (i /= 0) $ hPutBuf hdl buf i
      if (len < buf_size)
        then do
          fillBytes buf oc len
          writeFastMutInt r len
        else do
          -- fill a full buffer
          fillBytes buf oc buf_size
          -- flush it as many times as necessary
          let go Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
buf_size = do
                                       Handle -> Ptr Word8 -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr Word8
buf Int
buf_size
                                       Int -> IO ()
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
buf_size)
                   | Bool
otherwise     = FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r Int
n
          go len

bFlush :: BufHandle -> IO ()
bFlush :: BufHandle -> IO ()
bFlush (BufHandle Ptr Word8
buf FastMutInt
r Handle
hdl) = do
  i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  when (i > 0) $ hPutBuf hdl buf i
  free buf
  return ()