Pack.chs

module Codec.Compression.BZip.Pack ( compress
                                   , compressWith
                                   ) where

import Codec.Compression.BZip.Foreign
import Codec.Compression.BZip.Common
import Control.Applicative
import Control.Arrow (second)
import Control.Exception (bracket)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Unsafe as BS
import Foreign.C.Types (CInt)
import Foreign.Ptr (Ptr, castPtr, nullPtr)
import Foreign.Marshal.Alloc (free, mallocBytes)
import System.IO.Unsafe (unsafeDupablePerformIO)

#include <bzlib.h>

{-# NOINLINE compress #-}
-- | @since 0.1.1.0
compress :: BSL.ByteString -> BSL.ByteString
compress = compressWith 9 30

{-# NOINLINE compressWith #-}
-- | See [bzlib manual](https://www.sourceware.org/bzip2/manual/manual.html#bzcompress-init)
-- for information on compression parameters.
--
-- @since 0.1.1.0
compressWith :: CInt -- ^ Block size (@1-9@)
             -> CInt -- ^ Work factor (@0-250@)
             -> BSL.ByteString
             -> BSL.ByteString
compressWith blkSize wf bsl = unsafeDupablePerformIO $
    let bss = BSL.toChunks bsl in
    BSL.fromChunks <$> bracket
        (do { p <- bzStreamInit ; bzCompressInit blkSize wf p ; pure p })
        bZ2BzCompressEnd
        (\p -> bzCompressChunks p bss)

bzCompressChunks :: Ptr BzStream -> [BS.ByteString] -> IO [BS.ByteString]
bzCompressChunks p bs =

    bracket
        (mallocBytes bufSz)
        free
        (fmap snd . extractBuf bs)

    where

        -- corresponds to inner loop in zlib example
        fillBuf :: [BS.ByteString] -> BZAction -> Ptr a -> IO (BZError, [BS.ByteString])
        fillBuf bs' f bufOut = do

            {# set bz_stream.avail_out #} p bufSz
            {# set bz_stream.next_out #} p (castPtr bufOut)

            ret <- bZ2BzCompress p f

            szOut <- fromIntegral <$> {# get bz_stream->avail_out #} p

            let bytesAvail = bufSz - szOut

            newBSAp <- if bytesAvail /= 0
                then (:) <$> BS.packCStringLen (castPtr bufOut, bytesAvail)
                else pure id

            if ret == BzStreamEnd
                then pure (ret, newBSAp [])
                else if szOut == 0
                    then second newBSAp <$> fillBuf bs' f bufOut
                    else second newBSAp <$> extractBuf bs' bufOut

        -- corresponds to outer loop in zlib example
        extractBuf :: [BS.ByteString] -> Ptr a -> IO (BZError, [BS.ByteString])
        extractBuf [] bufOut = do

            (res, blocks) <- fillBuf [] BzFinish bufOut

            if res == BzStreamEnd
                then pure (BzStreamEnd, blocks)
                else extractBuf [] bufOut

        extractBuf (b:bs') bufOut =
            BS.unsafeUseAsCStringLen b $ \(buf, sz) -> do

                {# set bz_stream.avail_in #} p (fromIntegral sz)
                {# set bz_stream.next_in #} p buf

                (res, blocks) <- fillBuf bs' BzRun bufOut

                if res == BzStreamEnd
                    then pure (BzStreamEnd, blocks)
                    else extractBuf bs' bufOut

        bufSz :: Integral a => a
        bufSz = 32 * 1024

bzCompressInit :: CInt -> CInt -> Ptr BzStream -> IO ()
bzCompressInit blkSize wf p = do

    {# set bz_stream.next_in #} p nullPtr
    {# set bz_stream.avail_in #} p 0

    bZ2BzCompressInit p blkSize 0 wf