Streaming.hs

-- 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.

{-# LANGUAGE MultiWayIf #-}

-- |
-- Module      : Codec.Compression.Zstd.Streaming
-- Copyright   : (c) 2016-present, Facebook, Inc. All rights reserved.
--
-- License     : BSD3
-- Maintainer  : bryano@fb.com
-- Stability   : experimental
-- Portability : GHC
--
-- Streaming compression and decompression support for zstd.

module Codec.Compression.Zstd.Streaming
    (
      Result(..)
    , compress
    , decompress
    , maxCLevel
    ) where

import Codec.Compression.Zstd.FFI hiding (compress, decompress)
import Codec.Compression.Zstd.FFI.Types (peekPos)
import qualified Data.ByteString as B
import Data.ByteString.Internal (ByteString(..), mallocByteString)
import Foreign.Marshal.Alloc (finalizerFree, malloc)
import Foreign.C.Types (CSize)
import Foreign.ForeignPtr (ForeignPtr, FinalizerPtr, finalizeForeignPtr, touchForeignPtr, withForeignPtr, newForeignPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import Foreign.Storable (poke)
import Foreign.Ptr (Ptr, plusPtr)
import Data.Word (Word8)

-- | The result of a streaming compression or decompression step.
data Result
  = Produce ByteString (IO Result)
    -- ^ A single frame of transformed data, and an action that when
    -- executed will yield the next step in the streaming operation.
    -- The action is ephemeral; you should discard it as soon as you
    -- use it.
  | Consume (ByteString -> IO Result)
    -- ^ Provide the function with more input for the streaming
    -- operation to continue.  This function is ephemeral. You should
    -- call it exactly once, and discard it immediately after you call
    -- it.
    --
    -- To signal the end of a stream of data, supply an 'B.empty'
    -- input.
  | Error String String
    -- ^ An error has occurred. If an error occurs, the streaming
    -- operation cannot continue.
  | Done ByteString
    -- ^ The streaming operation has ended.  This payload may be
    -- empty. If it is not, it must be written out.
    --
    -- A non-empty payload consists of a frame epilogue, possibly
    -- preceded by any data left over from the final streaming step.

instance Show Result where
    show (Produce bs _) = "Produce " ++ show bs ++ " _"
    show (Consume _)    = "Consume _"
    show (Error n d)    = "Error " ++ show n ++ " " ++ show d
    show (Done bs)      = "Done " ++ show bs

-- | Begin a streaming compression operation.
--
-- The initial result will be either an 'Error' or a 'Consume'.
--
-- This will create and retain a (large)
-- t'Codec.Compression.Zstd.Efficient.CCtx' outside the RTS @-M@
-- heap until 'Error' or 'Done' is returned.
compress :: Int
         -- ^ Compression level. Must be >= 1 and <= 'maxCLevel'.
         -> IO Result
compress level
  | level < 1 || level > maxCLevel =
    return (Error "compress" "unsupported compression level")
  | otherwise =
  streaming
  createCStream
  p_freeCStream
  outSize
  (\cs -> initCStream cs (fromIntegral level))
  compressStream
  finish
 where
  outSize = fromIntegral cstreamOutSize
  finish cfp obfp _ dfp = do
    let cptr = unsafeForeignPtrToPtr cfp
        obuf = unsafeForeignPtrToPtr obfp
    check cfp "endStream" (endStream cptr obuf) $ \leftover -> do
      touchForeignPtr cfp
      touchForeignPtr obfp
      if | leftover <= 0 -> do -- leftover will never be <0, but compiler does not know that
             opos1 <- fromIntegral `fmap` peekPos obuf
             r <- shrink outSize dfp opos1
             -- The Result contract requires that cfp is not referenced anymore.
             -- We finalize the context asap, freeing significant memory
             finalizeForeignPtr cfp
             pure (Done r)
         | leftover > 0 -> do
             opos1 <- fromIntegral `fmap` peekPos obuf
             dfp1 <- mallocByteString (fromIntegral leftover)
             poke obuf (buffer (unsafeForeignPtrToPtr dfp1) leftover)
             touchForeignPtr obfp
             bs <- shrink outSize dfp opos1
             return (Produce bs (finish cfp obfp 0 dfp1))

type ConsumeBlock ctx io = Ptr ctx -> Ptr (Buffer Out)
                         -> Ptr (Buffer In) -> IO CSize

type Finish ctx io = ForeignPtr ctx -> ForeignPtr (Buffer Out)
                   -> Int -> ForeignPtr Word8 -> IO Result

streaming :: IO (Ptr ctx)
          -> FinalizerPtr ctx
          -> Int
          -> (Ptr ctx -> IO CSize)
          -> ConsumeBlock ctx io
          -> Finish ctx io
          -> IO Result
streaming createStream freeStream outSize initStream consumeBlock finish = do
  cx <- checkAlloc "createStream" createStream
  cxfp <- newForeignPtr freeStream cx
  check cxfp "initStream" (initStream cx) $ \_ -> do
    ibfp <- newForeignPtr finalizerFree =<< malloc
    obfp <- newForeignPtr finalizerFree =<< malloc
    dfp <- newOutput obfp
    advanceInput cxfp ibfp obfp 0 dfp
 where
  advanceInput cxfp ibfp obfp opos dfp = do
    let prompt (PS fp off len)
          | len == 0 = finish cxfp obfp opos dfp
          | otherwise = do
              withForeignPtr fp $ \sp0 ->
                withForeignPtr ibfp $ \ibuf ->
                  poke ibuf (buffer (sp0 `plusPtr` off) (fromIntegral len))
              consume cxfp ibfp 0 len obfp 0 dfp fp
    return (Consume prompt)
  newOutput obfp = do
    dfp <- mallocByteString outSize
    withForeignPtr dfp $ \dp ->
      withForeignPtr obfp $ \obuf ->
        poke obuf (buffer dp (fromIntegral outSize))
    return dfp
  consume cxfp ibfp ipos ilen obfp opos dfp fp = do
    if | fromIntegral ipos == ilen -> advanceInput cxfp ibfp obfp opos dfp
       | opos == outSize -> do
           let go = do
                 ndfp <- newOutput obfp
                 consume cxfp ibfp ipos ilen obfp 0 ndfp fp
           return (Produce (PS dfp 0 opos) go)
       | otherwise -> do
           let obuf = unsafeForeignPtrToPtr obfp
               ibuf = unsafeForeignPtrToPtr ibfp
           check cxfp "consumeBlock"
             (withForeignPtr cxfp $ \cptr ->
               consumeBlock cptr obuf ibuf <* touchForeignPtr fp) $ \_ -> do
             opos1 <- fromIntegral `fmap` peekPos obuf
             ipos1 <- peekPos ibuf
             touchForeignPtr obfp
             touchForeignPtr ibfp
             consume cxfp ibfp ipos1 ilen obfp opos1 dfp fp

-- | Begin a streaming decompression operation.
--
-- The initial result will be either an 'Error' or a 'Consume'.
--
-- This will create and retain a (large)
-- t'Codec.Compression.Zstd.Efficient.DCtx' outside the RTS @-M@
-- heap until 'Error' or 'Done' is returned.
decompress :: IO Result
decompress =
  streaming
  createDStream
  p_freeDStream
  outSize
  initDStream
  decompressStream
  finish
 where
  outSize = fromIntegral dstreamOutSize
  finish cfp _obfp opos dfp = do
    -- The Result contract requires that cfp is not referenced anymore,
    -- so we finalize the context asap, freeing significant memory.
    finalizeForeignPtr cfp
    Done `fmap` shrink outSize dfp opos

-- Returns a shrunk view of the buffer retaining the 'ForeignPtr',
-- or a copy when the view could create too much slop (unused but retained memory).
shrink :: Int -> ForeignPtr Word8 -> Int -> IO B.ByteString
shrink capacity dfp opos
  | opos == 0  = return B.empty
  | let unused = capacity - opos
    in unused >= 1024 || unused > capacity `rem` 8
               = return (B.copy (PS dfp 0 opos))
  | otherwise  = return (PS dfp 0 opos)

buffer :: Ptr a -> CSize -> Buffer io
buffer ptr size = Buffer ptr size 0

check :: ForeignPtr ctx -> String -> IO CSize -> (CSize -> IO Result) -> IO Result
check cfp name act onSuccess = do
  ret <- act
  if isError ret
    then do
      -- The Result contract requires that cfp is not referenced anymore,
      -- so we finalize the context asap, freeing significant memory.
      finalizeForeignPtr cfp
      return (Error name (getErrorName ret))
    else onSuccess ret