{-# LANGUAGE BangPatterns #-}

-----------------------------------------------------------------------------
--
-- 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 Pretty.printDoc.
--
-----------------------------------------------------------------------------

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

import GhcPrelude

import FastString
import 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

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

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

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

buf_size :: Int
buf_size :: Int
buf_size = 8192

bPutChar :: BufHandle -> Char -> IO ()
bPutChar :: BufHandle -> Char -> IO ()
bPutChar b :: BufHandle
b@(BufHandle buf :: Ptr Word8
buf r :: FastMutInt
r hdl :: Handle
hdl) !Char
c = do
  Int
i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  if (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
buf_size)
        then do Handle -> Ptr Word8 -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr Word8
buf Int
buf_size
                FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r 0
                BufHandle -> Char -> IO ()
bPutChar BufHandle
b Char
c
        else 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) :: Word8)
                FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)

bPutStr :: BufHandle -> String -> IO ()
bPutStr :: BufHandle -> String -> IO ()
bPutStr (BufHandle buf :: Ptr Word8
buf r :: FastMutInt
r hdl :: Handle
hdl) !String
str = do
  Int
i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  String -> Int -> IO ()
loop String
str Int
i
  where loop :: String -> Int -> IO ()
loop "" !Int
i = do FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r Int
i; () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        loop (c :: Char
c:cs :: 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) 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
+1)

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

bPutFZS :: BufHandle -> FastZString -> IO ()
bPutFZS :: BufHandle -> FastZString -> IO ()
bPutFZS b :: BufHandle
b fs :: 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 b :: BufHandle
b bs :: 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 buf :: Ptr Word8
buf r :: FastMutInt
r hdl :: Handle
hdl) cstr :: CStringLen
cstr@(ptr :: Ptr CChar
ptr, len :: Int
len) = do
  Int
i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  if (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
buf_size
        then do Handle -> Ptr Word8 -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr Word8
buf Int
i
                FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r 0
                if (Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
buf_size)
                    then Handle -> Ptr CChar -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr CChar
ptr Int
len
                    else BufHandle -> CStringLen -> IO ()
bPutCStringLen BufHandle
b CStringLen
cstr
        else do
                Ptr CChar -> Ptr CChar -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes (Ptr Word8
buf Ptr Word8 -> Int -> Ptr CChar
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
i) Ptr CChar
ptr Int
len
                FastMutInt -> Int -> IO ()
writeFastMutInt FastMutInt
r (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len)

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

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

bFlush :: BufHandle -> IO ()
bFlush :: BufHandle -> IO ()
bFlush (BufHandle buf :: Ptr Word8
buf r :: FastMutInt
r hdl :: Handle
hdl) = do
  Int
i <- FastMutInt -> IO Int
readFastMutInt FastMutInt
r
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> Ptr Word8 -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr Word8
buf Int
i
  Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
free Ptr Word8
buf
  () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()