Bz3.hs

module Codec.Bz3 ( Bz3Error (..)
                 , decompressFile
                 , compressFile
                 ) 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) <- decNext st bs; (dc:) <$> loop st rest}

-- will fail if bz3st cannot fit exotic big input
compressFile :: BSL.ByteString -> BSL.ByteString
compressFile 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 (0, [])
    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) <- if osz<64
            then pure (0xffffffff, b)
            else encN s b
        let chunk=BSL.toStrict $ runPut (putChunk (Chunk (fromIntegral csz) osz))
        bimap (max osz) ((chunk:).(e:)) <$> loop s bs
      where
        osz=fromIntegral (BS.length b)

    feed=fromIntegral (bz3Bound bSz)

    bSz :: Integral a => a
    bSz=16*1024*1024

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))

decNext :: ForeignPtr Bz3St -> BSL.ByteString -> LazyST.ST s (BS.ByteString, BSL.ByteString)
decNext st inp = LazyST.unsafeIOToST $ do
    let (next, off, Chunk csz osz) = y (runGetOrFail getChunk inp)
    if osz>=64
        then do
            let bufSz = bz3Bound (fromIntegral osz)
                csz32 = fromIntegral csz; osz32 = fromIntegral osz
            buf <- mallocForeignPtrBytes (fromIntegral bufSz)
            let bb=BSL.toStrict (BSL.drop 8 $ BSL.take off inp)
            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 let bb=BSL.toStrict (BSL.drop 8 $ BSL.take off inp) 
             in pure (bb, next)
  where
    y (Right x) = x; y (Left (_, _, e)) = error e