-- Copyright (c) 2016-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is licensed under the BSD-style license found in
-- the LICENSE file in the root directory of this source tree. An
-- additional grant of patent rights can be found in the PATENTS file
-- in the same directory.

-- |
-- Module      : Codec.Compression.Zstd.Internal
-- Copyright   : (c) 2016-present, Facebook, Inc. All rights reserved.
--
-- License     : BSD3
-- Maintainer  : bryano@fb.com
-- Stability   : experimental
-- Portability : GHC
--
-- A fast lossless compression algorithm, targeting real-time
-- compression scenarios at zlib-level and better compression ratios.

module Codec.Compression.Zstd.Internal
    (
      CCtx(..)
    , DCtx(..)
    , compressWith
    , decompressWith
    , decompressedSize
    , withCCtx
    , withDCtx
    , withDict
    , trainFromSamples
    , getDictID
    ) where

import Codec.Compression.Zstd.Types (Decompress(..), Dict(..))
import Control.Exception.Base (bracket)
import Data.ByteString.Internal (ByteString(..))
import Data.Word (Word, Word8)
import Foreign.C.Types (CInt, CSize)
import Foreign.Marshal.Array (withArray)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, castPtr, plusPtr)
import System.IO.Unsafe (unsafePerformIO)
import qualified Codec.Compression.Zstd.FFI as C
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as B

compressWith
    :: String
    -> (Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> CInt -> IO CSize)
    -> Int
    -> ByteString
    -> IO ByteString
compressWith :: String
-> (Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> CInt -> IO CSize)
-> Int
-> ByteString
-> IO ByteString
compressWith String
name Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> CInt -> IO CSize
compressor Int
level (PS ForeignPtr Word8
sfp Int
off Int
len)
  | Int
level Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
level Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
C.maxCLevel
              = String -> String -> IO ByteString
forall a. String -> String -> a
bail String
name String
"unsupported compression level"
  | Bool
otherwise =
  ForeignPtr Word8 -> (Ptr Word8 -> IO ByteString) -> IO ByteString
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO ByteString) -> IO ByteString)
-> (Ptr Word8 -> IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp -> do
    maxSize <- CSize -> IO CSize
C.compressBound (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    dfp <- B.mallocByteString (fromIntegral maxSize)
    withForeignPtr dfp $ \Ptr Word8
dst -> do
      let src :: Ptr b
src = Ptr Word8
sp Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off
      csz <- Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> CInt -> IO CSize
compressor Ptr Word8
dst CSize
maxSize Ptr Word8
forall {b}. Ptr b
src (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len) (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
level)
      handleError csz name $ do
        let size = CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
csz
        if csz < 128 || csz >= maxSize `div` 2
        then return (PS dfp 0 size)
        else B.create size $ \Ptr Word8
p -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
B.memcpy Ptr Word8
p Ptr Word8
dst Int
size

-- | Return the decompressed size of a compressed payload, as stored
-- in the payload's header.
--
-- The returned value will be `Nothing` if it is either not known
-- (probably because the payload was compressed using a streaming
-- API), empty, or too large to fit in an 'Int'.
--
-- /Note:/ this value should not be trusted, as it can be controlled
-- by an attacker.
decompressedSize :: ByteString -> Maybe Int
decompressedSize :: ByteString -> Maybe Int
decompressedSize (PS ForeignPtr Word8
fp Int
off Int
len) =
  IO (Maybe Int) -> Maybe Int
forall a. IO a -> a
unsafePerformIO (IO (Maybe Int) -> Maybe Int)
-> ((Ptr Word8 -> IO (Maybe Int)) -> IO (Maybe Int))
-> (Ptr Word8 -> IO (Maybe Int))
-> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr Word8 -> (Ptr Word8 -> IO (Maybe Int)) -> IO (Maybe Int)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO (Maybe Int)) -> Maybe Int)
-> (Ptr Word8 -> IO (Maybe Int)) -> Maybe Int
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> do
    sz <- Ptr (ZonkAny 2) -> CSize -> IO CULLong
forall src. Ptr src -> CSize -> IO CULLong
C.getDecompressedSize (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr (ZonkAny 2)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    return $ if sz == 0 || sz > fromIntegral (maxBound :: Int)
             then Nothing
             else Just (fromIntegral sz)

decompressWith :: (Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> IO CSize)
               -> ByteString
               -> IO Decompress
decompressWith :: (Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> IO CSize)
-> ByteString -> IO Decompress
decompressWith Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> IO CSize
decompressor (PS ForeignPtr Word8
sfp Int
off Int
len) = do
  ForeignPtr Word8 -> (Ptr Word8 -> IO Decompress) -> IO Decompress
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO Decompress) -> IO Decompress)
-> (Ptr Word8 -> IO Decompress) -> IO Decompress
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sp -> do
    let src :: Ptr b
src = Ptr Word8
sp Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off
    dstSize <- Ptr (ZonkAny 1) -> CSize -> IO CULLong
forall src. Ptr src -> CSize -> IO CULLong
C.getDecompressedSize Ptr (ZonkAny 1)
forall {b}. Ptr b
src (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
    if dstSize == 0
      then return Skip
      else if dstSize > fromIntegral (maxBound :: Int)
           then return (Error "invalid compressed payload size")
           else do
      dfp <- B.mallocByteString (fromIntegral dstSize)
      size <- withForeignPtr dfp $ \Ptr Word8
dst ->
        Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> IO CSize
decompressor Ptr Word8
dst (CULLong -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral CULLong
dstSize) Ptr Word8
forall {b}. Ptr b
src (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
      return $ if C.isError size
               then Error (C.getErrorName size)
               else Decompress (PS dfp 0 (fromIntegral size))

-- | Compression context.
newtype CCtx = CCtx { CCtx -> Ptr CCtx
getCCtx :: Ptr C.CCtx }

-- | Allocate a compression context, run an action that may reuse the
-- context as many times as it needs, then free the context.
withCCtx :: (CCtx -> IO a) -> IO a
withCCtx :: forall a. (CCtx -> IO a) -> IO a
withCCtx CCtx -> IO a
act =
  IO CCtx -> (CCtx -> IO ()) -> (CCtx -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket ((Ptr CCtx -> CCtx) -> IO (Ptr CCtx) -> IO CCtx
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ptr CCtx -> CCtx
CCtx (String -> IO (Ptr CCtx) -> IO (Ptr CCtx)
forall a. String -> IO (Ptr a) -> IO (Ptr a)
C.checkAlloc String
"withCCtx" IO (Ptr CCtx)
C.createCCtx))
          (Ptr CCtx -> IO ()
C.freeCCtx (Ptr CCtx -> IO ()) -> (CCtx -> Ptr CCtx) -> CCtx -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CCtx -> Ptr CCtx
getCCtx) CCtx -> IO a
act

-- | Decompression context.
newtype DCtx = DCtx { DCtx -> Ptr DCtx
getDCtx :: Ptr C.DCtx }

-- | Allocate a decompression context, run an action that may reuse the
-- context as many times as it needs, then free the context.
withDCtx :: (DCtx -> IO a) -> IO a
withDCtx :: forall a. (DCtx -> IO a) -> IO a
withDCtx DCtx -> IO a
act =
  IO DCtx -> (DCtx -> IO ()) -> (DCtx -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket ((Ptr DCtx -> DCtx) -> IO (Ptr DCtx) -> IO DCtx
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ptr DCtx -> DCtx
DCtx (String -> IO (Ptr DCtx) -> IO (Ptr DCtx)
forall a. String -> IO (Ptr a) -> IO (Ptr a)
C.checkAlloc String
"withDCtx" IO (Ptr DCtx)
C.createDCtx))
          (Ptr DCtx -> IO ()
C.freeDCtx (Ptr DCtx -> IO ()) -> (DCtx -> Ptr DCtx) -> DCtx -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DCtx -> Ptr DCtx
getDCtx) DCtx -> IO a
act

withDict :: Dict -> (Ptr dict -> CSize -> IO a) -> IO a
withDict :: forall dict a. Dict -> (Ptr dict -> CSize -> IO a) -> IO a
withDict (Dict (PS ForeignPtr Word8
fp Int
off Int
len)) Ptr dict -> CSize -> IO a
act =
  ForeignPtr Word8 -> (Ptr Word8 -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO a) -> IO a) -> (Ptr Word8 -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Ptr dict -> CSize -> IO a
act (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr dict
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)

-- | Create and train a compression dictionary from a collection of
-- samples.
--
-- To create a well-trained dictionary, here are some useful
-- guidelines to keep in mind:
--
-- * A reasonable dictionary size is in the region of 100 KB.  (Trying
--   to specify a dictionary size of less than a few hundred bytes will
--   probably fail.)
--
-- * To train the dictionary well, it is best to supply a few thousand
--   training samples.
--
-- * The combined size of all training samples should be 100 or more
--   times larger than the size of the dictionary.
trainFromSamples :: Int
                 -- ^ Maximum size of the compression dictionary to
                 -- create. The actual dictionary returned may be
                 -- smaller.
                 -> [ByteString]
                 -- ^ Samples to train with.
                 -> Either String Dict
trainFromSamples :: Int -> [ByteString] -> Either String Dict
trainFromSamples Int
capacity [ByteString]
samples = IO (Either String Dict) -> Either String Dict
forall a. IO a -> a
unsafePerformIO (IO (Either String Dict) -> Either String Dict)
-> IO (Either String Dict) -> Either String Dict
forall a b. (a -> b) -> a -> b
$
  [Int]
-> (Ptr Int -> IO (Either String Dict)) -> IO (Either String Dict)
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray ((ByteString -> Int) -> [ByteString] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ByteString -> Int
B.length [ByteString]
samples) ((Ptr Int -> IO (Either String Dict)) -> IO (Either String Dict))
-> (Ptr Int -> IO (Either String Dict)) -> IO (Either String Dict)
forall a b. (a -> b) -> a -> b
$ \Ptr Int
sizes -> do
    dfp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
capacity
    let PS sfp _ _ = B.concat samples
    withForeignPtr dfp $ \Ptr Word8
dict ->
      ForeignPtr Word8
-> (Ptr Word8 -> IO (Either String Dict))
-> IO (Either String Dict)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
sfp ((Ptr Word8 -> IO (Either String Dict)) -> IO (Either String Dict))
-> (Ptr Word8 -> IO (Either String Dict))
-> IO (Either String Dict)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
sampPtr -> do
        dsz <- Ptr Word8 -> CSize -> Ptr Word8 -> Ptr CSize -> CUInt -> IO CSize
forall dict samples.
Ptr dict -> CSize -> Ptr samples -> Ptr CSize -> CUInt -> IO CSize
C.trainFromBuffer
               Ptr Word8
dict (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
capacity)
               Ptr Word8
sampPtr (Ptr Int -> Ptr CSize
forall a b. Ptr a -> Ptr b
castPtr Ptr Int
sizes) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([ByteString] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ByteString]
samples))
        if C.isError dsz
          then return (Left (C.getErrorName dsz))
          else fmap (Right . Dict) $ do
            let size = CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
dsz
            if size < 128 || size >= capacity `div` 2
            then return (PS dfp 0 size)
            else B.create size $ \Ptr Word8
p -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
B.memcpy Ptr Word8
p Ptr Word8
dict Int
size

-- | Return the identifier for the given dictionary, or 'Nothing' if
-- not a valid dictionary.
getDictID :: Dict -> Maybe Word
getDictID :: Dict -> Maybe Word
getDictID Dict
dict = IO (Maybe Word) -> Maybe Word
forall a. IO a -> a
unsafePerformIO (IO (Maybe Word) -> Maybe Word) -> IO (Maybe Word) -> Maybe Word
forall a b. (a -> b) -> a -> b
$ do
  n <- Dict -> (Ptr (ZonkAny 0) -> CSize -> IO CUInt) -> IO CUInt
forall dict a. Dict -> (Ptr dict -> CSize -> IO a) -> IO a
withDict Dict
dict Ptr (ZonkAny 0) -> CSize -> IO CUInt
forall dict. Ptr dict -> CSize -> IO CUInt
C.getDictID
  return $! if n == 0
            then Nothing
            else Just (fromIntegral n)

handleError :: CSize -> String -> IO a -> IO a
handleError :: forall a. CSize -> String -> IO a -> IO a
handleError CSize
sizeOrError String
func IO a
act
  | CSize -> Bool
C.isError CSize
sizeOrError
              = String -> String -> IO a
forall a. String -> String -> a
bail String
func (CSize -> String
C.getErrorName CSize
sizeOrError)
  | Bool
otherwise = IO a
act

bail :: String -> String -> a
bail :: forall a. String -> String -> a
bail String
func String
str = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"Codec.Compression.Zstd." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
func String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
str