Lazy.hsc

{-# LANGUAGE BangPatterns, EmptyDataDecls, ForeignFunctionInterface #-}

-- |
-- Module:      Codec.Compression.Snappy
-- Copyright:   © 2011 MailRank, Inc.
-- License:     Apache
-- Maintainer:  ARJANEN Loïc Jean David <ljd@luigiscorner.mu>
-- Stability:   experimental
-- Portability: portable
--
-- This module provides fast, pure zero-copy compression and
-- decompression of lazy 'ByteString' data using the Snappy format.
--
-- Although these functions operate on lazy 'ByteString's, they
-- consume the data /strictly/: they do not produce any output until
-- they have consumed all of the input, and they produce the output in
-- a single large chunk.
--
-- If your data is already in the form of a lazy 'ByteString', it is
-- likely more efficient to use these functions than to convert your
-- data to and from strict ByteStrings, as you can avoid the
-- additional allocation and copying that would entail.

module Codec.Compression.Snappy.Lazy
    (
      compress
    , decompress
    ) where

#include "hs_snappy.h"

import Codec.Compression.Snappy.Internal (check, maxCompressedLength)
import Control.Exception (bracket)
import Data.ByteString.Internal hiding (ByteString(BS))
import Data.ByteString.Lazy.Internal (ByteString(..))
import Data.Word (Word8, Word32)
import Foreign.C.Types (CInt(..), CSize(..))
import Foreign.ForeignPtr (touchForeignPtr, withForeignPtr, newForeignPtr_)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Marshal.Array (withArray)
import Foreign.Marshal.Utils (with)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (Storable(..))
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L

newtype BS = BS B.ByteString

data BSSource

instance Storable BS where
    sizeOf _    = (#size struct BS)
    alignment _ = alignment (undefined :: Ptr CInt)
    poke ptr (BS (PS fp off len)) = withForeignPtr fp $ \p -> do
      (#poke struct BS, ptr) ptr (p `plusPtr` off)
      (#poke struct BS, off) ptr (0::CSize)
      (#poke struct BS, len) ptr len
    {-# INLINE poke #-}
    peek p = do
      pointer <- (#peek struct BS, ptr) p
      foreignPointer <- newForeignPtr_ pointer
      offset <- (#peek struct BS, off) p
      peekedLength <- (#peek struct BS, len) p
      return (BS (PS foreignPointer offset peekedLength))
    {-# INLINE peek #-}

-- | Compress data into the Snappy format.
compress :: ByteString -> ByteString
compress bs = unsafePerformIO . withChunks bs $ \chunkPtr numChunks len -> do
  let dlen0 = maxCompressedLength len
  dfp <- mallocByteString dlen0
  withForeignPtr dfp $ \dptr -> do
    with (fromIntegral dlen0) $ \dlenPtr -> do
      c_CompressChunks chunkPtr (fromIntegral numChunks)
                       (fromIntegral len) dptr dlenPtr
      dlen <- fromIntegral `fmap` peek dlenPtr
      if dlen == 0
        then return Empty
        else return (Chunk (PS dfp 0 dlen) Empty)

-- | Decompress data in the Snappy format.
--
-- If the input is not compressed or is corrupt, an exception will be
-- thrown.
decompress :: ByteString -> ByteString
decompress bs = unsafePerformIO . withChunks bs $ \chunkPtr numChunks len ->
  bracket (c_NewSource chunkPtr (fromIntegral numChunks) (fromIntegral len))
          c_DeleteSource $ \srcPtr -> do
    alloca $ \dlenPtr -> do
      check "Lazy.decompress" $ c_GetUncompressedLengthChunks srcPtr dlenPtr
      dlen <- fromIntegral `fmap` peek dlenPtr
      if dlen == 0
        then return L.empty
        else do
          dfp <- mallocByteString dlen
          withForeignPtr dfp $ \dptr -> do
            check "Lazy.decompress" $ c_UncompressChunks srcPtr dptr
            return (Chunk (PS dfp 0 dlen) Empty)

withChunks :: ByteString -> (Ptr BS -> Int -> Int -> IO a) -> IO a
withChunks bs act = do
  let len = fromIntegral (L.length bs)
  let chunks = L.toChunks bs
  r <- withArray (map BS chunks) $ \chunkPtr ->
       act chunkPtr (length chunks) len
  foldr (\(PS fp _ _) _ -> touchForeignPtr fp) (return ()) chunks
  return r

foreign import ccall unsafe "hs_snappy.h _hsnappy_CompressChunks"
    c_CompressChunks :: Ptr BS -> CSize -> CSize -> Ptr Word8 -> Ptr CSize
                     -> IO ()

foreign import ccall unsafe "hs_snappy.h _hsnappy_NewSource"
    c_NewSource :: Ptr BS -> CSize -> CSize -> IO (Ptr BSSource)

foreign import ccall unsafe "hs_snappy.h _hsnappy_DeleteSource"
    c_DeleteSource :: Ptr BSSource -> IO ()

foreign import ccall unsafe "hs_snappy.h _hsnappy_UncompressChunks"
    c_UncompressChunks :: Ptr BSSource -> Ptr Word8 -> IO Int

foreign import ccall unsafe "hs_snappy.h _hsnappy_GetUncompressedLengthChunks"
    c_GetUncompressedLengthChunks :: Ptr BSSource -> Ptr Word32 -> IO Int