Lz4.chs

{-# LANGUAGE BangPatterns #-}

-- | The functions in this module throw exceptions on error.
--
-- 'decompress' and 'compress' are fully lazy, i.e. memory efficient.
module Codec.Lz4 ( -- * Functions for working with blocks
                   compressBlock
                 , decompressBlockSz
                 , lZ4MaxInputSize
                 , compressBlockHC
                 , lZ4HCClevelMax
                 -- * Functions for working with frames
                 , compress
                 , compressSz
                 , decompress
                 , decompressBufSz
                 -- * Version info
                 , lZ4VersionNumber
                 , lZ4VersionString
                 ) where

import           Codec.Lz4.Foreign
import           Control.Monad                (when)
import           Control.Monad.ST.Lazy        (runST)
import qualified Control.Monad.ST.Lazy        as LazyST
import qualified Control.Monad.ST.Lazy.Unsafe as LazyST
import qualified Data.ByteString              as BS
import qualified Data.ByteString.Internal     as BS
import qualified Data.ByteString.Lazy         as BSL
import qualified Data.ByteString.Unsafe       as BS
import           Foreign.C.String             (CString)
import           Foreign.C.Types              (CInt)
import           Foreign.Ptr                  (castPtr, nullPtr)
import           Foreign.ForeignPtr           (ForeignPtr, castForeignPtr,
                                               mallocForeignPtrBytes,
                                               newForeignPtr, withForeignPtr)
import           Foreign.Marshal.Alloc        (alloca)
import           Foreign.Storable             (peek, poke)
import           System.IO.Unsafe             (unsafePerformIO)

#include <lz4frame.h>

check :: LZ4FErrorCode -> IO ()
check err = when (lZ4FIsError err) $ error (lZ4FGetErrorName err)

-- | Lazily decompress a frame
decompress :: BSL.ByteString -> BSL.ByteString
decompress = decompressBufSz (64 * 1014)

-- | @since 0.1.3.0
decompressBufSz :: Int -- ^ Size of the output buffer
                -> BSL.ByteString
                -> BSL.ByteString
decompressBufSz bufSz bs = runST $ do

    let bss = BSL.toChunks bs

    (ctx, buf) <- LazyST.unsafeIOToST $ do
        (err, preCtx) <- lZ4FCreateDecompressionContext lZ4FGetVersion
        ctx <- castForeignPtr <$> newForeignPtr lZ4FFreeDecompressionContext (castPtr preCtx)
        check err
        dstBuf <- mallocForeignPtrBytes bufSz
        pure (ctx, dstBuf)

    BSL.fromChunks <$> loop ctx buf bss

    where loop :: LzDecompressionCtxPtr -> ForeignPtr a -> [BS.ByteString] -> LazyST.ST s [BS.ByteString]
          loop _ _ [] = pure []
          loop ctx buf (b:bs') = do
                (nxt, res) <- stepChunk ctx buf b
                case nxt of
                    Nothing   -> (res:) <$> loop ctx buf bs'
                    Just next -> (res:) <$> loop ctx buf (next:bs')

          stepChunk :: LzDecompressionCtxPtr -> ForeignPtr a -> BS.ByteString -> LazyST.ST s (Maybe BS.ByteString, BS.ByteString)
          stepChunk !ctx !dst b = LazyST.unsafeIOToST $
            BS.unsafeUseAsCStringLen b $ \(buf, sz) ->
                withForeignPtr dst $ \d ->
                    alloca $ \dSzPtr ->
                        alloca $ \szPtr -> do
                            poke dSzPtr (fromIntegral bufSz)
                            poke szPtr (fromIntegral sz)
                            res <- lZ4FDecompress ctx d dSzPtr buf szPtr nullPtr
                            check res
                            bRead <- peek szPtr
                            bWritten <- peek dSzPtr
                            outBs <- BS.packCStringLen (castPtr d, fromIntegral bWritten)
                            let remBs = if fromIntegral bRead == sz
                                then Nothing
                                else Just (BS.drop (fromIntegral bRead) b)
                            pure (remBs, outBs)

-- | Lazily compress a frame.
compress :: BSL.ByteString -> BSL.ByteString
compress = compressSz 0

-- | @since 0.1.4.0
compressSz :: Int -- ^ Compression level
           -> BSL.ByteString
           -> BSL.ByteString
compressSz lvl bs = runST $ do

    let bss = BSL.toChunks bs

    (ctx, pref, header) <- initCtx
    rest <- loop ctx pref bss

    pure $ BSL.fromChunks (header:rest)

    where initCtx :: LazyST.ST s (LzCtxPtr, LzPreferencesPtr, BS.ByteString)
          initCtx = LazyST.unsafeIOToST $ do
                (err, preCtx) <- lZ4FCreateCompressionContext lZ4FGetVersion
                ctx <- castForeignPtr <$> newForeignPtr lZ4FFreeCompressionContext (castPtr preCtx)
                check err
                dst <- mallocForeignPtrBytes lZ4FHeaderSizeMax
                pref <- mallocForeignPtrBytes {# sizeof LZ4F_preferences_t #}
                preferencesPtr pref lvl
                header <- withForeignPtr dst $ \d -> do
                    res <- lZ4FCompressBegin ctx d lZ4FHeaderSizeMax pref
                    check res
                    BS.packCStringLen (castPtr d, fromIntegral res)
                pure (ctx, pref, header)

          loop :: LzCtxPtr -> LzPreferencesPtr -> [BS.ByteString] -> LazyST.ST s [BS.ByteString]
          loop ctx pref []      = pure <$> finish ctx pref
          loop ctx pref (b:bs') = (:) <$> update ctx pref b <*> loop ctx pref bs'

          finish :: LzCtxPtr -> LzPreferencesPtr -> LazyST.ST s BS.ByteString
          finish ctx pref = LazyST.unsafeIOToST $ do
            let expectedSz = lZ4FCompressBound 0 pref
            dst <- mallocForeignPtrBytes (fromIntegral expectedSz)
            withForeignPtr dst $ \d -> do
                res <- lZ4FCompressEnd ctx d expectedSz nullPtr
                check res
                BS.packCStringLen (castPtr d, fromIntegral res)

          update :: LzCtxPtr -> LzPreferencesPtr -> BS.ByteString -> LazyST.ST s BS.ByteString
          update !ctx !pref b = LazyST.unsafeIOToST $
            BS.unsafeUseAsCStringLen b $ \(buf, sz) -> do
                let expectedSz = lZ4FCompressBound (fromIntegral sz) pref
                dst <- mallocForeignPtrBytes (fromIntegral expectedSz)
                withForeignPtr dst $ \d -> do
                    res <- lZ4FCompressUpdate ctx d expectedSz buf (fromIntegral sz) nullPtr
                    check res
                    BS.packCStringLen (castPtr d, fromIntegral res)

{-# NOINLINE compressBlock #-}
compressBlock :: BS.ByteString -> BS.ByteString
compressBlock = unsafePerformIO . compressBlockGeneric lZ4CompressDefault

{-# NOINLINE compressBlockHC #-}
-- | @since 0.1.1.0
compressBlockHC :: Int -- ^ Compression level (must be less than 'lZ4HCClevelMax')
                -> BS.ByteString
                -> BS.ByteString
compressBlockHC lvl = unsafePerformIO . compressBlockGeneric (\src dst ssz dsz -> lZ4CompressHC src dst ssz dsz (fromIntegral lvl))

compressBlockGeneric :: (CString -> CString -> CInt -> CInt -> IO CInt) -> BS.ByteString -> IO BS.ByteString
compressBlockGeneric cfun bs =
    BS.unsafeUseAsCStringLen bs $ \(buf, sz) -> do
        let resSz = lZ4CompressBound (fromIntegral sz)
        dst <- mallocForeignPtrBytes (fromIntegral resSz)
        withForeignPtr dst $ \d -> do
            bWritten <- cfun buf d (fromIntegral sz) resSz
            when (bWritten == 0) $ error "Compression error"
            pure $ BS.fromForeignPtr (castForeignPtr dst) 0 (fromIntegral bWritten)

{-# NOINLINE decompressBlockSz #-}
-- | Decompress a block. The size of the uncompressed data must be known.
decompressBlockSz :: BS.ByteString
                  -> Int -- ^ Decompressed size
                  -> BS.ByteString
decompressBlockSz bs expectedSz = unsafePerformIO $
    BS.unsafeUseAsCStringLen bs $ \(buf, sz) -> do
        dst <- mallocForeignPtrBytes expectedSz
        withForeignPtr dst $ \d -> do
            bWritten <- lZ4DecompressSafe buf d (fromIntegral sz) (fromIntegral expectedSz)
            when (bWritten < 0) $ error "Decompression error"
        pure $ BS.fromForeignPtr (castForeignPtr dst) 0 expectedSz

cint :: Enum a => a -> CInt
cint = fromIntegral . fromEnum

preferencesPtr :: LzPreferencesPtr -> Int -> IO ()
preferencesPtr fp i =
    withForeignPtr fp $ \p -> do

        {# set LZ4F_preferences_t.frameInfo.blockSizeID #} p (cint Lz4fDefault)
        {# set LZ4F_preferences_t.frameInfo.blockMode #} p (cint Lz4fBlocklinked)
        {# set LZ4F_preferences_t.frameInfo.contentChecksumFlag #} p (cint Lz4fNocontentchecksum)
        {# set LZ4F_preferences_t.frameInfo.frameType #} p (cint Lz4fFrame)
        {# set LZ4F_preferences_t.frameInfo.contentSize #} p 0
        {# set LZ4F_preferences_t.frameInfo.dictID #} p 0
        {# set LZ4F_preferences_t.frameInfo.blockChecksumFlag #} p (cint Lz4fNoblockchecksum)

        {# set LZ4F_preferences_t.compressionLevel #} p (fromIntegral i)
        {# set LZ4F_preferences_t.autoFlush #} p 0
        {# set LZ4F_preferences_t.favorDecSpeed #} p 0