{-# 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,
        bPutLitString,
        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 hdl = do
  ptr <- mallocBytes buf_size
  r <- newFastMutInt
  writeFastMutInt r 0
  return (BufHandle ptr r hdl)

buf_size :: Int
buf_size = 8192

bPutChar :: BufHandle -> Char -> IO ()
bPutChar b@(BufHandle buf r hdl) !c = do
  i <- readFastMutInt 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)

bPutStr :: BufHandle -> String -> IO ()
bPutStr (BufHandle buf r hdl) !str = do
  i <- readFastMutInt r
  loop str i
  where loop "" !i = do writeFastMutInt r i; return ()
        loop (c:cs) !i
           | i >= buf_size = do
                hPutBuf hdl buf buf_size
                loop (c:cs) 0
           | otherwise = do
                pokeElemOff buf i (fromIntegral (ord c))
                loop cs (i+1)

bPutFS :: BufHandle -> FastString -> IO ()
bPutFS b fs = bPutBS b $ fastStringToByteString fs

bPutFZS :: BufHandle -> FastZString -> IO ()
bPutFZS b fs = bPutBS b $ fastZStringToByteString fs

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

bPutCStringLen :: BufHandle -> CStringLen -> IO ()
bPutCStringLen b@(BufHandle buf r hdl) cstr@(ptr, len) = do
  i <- readFastMutInt 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)

bPutLitString :: BufHandle -> LitString -> IO ()
bPutLitString b@(BufHandle buf r hdl) l@(LitString a len) = l `seq` do
  i <- readFastMutInt 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 bPutLitString 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 buf r hdl) len c = do
  i <- readFastMutInt r
  let oc = fromIntegral (ord 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 n | n >= buf_size = do
                                       hPutBuf hdl buf buf_size
                                       go (n-buf_size)
                   | otherwise     = writeFastMutInt r n
          go len

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