hs-brotli-0.1.0.0: src/Codec/Compression/Brotli.hs
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
module Codec.Compression.Brotli (
Compress(..)
, compress
, decompress
, BrotliStream(..)
, Chunk(..)
, compressor
, decompressor
, maxCompressedSize
, CompressionSettings(..)
, defaultCompressionSettings
, setCompressionSettings
, I.minWindowBits
, I.maxWindowBits
, I.minInputBlockBits
, I.maxInputBlockBits
, I.minQuality
, I.maxQuality
, I.encoderModeGeneric
, I.encoderModeText
, I.encoderModeFont
) where
import Control.Monad (when, unless, forM)
import Control.Exception (SomeException, assert, handle, bracket, throw)
import qualified Codec.Compression.Brotli.Internal as I
import qualified Data.ByteString as B
import qualified Data.ByteString.Unsafe as B
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Internal as LI
import Foreign.ForeignPtr
import Data.Maybe (catMaybes)
import Data.IORef
import Data.Int
import Data.Word
import Foreign.C
import Foreign.Ptr (Ptr, castPtr, nullPtr, plusPtr)
import Foreign.Marshal (alloca, allocaBytes, callocBytes, free)
import Foreign.Storable (sizeOf, peek, poke)
import GHC.Int
import GHC.Types
import System.IO.Unsafe
data CompressionSettings = CompressionSettings
{ compressionQuality :: !Word32
, compressionWindow :: !Word32
, compressionMode :: !I.BrotliEncoderMode
, compressionBufferSize :: !Int
, compressionBlockSize :: !(Maybe Word32)
, compressionDisableLiteralContextModeling :: !(Maybe Word32)
, compressionSizeHint :: !(Maybe Word32)
}
-- | The default set of parameters for compression. This is typically used with the compressWith function with specific parameters overridden.
defaultCompressionSettings :: CompressionSettings
defaultCompressionSettings = CompressionSettings
{ compressionQuality = I.defaultQuality
, compressionWindow = I.defaultWindow
, compressionMode = I.defaultMode
, compressionBufferSize = 0
, compressionBlockSize = Nothing
, compressionDisableLiteralContextModeling = Nothing
, compressionSizeHint = Nothing
}
setCompressionSettings :: I.BrotliEncoderState -> CompressionSettings -> IO ()
setCompressionSettings st CompressionSettings{..}= do
r1 <- I.encoderSetParameter st I.mode $ fromIntegral $ I.fromBrotliEncoderMode compressionMode
r2 <- I.encoderSetParameter st I.quality compressionQuality
r3 <- I.encoderSetParameter st I.lz77WindowSize compressionWindow
mr4 <- forM compressionBlockSize $ I.encoderSetParameter st I.lz77BlockSize
mr5 <- forM compressionDisableLiteralContextModeling $ I.encoderSetParameter st I.disableLiteralContextModeling
mr6 <- forM compressionSizeHint $ I.encoderSetParameter st I.brotliParamSizeHint
if any (/= 1) $ r1 : r2 : r3 : catMaybes [mr4, mr5, mr6]
then error "Invalid compression setting parameter"
else pure ()
-- | Like compress but with the ability to specify various compression parameters. Typical usage:
--
-- > compressWith defaultCompressionSettings { ... }
--
-- In particular you can set the compression level:
--
-- > compressWith defaultCompressParams { compressionQuality = minQuality }
class Compress a b where
compressWith :: CompressionSettings -> a-> b
-- | Compress a stream of data into the brotli format.
--
-- This uses the default compression parameters. In particular it uses the highest compression level which favours a higher compression ratio over compression speed.
--
-- Use compressWith to adjust the compression level or other compression parameters.
compress :: Compress a b => a -> b
compress = compressWith defaultCompressionSettings
instance Compress B.ByteString B.ByteString where
compressWith CompressionSettings{..} b = unsafePerformIO $ do
let estBufSize = fromIntegral $ I.maxCompressedSize $ fromIntegral $ B.length b
res <- alloca $ \outSize -> do
poke outSize estBufSize
BI.createAndTrim (fromIntegral estBufSize) $ \outBuf -> B.unsafeUseAsCStringLen b $ \(inPtr, inLen) -> do
ok <- I.encoderCompress (fromIntegral compressionQuality) (fromIntegral compressionWindow) compressionMode (fromIntegral inLen) (castPtr inPtr) outSize outBuf
os <- peek outSize
if (ok /= 1)
then error "Compression error or output buffer is too small"
else pure $ fromIntegral os
pure res
isTrue :: CInt -> Bool
isTrue (CInt (I32# x)) = isTrue# x
maxCompressedSize :: Int -> Int
maxCompressedSize = fromIntegral . I.maxCompressedSize . fromIntegral
hasMoreOutput :: I.BrotliEncoderState -> IO Bool
hasMoreOutput = fmap isTrue . I.encoderHasMoreOutput
isFinished :: I.BrotliEncoderState -> IO Bool
isFinished = fmap isTrue . I.encoderIsFinished
takeOutput :: I.BrotliEncoderState -> IO B.ByteString
takeOutput st = alloca $ \sizeP -> do
poke sizeP 0
ptr <- I.encoderTakeOutput st sizeP
takeSize <- peek sizeP
B.packCStringLen (castPtr ptr, fromIntegral takeSize)
data StreamVars = StreamVars
{ availableIn :: !(Ptr CSize)
, availableOut :: !(Ptr CSize)
, totalOut :: !(Ptr CSize)
, nextIn :: !(Ptr (Ptr Word8))
, nextOut :: !(Ptr (Ptr Word8))
} deriving (Show)
freeStreamVars :: StreamVars -> IO ()
freeStreamVars = free . availableIn
createStreamVars :: IO StreamVars
createStreamVars = do
bs <- callocBytes (3 * sizeOf (0 :: CSize) + 2 * sizeOf (nullPtr :: Ptr Word8))
let aiPtr = bs
aoPtr = plusPtr aiPtr (sizeOf (0 :: CSize))
toPtr = plusPtr aoPtr (sizeOf (0 :: CSize))
niPtr = castPtr $ plusPtr toPtr (sizeOf (nullPtr :: Ptr Word8))
noPtr = plusPtr niPtr (sizeOf (nullPtr :: Ptr Word8))
vs@StreamVars{..} = StreamVars aiPtr aoPtr toPtr niPtr noPtr
pure vs
newtype EncoderFeedResponse = EncoderFeedResponse
{ pendingInput :: B.ByteString
}
-- CAUTION: we aren't ensuring that bytestrings stay alive since the contents are
-- poked in and outlive the function (streaming and all)
--
-- Must use only within the context of the bytestring be alive via an external unsafeUseAsCStringLen
feedEncoder' :: I.BrotliEncoderState -> StreamVars -> Int -> B.ByteString -> IO CSize
feedEncoder' st vs@StreamVars{..} bufSize bs = B.unsafeUseAsCStringLen bs $ \(bsP, len) -> do
poke availableIn (fromIntegral len)
poke nextIn (castPtr bsP)
res <- isTrue <$> I.encoderCompressStream st I.encoderOperationProcess availableIn nextIn availableOut nextOut totalOut
unless res $ error "Unknown stream encoding failure"
peek availableIn
-- CAUTION: we aren't ensuring that bytestrings stay alive since the contents are
-- poked in and outlive the function (streaming and all)
--
-- Must use only within the context of the bytestring be alive via an external unsafeUseAsCStringLen
feedEncoder :: I.BrotliEncoderState -> StreamVars -> Int -> B.ByteString -> IO EncoderFeedResponse
feedEncoder st vs@StreamVars{..} bufSize bs = B.unsafeUseAsCStringLen bs $ \(bsP, len) -> do
poke availableIn (fromIntegral len)
poke nextIn (castPtr bsP)
res <- isTrue <$> I.encoderCompressStream st I.encoderOperationProcess availableIn nextIn availableOut nextOut totalOut
unless res $ error "Unknown stream encoding failure"
unconsumedBytesCount <- peek availableIn
unconsumedBytesP <- peek nextIn
unusedInput <- if unconsumedBytesCount == 0
then pure B.empty
else B.packCStringLen (castPtr unconsumedBytesP, fromIntegral unconsumedBytesCount)
pure $ EncoderFeedResponse unusedInput
encoderMaybeTakeOutput :: I.BrotliEncoderState -> Int -> IO (Maybe B.ByteString)
encoderMaybeTakeOutput st bufSize = do
takeOut <- isTrue <$> I.encoderHasMoreOutput st
if takeOut
then alloca $ \s -> do
poke s $ fromIntegral bufSize
bsp <- I.encoderTakeOutput st s
len <- peek s
Just <$> B.packCStringLen (castPtr bsp, fromIntegral len)
else pure Nothing
encoderTakeRestAvailable :: I.BrotliEncoderState -> IO () -> Int -> L.ByteString -> IO L.ByteString
encoderTakeRestAvailable st cleanup bufSize graft = do
out <- encoderMaybeTakeOutput st bufSize
case out of
Nothing -> pure graft
Just bs -> do
rest <- unsafeInterleaveIO $ encoderTakeRestAvailable st cleanup bufSize graft
pure $ LI.Chunk bs rest
-- | Note that this should be called until returned bytestring is empty. Once is not enough.
finishStream :: I.BrotliEncoderState -> StreamVars -> IO () -> Int -> IO L.ByteString
finishStream st StreamVars{..} cleanup bufSize = do
poke availableIn 0
poke nextIn nullPtr
res <- isTrue <$> I.encoderCompressStream st I.encoderOperationFinish availableIn nextIn availableOut nextOut totalOut
unless res $ error "Unknown stream encoding failure"
encoderTakeRestAvailable st cleanup bufSize L.empty
pOff :: Int -> Ptr a -> Ptr b
pOff n p = castPtr $ plusPtr p (n * sizeOf p)
pushNoCheck :: B.ByteString -> L.ByteString -> L.ByteString
pushNoCheck = LI.Chunk
instance Compress L.ByteString L.ByteString where
compressWith settings b = unsafePerformIO $ do
inst <- I.createEncoder
setCompressionSettings inst settings
vars <- createStreamVars
poke (availableOut vars) 0
poke (nextOut vars) nullPtr
let cleanup = freeStreamVars vars >> I.destroyEncoder inst
lazyCompress cleanup inst vars b
where
lazyCompress cleanup st vars c = unsafeInterleaveIO $ readChunks cleanup st vars c
readChunks cleanup st vars c = do
case c of
LI.Chunk bs next -> handle (\(e :: SomeException) -> cleanup >> throw e) $ do
(EncoderFeedResponse unusedInput) <- feedEncoder st vars (compressionBufferSize settings) bs
-- NOTE: LI.chunk checks for empty string results so we don't have to worry about empty chunks ourselves.
rest <- lazyCompress cleanup st vars $ LI.chunk unusedInput next
encoderTakeRestAvailable st cleanup (compressionBufferSize settings) rest
LI.Empty -> finishStream st vars (freeStreamVars vars >> I.destroyEncoder st) $ compressionBufferSize settings
data DecompressionVars = DecompressionVars
{ dAvailableInput :: !(Ptr CSize)
, dAvailableOut :: !(Ptr CSize)
, dTotalOut :: !(Ptr CSize)
, dNextIn :: !(Ptr (Ptr Word8))
, dNextOut :: !(Ptr (Ptr Word8))
}
createDecompressionVars :: IO DecompressionVars
createDecompressionVars = do
bs <- callocBytes (3 * sizeOf (0 :: CSize) + 2 * sizeOf (nullPtr :: Ptr Word8))
let aiPtr = bs
aoPtr = plusPtr aiPtr (sizeOf (0 :: CSize))
toPtr = plusPtr aoPtr (sizeOf (0 :: CSize))
niPtr = castPtr $ plusPtr toPtr (sizeOf (nullPtr :: Ptr Word8))
noPtr = plusPtr niPtr (sizeOf (nullPtr :: Ptr Word8))
pure $ DecompressionVars aiPtr aoPtr toPtr niPtr noPtr
destroyDecompressionVars :: DecompressionVars -> IO ()
destroyDecompressionVars = free . dAvailableInput
-- | Decompress a stream of data in the brotli format.
--
-- There are a number of errors that can occur. In each case an exception will
-- be thrown.
--
-- Note that the decompression is performed /lazily/. Errors in the data stream
-- may not be detected until the end of the stream is demanded (since it is
-- only at the end that the final checksum can be checked). If this is
-- important to you, you must make sure to consume the whole decompressed
-- stream before doing any IO action that depends on it.
--
decompress :: L.ByteString -> L.ByteString
decompress b = unsafePerformIO $ do
st <- I.createDecoder
vs <- createDecompressionVars
poke (dAvailableOut vs) 0
poke (dNextOut vs) nullPtr
lazyDecompress st vs b
where
lazyDecompress st vs rest = unsafeInterleaveIO $ writeChunks st vs rest
writeChunks st vs@DecompressionVars{..} lbs = do
case lbs of
LI.Chunk bs rest -> do
v@(res, unconsumed') <- B.unsafeUseAsCStringLen bs $ \(strP, strLen) -> do
poke dAvailableInput $ fromIntegral strLen
poke dNextIn $ castPtr strP
res <- I.decoderDecompressStream st dAvailableInput dNextIn dAvailableOut dNextOut dTotalOut
remainingInputBytes <- peek dAvailableInput
compressedBytesPtr <- peek dNextIn
unconsumed' <- B.packCStringLen (castPtr compressedBytesPtr, fromIntegral remainingInputBytes)
pure (res, unconsumed')
case I.decoderResult res of
I.Success -> do
-- allTheRest <- takeRestAvailable st (I.destroyDecoder st >> destroyDecompressionVars vs) L.empty
-- return allTheRest
pure L.empty
I.NeedsMoreInput -> do
lazyDecompress st vs $ LI.chunk unconsumed' rest
I.NeedsMoreOutput -> do
-- Sneak invariant breaking here by pushing what is quite possibly an empty Chunk.
-- this is intentional because we need one last empty string to trigger either success or error
-- depending on whether the string shouldn't have ended there
afterOut <- lazyDecompress st vs $ LI.Chunk unconsumed' rest
decoderTakeRestAvailable st (pure ()) afterOut
-- TODO need a safe version of decompress too
I.DecoderError e -> I.destroyDecoder st >> throw e
LI.Empty -> do
I.destroyDecoder st
destroyDecompressionVars vs
throw $ I.BrotliDecoderErrorCode 2
decoderMaybeTakeOutput :: I.BrotliDecoderState -> IO (Maybe B.ByteString)
decoderMaybeTakeOutput st = do
takeIn <- isTrue <$> I.decoderHasMoreOutput st
if takeIn
then alloca $ \s -> do
poke s 0 -- TODO settings
bsp <- I.decoderTakeOutput st s
len <- peek s
Just <$> B.packCStringLen (castPtr bsp, fromIntegral len)
else pure Nothing
decoderTakeRestAvailable :: I.BrotliDecoderState -> IO () -> L.ByteString -> IO L.ByteString
decoderTakeRestAvailable st cleanup graft = do
out <- decoderMaybeTakeOutput st
case out of
Nothing -> cleanup >> pure graft
Just bs -> do
rest <- unsafeInterleaveIO $ decoderTakeRestAvailable st cleanup graft
pure $ LI.Chunk bs rest
data Chunk
= Chunk !B.ByteString
| Flush
deriving (Show, Eq)
data BrotliStream input
= Produce !B.ByteString (IO (BrotliStream input))
| Consume (input -> IO (BrotliStream input))
| Error
| Done
withEncoder :: ForeignPtr I.BrotliEncoderState -> (I.BrotliEncoderState -> IO a) -> IO a
withEncoder p f = withForeignPtr p (f . I.BrotliEncoderState)
-- | A strict, streaming compressor. This allows compressing
-- values in constant memory in addition to giving fine-grained
-- control of when to flush data in the stream.
compressor :: CompressionSettings -> IO (BrotliStream Chunk)
compressor settings = do
(I.BrotliEncoderState encoder) <- I.createEncoder
efp <- newForeignPtr I.destroyEncoder_ptr encoder
withEncoder efp $ \encoder -> setCompressionSettings encoder settings
vars <- createStreamVars
pure $ go efp vars
where
go efp vars = Consume (consume B.empty)
where
consume
:: B.ByteString
{- ^ Presently poked bytestring
Must `touch` underlying ForeignPtr until
the underlying contents have been consumed
In other words, consumption must be wrapped in an unsafeUseAsCStringLen to keep it alive.
We could be a bit more fine-grained on things if we use bytestring internals, but maybe
not worth the hassle?
-}
-> Chunk
-> IO (BrotliStream Chunk)
consume !currentBs !chunk = B.unsafeUseAsCStringLen currentBs $ \_ -> do
let StreamVars{..} = vars
case chunk of
Chunk bs -> do
-- print vars
previousUnconsumedSize <- peek availableIn
if previousUnconsumedSize == 0 && B.null bs
then done
else withEncoder efp $ \encoder -> do
unconsumedSize <-
if previousUnconsumedSize == 0
then feedEncoder'
encoder
vars
(compressionBufferSize settings)
bs
else do
_ <- I.encoderCompressStream
encoder
I.encoderOperationProcess
availableIn
nextIn
availableOut
nextOut
totalOut
-- TODO assert result
peek availableIn
let bytestringRef = if unconsumedSize > 0 then currentBs else B.empty
moreOutput <- hasMoreOutput encoder
if moreOutput
then produce bytestringRef $ if previousUnconsumedSize == 0
then Just chunk
else Nothing
else pure $ Consume (consume bytestringRef)
Flush -> withEncoder efp $ \encoder -> do
-- TODO make appropriate assertions here
-- around in state & out state
I.encoderCompressStream
encoder
I.encoderOperationFlush
availableIn
nextIn
availableOut
nextOut
totalOut
-- ls <- (,) <$> peek availableOut <*> peek availableIn
-- print ls
hasMore <- hasMoreOutput encoder
produce B.empty Nothing
produce :: B.ByteString -> Maybe Chunk -> IO (BrotliStream Chunk)
produce currentInput unusedInput = withEncoder efp $ \encoder -> do
-- assert: this function is only called from other functions
-- when output is guaranteed
out <- takeOutput encoder
hasMore <- hasMoreOutput encoder
pure $ Produce out $ if hasMore
then produce currentInput unusedInput
else
maybe
(pure $ Consume $ consume B.empty)
(\c -> consume (case c of
Chunk bs -> bs
Flush -> B.empty) c)
unusedInput
-- err = undefined
done :: IO (BrotliStream Chunk)
done = withEncoder efp $ \encoder -> do
let StreamVars {..} = vars
poke availableIn 0
poke nextIn nullPtr
I.encoderCompressStream
encoder
I.encoderOperationFinish
availableIn
nextIn
availableOut
nextOut
totalOut
done'
done' = withEncoder efp $ \encoder -> do
out <- encoderMaybeTakeOutput encoder (compressionBufferSize settings)
case out of
Nothing -> do
freeStreamVars vars
pure Done
Just str -> pure $ Produce str done'
-- | A strict, streaming decompressor. This allows decompressing
-- values in constant memory as long as you don't need the full output
-- at one time.
decompressor :: IO (BrotliStream B.ByteString)
decompressor = do
st <- I.createDecoder
vs <- createDecompressionVars
poke (dAvailableOut vs) 0
poke (dNextOut vs) nullPtr
go st vs
where
go st vs@DecompressionVars {..} = pure $ consume B.empty
where
consume :: B.ByteString -> BrotliStream B.ByteString
consume leftover =
Consume $ \bs -> do
v@(res, unconsumed') <-
B.unsafeUseAsCStringLen
(if B.null leftover
then bs
else if B.null bs
then B.empty
else leftover `mappend` bs) $ \(strP, strLen) -> do
poke dAvailableInput (fromIntegral strLen)
poke dNextIn (castPtr strP)
res <-
I.decoderDecompressStream
st
dAvailableInput
dNextIn
dAvailableOut
dNextOut
dTotalOut
remainingInputBytes <- peek dAvailableInput
compressedBytesPtr <- peek dNextIn
unconsumed' <-
B.packCStringLen
( castPtr compressedBytesPtr
, fromIntegral remainingInputBytes)
pure (res, unconsumed')
case I.decoderResult res of
I.Success -> done
I.NeedsMoreInput -> pure $ consume B.empty
I.NeedsMoreOutput -> produce unconsumed'
I.DecoderError e -> err -- I.destroyDecoder st >> throw e
produce rem = do
out <- decoderMaybeTakeOutput st
case out of
Nothing -> pure (consume rem)
Just bs -> pure $ Produce bs (produce rem)
err = pure Error
done = do
I.destroyDecoder st
destroyDecompressionVars vs
pure Done