    {-# LANGUAGE ScopedTypeVariables, RecordWildCards #-}

    -- I tried switching to blaze-builder, but this buffer is massively faster
    module Hoogle.Store.WriteBuffer(
        Buffer, withBuffer,
        putStorable, putByteString,
        patch, getPos
        ) where

    import General.Base
    import General.System
    import Data.IORef
    import Foreign
    import qualified Data.ByteString as BS
    import qualified Data.ByteString.Internal as BS
    import General.Util

    bufferSize = 10000 :: Word32

    -- (number in file, number in buffer)
    data Buffer = Buffer
        {handle :: Handle -- the handle we are writing to
        ,ptr :: Ptr () -- the current buffer
        ,inFile :: IORef Word32 -- the number of bytes on the disk
        ,inBuffer :: IORef Word32 -- the number of bytes in the buffer
        ,patchup :: IORef [Patchup]
        }

    data Patchup = !Word32 := !Word32

    writeRef ref v = v `seq` writeIORef ref v
    modifyRef ref f = writeRef ref . f =<< readIORef ref

    withBuffer :: Handle -> (Buffer -> IO a) -> IO a
    withBuffer handle f = do
        inFile <- newIORef . fromInteger =<< hTell handle
        inBuffer <- newIORef 0
        patchup <- newIORef []
        allocaBytes (fromIntegral bufferSize) $ \ptr -> do
            res <- f $ Buffer handle ptr inFile inBuffer patchup
            inBuf <- readIORef inBuffer
            when (inBuf > 0) $ hPutBuf handle ptr (fromIntegral inBuf)
            xs <- fmap (sortOn $ \(a := b) -> a) $ readIORef patchup
            forM_ xs $ \(pos := val) -> do
                hSeek handle AbsoluteSeek $ toInteger pos
                poke (castPtr ptr) val
                hPutBuf handle ptr $ sizeOf val
            return res

    put :: Buffer -> Word32 -> (Handle -> IO ()) -> (Ptr a -> Int -> IO ()) -> IO ()
    put _ 0 _ _ = return ()
    put Buffer{..} sz toFile toBuffer = do
        inBuf <- readIORef inBuffer
        if inBuf + sz >= bufferSize then do
            when (inBuf > 0) $ hPutBuf handle ptr $ fromIntegral inBuf
            if sz >= bufferSize `div` 2 then do
                toFile handle
                modifyRef inFile (+ (inBuf+sz))
                writeRef inBuffer 0
             else do
                toBuffer (castPtr ptr) 0
                modifyRef inFile (+inBuf)
                writeRef inBuffer sz
         else do
            toBuffer (castPtr ptr) $ fromIntegral inBuf
            writeIORef inBuffer (inBuf+sz)

    putStorable :: Storable a => Buffer -> a -> IO ()
    putStorable buf x = put buf (fromIntegral sz)
        (\h -> allocaBytes (sizeOf x) $ \ptr -> poke ptr x >> hPutBuf h ptr sz)
        (\ptr pos -> pokeByteOff ptr pos x)
        where sz = sizeOf x

    putByteString :: Buffer -> BS.ByteString -> IO ()
    putByteString buf x = put buf (fromIntegral $ BS.length x) (`BS.hPut` x) $
        \ptr pos -> let (fp,offset,len) = BS.toForeignPtr x in
                    withForeignPtr fp $ \p -> BS.memcpy (plusPtr ptr pos) (plusPtr p offset) (fromIntegral len)

    getPos :: Buffer -> IO Word32
    getPos Buffer{..} = liftM2 (+) (readIORef inFile) (readIORef inBuffer)

    -- Patch at position p, with value v. p must be in the past.
    -- Return True if you succeeded, False if that is already on disk
    patch :: Buffer -> Word32 -> Word32 -> IO ()
    patch Buffer{..} p v = do
        i <- readIORef inFile
        if p >= i then
            pokeByteOff ptr (fromIntegral $ p-i) v
         else
            modifyRef patchup $ (:) (p := v)