module Data.Git.Storage.FileReader
( FileReader
, fileReaderNew
, fileReaderClose
, withFileReader
, withFileReaderDecompress
, fileReaderGetPos
, fileReaderGet
, fileReaderGetLBS
, fileReaderGetBS
, fileReaderGetVLF
, fileReaderSeek
, fileReaderParse
, fileReaderInflateToSize
) where
import Control.Applicative ((<$>))
import Control.Exception (bracket, throwIO)
import Control.Monad
import Data.Attoparsec (parseWith, Parser, IResult(..))
import qualified Data.Attoparsec as A
import Data.Bits
import Data.ByteString (ByteString)
import Data.ByteString.Unsafe
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import Data.ByteString.Lazy.Internal (defaultChunkSize)
import Data.IORef
import Data.Data
import Data.Word
import Codec.Zlib
import Codec.Zlib.Lowlevel
import Foreign.ForeignPtr
import qualified Control.Exception as E
import System.IO (hClose, hSeek, SeekMode(..))
import Filesystem
import Filesystem.Path
import Prelude hiding (FilePath)
data FileReader = FileReader
{ fbHandle :: Handle
, fbUseInflate :: Bool
, fbInflate :: Inflate
, fbRemaining :: IORef ByteString
, fbPos :: IORef Word64
}
data InflateException = InflateException Word64 Word64 String
deriving (Show,Eq,Typeable)
instance E.Exception InflateException
fileReaderNew :: Bool -> Handle -> IO FileReader
fileReaderNew decompress handle = do
ref <- newIORef B.empty
pos <- newIORef 0
inflate <- initInflate defaultWindowBits
return $ FileReader handle decompress inflate ref pos
fileReaderClose :: FileReader -> IO ()
fileReaderClose = hClose . fbHandle
withFileReader :: FilePath -> (FileReader -> IO a) -> IO a
withFileReader path f =
bracket (openFile path ReadMode) (hClose) $ \handle ->
bracket (fileReaderNew False handle) (\_ -> return ()) f
withFileReaderDecompress :: FilePath -> (FileReader -> IO a) -> IO a
withFileReaderDecompress path f =
bracket (openFile path ReadMode) (hClose) $ \handle ->
bracket (fileReaderNew True handle) (\_ -> return ()) f
fileReaderGetNext :: FileReader -> IO ByteString
fileReaderGetNext fb = do
bs <- if fbUseInflate fb then inflateTillPop else B.hGet (fbHandle fb) 8192
modifyIORef (fbPos fb) (\pos -> pos + (fromIntegral $ B.length bs))
return bs
where inflateTillPop = do
b <- B.hGet (fbHandle fb) 4096
if B.null b
then finishInflate (fbInflate fb)
else (>>= maybe inflateTillPop return) =<< feedInflate (fbInflate fb) b
fileReaderGetPos :: FileReader -> IO Word64
fileReaderGetPos fr = do
storeLeft <- B.length <$> readIORef (fbRemaining fr)
pos <- readIORef (fbPos fr)
return (pos fromIntegral storeLeft)
fileReaderFill :: FileReader -> IO ()
fileReaderFill fb = fileReaderGetNext fb >>= writeIORef (fbRemaining fb)
fileReaderGet :: Int -> FileReader -> IO [ByteString]
fileReaderGet size fb@(FileReader { fbRemaining = ref }) = loop size
where loop left = do
b <- readIORef ref
if B.length b >= left
then do
let (b1, b2) = B.splitAt left b
writeIORef ref b2
return [b1]
else do
let nleft = left B.length b
fileReaderFill fb
liftM (b :) (loop nleft)
fileReaderGetLBS :: Int -> FileReader -> IO L.ByteString
fileReaderGetLBS size fb = L.fromChunks <$> fileReaderGet size fb
fileReaderGetBS :: Int -> FileReader -> IO ByteString
fileReaderGetBS size fb = B.concat <$> fileReaderGet size fb
fileReaderSeek :: FileReader -> Word64 -> IO ()
fileReaderSeek (FileReader { fbHandle = handle, fbRemaining = ref, fbPos = pos }) absPos = do
writeIORef ref B.empty >> writeIORef pos absPos >> hSeek handle AbsoluteSeek (fromIntegral absPos)
fileReaderParse :: FileReader -> Parser a -> IO a
fileReaderParse fr@(FileReader { fbRemaining = ref }) parseF = do
initBS <- readIORef ref
result <- parseWith (fileReaderGetNext fr) parseF initBS
case result of
Done remaining a -> writeIORef ref remaining >> return a
Partial _ -> error "parsing failed: partial with a handle, reached EOF ?"
Fail _ ctxs err -> error ("parsing failed: " ++ show ctxs ++ " : " ++ show err)
fileReaderGetVLF :: FileReader -> IO [Word8]
fileReaderGetVLF fr = fileReaderParse fr $ do
bs <- A.takeWhile (\w -> w `testBit` 7)
l <- A.anyWord8
return $ (map (\w -> w `clearBit` 7) $ B.unpack bs) ++ [l]
fileReaderInflateToSize :: FileReader -> Word64 -> IO L.ByteString
fileReaderInflateToSize fb@(FileReader { fbRemaining = ref }) outputSize = do
inflate <- inflateNew
l <- loop inflate outputSize
return $ L.fromChunks l
where loop inflate left = do
rbs <- readIORef ref
let maxToInflate = min left (16 * 1024)
let lastBlock = if left == maxToInflate then True else False
(dbs,remaining) <- inflateToSize inflate (fromIntegral maxToInflate) lastBlock rbs (fileReaderGetNext fb)
`E.catch` augmentAndRaise left
writeIORef ref remaining
let nleft = left fromIntegral (B.length dbs)
if nleft > 0
then liftM (dbs:) (loop inflate nleft)
else return [dbs]
augmentAndRaise :: Word64 -> E.SomeException -> IO a
augmentAndRaise left exn = throwIO $ InflateException outputSize left (show exn)
inflateNew = do
zstr <- zstreamNew
inflateInit2 zstr defaultWindowBits
newForeignPtr c_free_z_stream_inflate zstr
inflateToSize inflate sz isLastBlock ibs nextBs = withForeignPtr inflate $ \zstr -> do
let boundSz = min defaultChunkSize sz
fbuff <- mallocForeignPtrBytes boundSz
withForeignPtr fbuff $ \buff -> do
c_set_avail_out zstr buff (fromIntegral boundSz)
rbs <- loop zstr ibs
bs <- B.packCStringLen (buff, boundSz)
return (bs, rbs)
where
loop zstr nbs = do
(ai, streamEnd) <- inflateOneInput zstr nbs
ao <- c_get_avail_out zstr
if (isLastBlock && streamEnd) || (not isLastBlock && ao == 0)
then return $ bsTakeLast ai nbs
else do
(if ai == 0
then nextBs
else return (bsTakeLast ai nbs)) >>= loop zstr
inflateOneInput zstr bs = unsafeUseAsCStringLen bs $ \(istr, ilen) -> do
c_set_avail_in zstr istr $ fromIntegral ilen
r <- c_call_inflate_noflush zstr
when (r < 0 && r /= (5)) $ do
throwIO $ ZlibException $ fromIntegral r
ai <- c_get_avail_in zstr
return (ai, r == 1)
bsTakeLast len bs = B.drop (B.length bs fromIntegral len) bs