Bz3.hs

module Codec.Bz3 ( Bz3Error (..)
                 , decompressFile
                 , compressFile
                 , compressFileDefault
                 ) where

import           Codec.Bz3.Binary
import           Codec.Bz3.Foreign
import qualified Control.Monad.ST.Lazy        as LazyST
import qualified Control.Monad.ST.Lazy.Unsafe as LazyST
import           Data.Bifunctor               (bimap)
import           Data.Binary.Get              (runGetOrFail)
import           Data.Binary.Put              (runPut)
import qualified Data.ByteString              as BS
import           Data.ByteString.Internal     as BS
import qualified Data.ByteString.Lazy         as BSL
import qualified Data.ByteString.Unsafe       as BS
import           Data.Int                     (Int32)
import           Foreign.ForeignPtr           (ForeignPtr, castForeignPtr,
                                               mallocForeignPtrBytes,
                                               newForeignPtr, withForeignPtr)
import           Foreign.Marshal.Array        (copyArray)
import           Foreign.Ptr                  (castPtr)

newBz3StForeign :: Int32 -> IO (ForeignPtr Bz3St)
newBz3StForeign oSz = castForeignPtr <$> (newForeignPtr bz3Free . castPtr =<< bz3New oSz)

decompressFile :: BSL.ByteString -> BSL.ByteString
decompressFile contents = LazyST.runST $ BSL.fromChunks <$> do
    let Right (bs, _, bSz) = runGetOrFail getFileH contents
    st <- LazyST.unsafeIOToST $ newBz3StForeign (fromIntegral bSz)
    loop st bs
  where
    loop st bs | BSL.null bs = pure []
               | otherwise = do {(dc, rest) <- decN st bs; (dc:) <$> loop st rest}

-- | @since 1.0.0.0
compressFileDefault :: BSL.ByteString -> BSL.ByteString
compressFileDefault = compressFile (16*1024*1024)

-- will fail if bz3st cannot fit exotic big input
compressFile :: Int32 -- ^ Block size, 65kB to 511MB
             -> BSL.ByteString -> BSL.ByteString
compressFile ibSz d = LazyST.runST $ do
    st <- LazyST.unsafeIOToST $ newBz3StForeign bSz
    (sz,bb) <- loop st (BSL.toChunks d)
    let fileH = runPut (putFileH sz)
    pure (fileH<>BSL.fromChunks bb)
  where
    loop _ []     = pure (bSz, [])
    loop s (b:bs) | BS.length b > feed = let (next,b') = BS.splitAt feed b in loop s (next:b':bs)
    loop s (b:bs) = do
        (csz, e) <- encN s b
        let ucsz=fromIntegral csz; chunk=BSL.toStrict $ runPut (putChunk (Chunk ucsz osz))
        bimap (max ucsz) ((chunk:).(e:)) <$> loop s bs
      where
        osz=fromIntegral (BS.length b)

    feed=fromIntegral (bz3Bound bSz)

    bSz :: Integral a => a
    bSz=fromIntegral ibSz

encN :: ForeignPtr Bz3St -> BS.ByteString -> LazyST.ST s (Int32, BS.ByteString)
encN st inp = LazyST.unsafeIOToST $
    BS.unsafeUseAsCStringLen inp $ \(d,sz) -> do
        let bufSz = bz3Bound (fromIntegral sz)
        buf <- mallocForeignPtrBytes (fromIntegral bufSz)
        enc <- withForeignPtr buf $ \bb -> do
            copyArray bb d sz
            bz3EncodeBlock st (castPtr bb) (fromIntegral sz)
        pure (enc, BS.BS (castForeignPtr buf) (fromIntegral enc))

decN :: ForeignPtr Bz3St -> BSL.ByteString -> LazyST.ST s (BS.ByteString, BSL.ByteString)
decN st inp = LazyST.unsafeIOToST $ do
    let (next, off, Chunk csz osz) = y (runGetOrFail getChunk inp)
        bb=BSL.toStrict (BSL.drop 8 $ BSL.take off inp)
    if osz>=64
        then do
            let bufSz = bz3Bound (fromIntegral osz)
                csz32 = fromIntegral csz; osz32 = fromIntegral osz
            buf <- mallocForeignPtrBytes (fromIntegral bufSz)
            res <- withForeignPtr buf $ \b -> do
                BS.unsafeUseAsCStringLen bb $ \(p,isz) -> copyArray b (castPtr p) isz
                bz3DecodeBlock st b bufSz csz32 osz32
            case res of
                Right{} -> pure ()
                Left e | e==Bz3Ok -> pure ()
                       | otherwise -> error =<< bz3Strerror st
            pure (BS.BS (castForeignPtr buf) (fromIntegral osz), next)
        else pure (bb, next)
  where
    y (Right x) = x; y (Left (_, _, e)) = error e