Lz4.hs

-- | 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
                 , decompress
                 -- * 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.ForeignPtr           (ForeignPtr, castForeignPtr,
                                               mallocForeignPtrBytes,
                                               newForeignPtr, withForeignPtr)
import           Foreign.Marshal.Alloc        (alloca)
import           Foreign.Ptr                  (castPtr, nullPtr)
import           Foreign.Storable             (peek, poke)
import           System.IO.Unsafe             (unsafePerformIO)

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

-- | Lazily decompress a frame
decompress :: BSL.ByteString -> BSL.ByteString
decompress bs = runST $ do

    let bss = BSL.toChunks bs

    (ctx, buf) <- LazyST.unsafeIOToST $ do
        (err, preCtx) <- lZ4FCreateDecompressionContext lZ4FGetVersion
        check err
        ctx <- castForeignPtr <$> newForeignPtr lZ4FFreeCompressionContext (castPtr preCtx)
        dstBuf <- mallocForeignPtrBytes (64 * 1024)
        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 (64 * 1024)
                            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 bs = runST $ do

    let bss = BSL.toChunks bs

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

    pure $ BSL.fromChunks (header:rest)

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

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

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

          update :: LzCtxPtr -> BS.ByteString -> LazyST.ST s BS.ByteString
          update ctx b = LazyST.unsafeIOToST $
            BS.unsafeUseAsCStringLen b $ \(buf, sz) -> do
                let expectedSz = lZ4FCompressBound (fromIntegral sz) nullPtr
                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