Add compressFromHandle / decompressFromHandle

Author: Bodigrim

zlib/Codec/Compression/Zlib/Internal.hs

@@ -1,6 +1,8 @@
 {-# LANGUAGE CPP, RankNTypes, BangPatterns #-}
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE Trustworthy #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 -----------------------------------------------------------------------------
 -- |
 -- Copyright   :  (c) 2006-2015 Duncan Coutts
@@ -15,7 +17,9 @@ module Codec.Compression.Zlib.Internal (
 
   -- * Pure interface
   compress,
+  compressFromHandle,
   decompress,
+  decompressFromHandle,
 
   -- * Monadic incremental interface
   -- $incremental-compression
@@ -76,7 +80,7 @@ module Codec.Compression.Zlib.Internal (
 
 import Prelude hiding (length)
 import Control.Monad (when)
-import Control.Exception (Exception, throw, assert)
+import Control.Exception (Exception, throw, assert, throwIO)
 import Control.Monad.ST.Lazy hiding (stToIO)
 import Control.Monad.ST.Strict (stToIO)
 import qualified Control.Monad.ST.Unsafe as Unsafe (unsafeIOToST)
@@ -94,6 +98,8 @@ import GHC.IO (noDuplicate)
 import qualified Codec.Compression.Zlib.Stream as Stream
 import Codec.Compression.Zlib.ByteStringCompat (mkBS, withBS)
 import Codec.Compression.Zlib.Stream (Stream)
+import System.IO (Handle, hIsSeekable, hSeek, SeekMode(..))
+import Data.ByteString.Builder.Extra (defaultChunkSize)
 
 -- | The full set of parameters for compression. The defaults are
 -- 'defaultCompressParams'.
@@ -483,6 +489,26 @@ compress   format params = foldCompressStreamWithInput
 compressST format params = compressStreamST  format params
 compressIO format params = compressStreamIO  format params
 
+compressFromHandle
+  :: forall acc.
+     Stream.Format
+  -> CompressParams
+  -> Handle
+  -> (acc -> S.ByteString -> IO acc) -- TODO add early exit?
+  -> acc
+  -> IO acc
+compressFromHandle format params hndl cons nil = go nil (compressStreamIO format params)
+  where
+    go :: acc -> CompressStream IO -> IO acc
+    go !acc = \case
+      CompressInputRequired next ->
+        S.hGetSome hndl defaultChunkSize >>= next >>= go acc
+      CompressOutputAvailable outchunk next -> do
+        acc' <- acc `cons` outchunk
+        next >>= go acc'
+      CompressStreamEnd ->
+        pure acc
+
 -- | Chunk size must fit into t'CUInt'.
 compressStream :: Stream.Format -> CompressParams -> S.ByteString
                -> Stream (CompressStream Stream)
@@ -617,6 +643,33 @@ decompress   format params = foldDecompressStreamWithInput
 decompressST format params = decompressStreamST  format params
 decompressIO format params = decompressStreamIO  format params
 
+decompressFromHandle
+  :: forall acc.
+     Stream.Format
+  -> DecompressParams
+  -> Handle
+  -> (acc -> S.ByteString -> IO acc) -- TODO add early exit?
+  -> acc
+  -> IO (acc, S.ByteString)
+decompressFromHandle format params hndl cons nil = go nil (decompressStreamIO format params)
+  where
+    go :: acc -> DecompressStream IO -> IO (acc, S.ByteString)
+    go !acc = \case
+      DecompressInputRequired next ->
+        S.hGetSome hndl defaultChunkSize >>= next >>= go acc
+      DecompressOutputAvailable outchunk next -> do
+        acc' <- acc `cons` outchunk
+        next >>= go acc'
+      DecompressStreamEnd unconsumed -> do
+        isHndlSeekable <- hIsSeekable hndl
+        if isHndlSeekable then do
+          hSeek hndl RelativeSeek (negate $ toInteger $ S.length unconsumed)
+          pure (acc, mempty)
+        else
+          pure (acc, unconsumed)
+      DecompressStreamError err ->
+        throwIO err
+
 -- | Chunk size must fit into t'CUInt'.
 decompressStream :: Stream.Format -> DecompressParams
                  -> Bool -> S.ByteString