110 lines
4.0 KiB
Haskell
110 lines
4.0 KiB
Haskell
{-# LANGUAGE CPP #-}
|
|
{-# LANGUAGE BangPatterns #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
module Codec.Archive.Zip.Conduit.Internal
|
|
( osVersion, zipVersion
|
|
, zipError
|
|
, idConduit
|
|
, sizeCRC
|
|
, outputSize
|
|
, inputSize
|
|
, maxBound32
|
|
, compressStream, decompressStream
|
|
) where
|
|
|
|
import Control.Monad.Catch (MonadThrow, throwM)
|
|
import qualified Data.ByteString as BS
|
|
import qualified Data.Conduit as C
|
|
import qualified Data.Conduit.Internal as CI
|
|
import Data.Digest.CRC32 (crc32Update)
|
|
import Data.Word (Word8, Word32, Word64)
|
|
import qualified Codec.Compression.Zlib.Raw as Z
|
|
import qualified Codec.Compression.Zlib.Internal as Z
|
|
import Control.Monad.Primitive
|
|
import Data.Maybe (fromMaybe)
|
|
import Control.Monad.Trans.Class (MonadTrans(lift))
|
|
import qualified Control.Monad.ST.Lazy as STL
|
|
|
|
import Codec.Archive.Zip.Conduit.Types
|
|
|
|
#if MIN_VERSION_conduit(1,3,0)
|
|
#define ConduitM ConduitT
|
|
#define PRE13(x)
|
|
#else
|
|
#define PRE13(x) x
|
|
#endif
|
|
|
|
-- | The version of this zip program, really just rough indicator of compatibility
|
|
zipVersion :: Word8
|
|
zipVersion = 48
|
|
|
|
-- | The OS this implementation tries to be compatible to
|
|
osVersion :: Word8
|
|
osVersion = 0 -- DOS
|
|
|
|
zipError :: MonadThrow m => String -> m a
|
|
zipError = throwM . ZipError
|
|
|
|
idConduit :: Monad m => C.ConduitM a a m ()
|
|
idConduit = C.awaitForever C.yield
|
|
|
|
passthroughFold :: Monad m => (a -> b -> a) -> a -> C.ConduitM b b m a
|
|
passthroughFold f !z = C.await >>= maybe
|
|
(return z)
|
|
(\x -> do
|
|
C.yield x
|
|
passthroughFold f (f z x))
|
|
|
|
sizeCRC :: Monad m => C.ConduitM BS.ByteString BS.ByteString m (Word64, Word32)
|
|
sizeCRC = passthroughFold (\(!l, !c) b -> (l + fromIntegral (BS.length b), crc32Update c b)) (0, 0)
|
|
|
|
sizeC :: Monad m => C.ConduitM BS.ByteString BS.ByteString m Word64
|
|
sizeC = passthroughFold (\l b -> l + fromIntegral (BS.length b)) 0 -- fst <$> sizeCRC
|
|
|
|
outputSize :: Monad m => C.ConduitM i BS.ByteString m () -> C.ConduitM i BS.ByteString m Word64
|
|
outputSize = (C..| sizeC)
|
|
|
|
inputSize :: Monad m => C.ConduitM BS.ByteString o m () -> C.ConduitM BS.ByteString o m Word64
|
|
-- inputSize = fuseUpstream sizeC -- won't work because we need to deal with leftovers properly
|
|
inputSize (CI.ConduitM src) = CI.ConduitM $ \rest -> let
|
|
go n (CI.Done ()) = rest n
|
|
go n (CI.PipeM m) = CI.PipeM $ go n <$> m
|
|
go n (CI.Leftover p b) = CI.Leftover (go (n - fromIntegral (BS.length b)) p) b
|
|
go n (CI.HaveOutput p PRE13(f) o) = CI.HaveOutput (go n p) PRE13(f) o
|
|
go n (CI.NeedInput p q) = CI.NeedInput (\b -> go (n + fromIntegral (BS.length b)) (p b)) (go n . q)
|
|
in go 0 (src CI.Done)
|
|
|
|
maxBound32 :: Integral n => n
|
|
maxBound32 = fromIntegral (maxBound :: Word32)
|
|
|
|
|
|
awaitNonNull :: forall m o. Monad m => C.ConduitT BS.ByteString o m (Maybe BS.ByteString)
|
|
awaitNonNull = do
|
|
next <- C.await
|
|
case next of
|
|
Nothing -> return Nothing
|
|
Just bs
|
|
| BS.null bs -> awaitNonNull
|
|
| otherwise -> return $ Just bs
|
|
|
|
compressStream :: forall m.
|
|
PrimMonad m
|
|
=> Z.CompressParams
|
|
-> C.ConduitT BS.ByteString BS.ByteString m ()
|
|
compressStream params = C.transPipe primToPrim . go $ Z.compressST Z.rawFormat params
|
|
where
|
|
go Z.CompressStreamEnd = return ()
|
|
go (Z.CompressOutputAvailable outBS cont) = C.yield outBS >> lift cont >>= go
|
|
go (Z.CompressInputRequired cont) = awaitNonNull >>= lift . cont . fromMaybe BS.empty >>= go
|
|
|
|
decompressStream :: forall m.
|
|
( MonadThrow m, PrimMonad m )
|
|
=> C.ConduitT BS.ByteString BS.ByteString m ()
|
|
decompressStream = go $ Z.decompressST Z.rawFormat Z.defaultDecompressParams
|
|
where
|
|
go :: Z.DecompressStream (STL.ST (PrimState m)) -> C.ConduitT BS.ByteString BS.ByteString m ()
|
|
go (Z.DecompressStreamEnd unconsumed) = C.leftover unconsumed
|
|
go (Z.DecompressOutputAvailable outBS cont) = C.yield outBS >> lift (primToPrim cont) >>= go
|
|
go (Z.DecompressInputRequired cont) = awaitNonNull >>= lift . primToPrim . cont . fromMaybe BS.empty >>= go
|
|
go (Z.DecompressStreamError err) = throwM err
|