content: -- |
  -- Module      : Data.Git.Storage.FileReader
  -- License     : BSD-style
  -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
  -- Stability   : experimental
  -- Portability : unix
  --
  {-# LANGUAGE DeriveDataTypeable #-}
  module Data.Git.Storage.FileReader
          ( FileReader
          , fileReaderNew
          , fileReaderClose
          , withFileReader
          , withFileReaderDecompress
          , fileReaderGetPos
          , fileReaderGet
          , fileReaderGetLBS
          , fileReaderGetBS
          , fileReaderGetRef
          , fileReaderGetVLF
          , fileReaderSeek
          , fileReaderParse
          , fileReaderInflateToSize
          ) where
  
  
  import Control.Exception (bracket, throwIO)
  
  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.Git.Imports
  import Data.Git.OS
  import Data.Git.Ref
  import qualified Data.Git.Parser as P
  
  import Data.Data
  import Data.Word
  
  import Codec.Zlib
  import Codec.Zlib.Lowlevel
  import Crypto.Hash
  import Foreign.ForeignPtr
  import qualified Control.Exception as E
  
  import System.IO (hSeek, SeekMode(..))
  
  data FileReader = FileReader
          { fbHandle     :: Handle
          , fbUseInflate :: Bool
          , fbInflate    :: Inflate
          , fbRemaining  :: IORef (Maybe 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 (Just B.empty)
          pos <- newIORef 0
          inflate <- initInflate defaultWindowBits
          return $ FileReader handle decompress inflate ref pos
  
  fileReaderClose :: FileReader -> IO ()
  fileReaderClose = hClose . fbHandle
  
  withFileReader :: LocalPath -> (FileReader -> IO a) -> IO a
  withFileReader path f =
          bracket (openFile path ReadMode) (hClose) $ \handle ->
                  bracket (fileReaderNew False handle) (\_ -> return ()) f
  
  withFileReaderDecompress :: LocalPath -> (FileReader -> IO a) -> IO a
  withFileReaderDecompress path f =
          bracket (openFile path ReadMode) (hClose) $ \handle ->
                  bracket (fileReaderNew True handle) (\_ -> return ()) f
  
  fileReaderGetNext :: FileReader -> IO (Maybe 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 $ nothingOnNull 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
      nothingOnNull b
          | B.null b  = Nothing
          | otherwise = Just b
  
  fileReaderGetPos :: FileReader -> IO Word64
  fileReaderGetPos fr = do
      storeLeft <- maybe 0 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 <- maybe B.empty id <$> readIORef ref
          if B.length b >= left
              then do
                  let (b1, b2) = B.splitAt left b
                  writeIORef ref (Just 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
  
  fileReaderGetRef :: HashAlgorithm hash => hash -> FileReader -> IO (Ref hash)
  fileReaderGetRef alg fr = fromBinary <$> fileReaderGetBS (hashDigestSize alg) fr
  
  -- | seek in a handle, and reset the remaining buffer to empty.
  fileReaderSeek :: FileReader -> Word64 -> IO ()
  fileReaderSeek (FileReader { fbHandle = handle, fbRemaining = ref, fbPos = pos }) absPos = do
          writeIORef ref (Just B.empty) >> writeIORef pos absPos >> hSeek handle AbsoluteSeek (fromIntegral absPos)
  
  -- | parse from a filebuffer
  fileReaderParse :: FileReader -> P.Parser a -> IO a
  fileReaderParse fr@(FileReader { fbRemaining = ref }) parseF = do
      initBS <- maybe B.empty id <$> readIORef ref
      result <- P.parseFeed (fileReaderGetNext fr) parseF initBS
      case result of
          P.ParseOK remaining a -> writeIORef ref (Just remaining) >> return a
          P.ParseMore _         -> error "parsing failed: partial with a handle, reached EOF ?"
          P.ParseFail err       -> error ("parsing failed: " ++ err)
  
  -- | get a Variable Length Field. get byte as long as MSB is set, and one byte after
  fileReaderGetVLF :: FileReader -> IO [Word8]
  fileReaderGetVLF fr = fileReaderParse fr P.vlf
  
  fileReaderInflateToSize :: FileReader -> Word64 -> IO L.ByteString
  fileReaderInflateToSize fb@(FileReader { fbRemaining = ref }) outputSize = do
          --pos <- fileReaderGetPos fb
          --putStrLn ("inflate to size " ++ show outputSize ++ " " ++ show pos)
          inflate <- inflateNew
          l       <- loop inflate outputSize
          --posend <- fileReaderGetPos fb
          --putStrLn ("inflated input " ++ show posend)
          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 (maybe B.empty id rbs) (maybe B.empty id <$> fileReaderGetNext fb)
                                     `E.catch` augmentAndRaise left
                  writeIORef ref (Just 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)
  
  -- lowlevel helpers to inflate only to a specific size.
  inflateNew :: IO (ForeignPtr ZStreamStruct)
  inflateNew = do
          zstr <- zstreamNew
          inflateInit2 zstr defaultWindowBits
          newForeignPtr c_free_z_stream_inflate zstr
  
  inflateToSize :: ForeignPtr ZStreamStruct -> Int -> Bool -> ByteString -> IO ByteString -> IO (ByteString, ByteString)
  inflateToSize inflate sz isLastBlock ibs nextBs = withForeignPtr inflate $ \zstr -> do
          let boundSz = min defaultChunkSize sz
          -- create an output buffer
          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
                                          --when (ai /= 0) $ error ("input not consumed completly: ai" ++ show ai)
                                          (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
  