{-- Originally from the lzma-conduit package
 -- https://hackage.haskell.org/package/lzma-conduit
 --
 -- Adapted & updated
 -}
module Data.Conduit.Lzma2(compress, compressWith, decompress, decompressWith) where

import qualified Codec.Compression.Lzma       as Lzma
import           Control.Monad.IO.Class       (MonadIO (liftIO))
import           Control.Monad.Trans.Resource
import           Data.ByteString              (ByteString)
import qualified Data.ByteString              as B
import           Data.Conduit
import           Data.Conduit.List            (peek)
import           Data.Maybe                   (fromMaybe)
import           Data.Word

prettyRet
  :: Lzma.LzmaRet
  -> String
prettyRet :: LzmaRet -> String
prettyRet LzmaRet
r = case LzmaRet
r of
  LzmaRet
Lzma.LzmaRetOK               -> String
"Operation completed successfully"
  LzmaRet
Lzma.LzmaRetStreamEnd        -> String
"End of stream was reached"
  LzmaRet
Lzma.LzmaRetUnsupportedCheck -> String
"Cannot calculate the integrity check"
  LzmaRet
Lzma.LzmaRetGetCheck         -> String
"Integrity check type is now available"
  LzmaRet
Lzma.LzmaRetMemError         -> String
"Cannot allocate memory"
  LzmaRet
Lzma.LzmaRetMemlimitError    -> String
"Memory usage limit was reached"
  LzmaRet
Lzma.LzmaRetFormatError      -> String
"File format not recognized"
  LzmaRet
Lzma.LzmaRetOptionsError     -> String
"Invalid or unsupported options"
  LzmaRet
Lzma.LzmaRetDataError        -> String
"Data is corrupt"
  LzmaRet
Lzma.LzmaRetBufError         -> String
"No progress is possible"
  LzmaRet
Lzma.LzmaRetProgError        -> String
"Programming error"


-- | Decompress a 'ByteString' from a lzma or xz container stream.
decompress
  :: (MonadThrow m, MonadIO m)
  => Maybe Word64 -- ^ Memory limit, in bytes.
  -> ConduitM ByteString ByteString m ()
decompress :: forall (m :: * -> *).
(MonadThrow m, MonadIO m) =>
Maybe Word64 -> ConduitM ByteString ByteString m ()
decompress Maybe Word64
memlimit =
    DecompressParams -> ConduitM ByteString ByteString m ()
forall (m :: * -> *).
(MonadThrow m, MonadIO m) =>
DecompressParams -> ConduitM ByteString ByteString m ()
decompressWith DecompressParams
Lzma.defaultDecompressParams
                   { Lzma.decompressMemLimit     = fromMaybe maxBound memlimit
                   , Lzma.decompressAutoDecoder  = True
                   , Lzma.decompressConcatenated = True
                   }

decompressWith
  :: (MonadThrow m, MonadIO m)
  => Lzma.DecompressParams
  -> ConduitM ByteString ByteString m ()
decompressWith :: forall (m :: * -> *).
(MonadThrow m, MonadIO m) =>
DecompressParams -> ConduitM ByteString ByteString m ()
decompressWith DecompressParams
parms = do
    Maybe ByteString
c <- ConduitT ByteString ByteString m (Maybe ByteString)
forall (m :: * -> *) a o. Monad m => ConduitT a o m (Maybe a)
peek
    case Maybe ByteString
c of
      Maybe ByteString
Nothing -> IOError -> ConduitM ByteString ByteString m ()
forall e a.
(HasCallStack, Exception e) =>
e -> ConduitT ByteString ByteString m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IOError -> ConduitM ByteString ByteString m ())
-> IOError -> ConduitM ByteString ByteString m ()
forall a b. (a -> b) -> a -> b
$ String -> IOError
userError String
"Data.Conduit.Lzma.decompress: invalid empty input"
      Just ByteString
_  -> IO (DecompressStream IO)
-> ConduitT ByteString ByteString m (DecompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (DecompressParams -> IO (DecompressStream IO)
Lzma.decompressIO DecompressParams
parms) ConduitT ByteString ByteString m (DecompressStream IO)
-> (DecompressStream IO -> ConduitM ByteString ByteString m ())
-> ConduitM ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DecompressStream IO -> ConduitM ByteString ByteString m ()
forall {m :: * -> *}.
(MonadIO m, MonadThrow m) =>
DecompressStream IO -> ConduitT ByteString ByteString m ()
go
  where
    go :: DecompressStream IO -> ConduitT ByteString ByteString m ()
go s :: DecompressStream IO
s@(Lzma.DecompressInputRequired ByteString -> IO (DecompressStream IO)
more) = do
        Maybe ByteString
mx <- ConduitT ByteString ByteString m (Maybe ByteString)
forall (m :: * -> *) a o. Monad m => ConduitT a o m (Maybe a)
await
        case Maybe ByteString
mx of
          Just ByteString
x
            | ByteString -> Bool
B.null ByteString
x  -> DecompressStream IO -> ConduitT ByteString ByteString m ()
go DecompressStream IO
s -- ignore/skip empty bytestring chunks
            | Bool
otherwise -> IO (DecompressStream IO)
-> ConduitT ByteString ByteString m (DecompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ByteString -> IO (DecompressStream IO)
more ByteString
x) ConduitT ByteString ByteString m (DecompressStream IO)
-> (DecompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DecompressStream IO -> ConduitT ByteString ByteString m ()
go
          Maybe ByteString
Nothing       -> IO (DecompressStream IO)
-> ConduitT ByteString ByteString m (DecompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ByteString -> IO (DecompressStream IO)
more ByteString
B.empty) ConduitT ByteString ByteString m (DecompressStream IO)
-> (DecompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DecompressStream IO -> ConduitT ByteString ByteString m ()
go
    go (Lzma.DecompressOutputAvailable ByteString
output IO (DecompressStream IO)
cont) = do
        ByteString -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield ByteString
output
        IO (DecompressStream IO)
-> ConduitT ByteString ByteString m (DecompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (DecompressStream IO)
cont ConduitT ByteString ByteString m (DecompressStream IO)
-> (DecompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DecompressStream IO -> ConduitT ByteString ByteString m ()
go
    go (Lzma.DecompressStreamEnd ByteString
rest)
        | ByteString -> Bool
B.null ByteString
rest = () -> ConduitT ByteString ByteString m ()
forall a. a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        | Bool
otherwise   = ByteString -> ConduitT ByteString ByteString m ()
forall i o (m :: * -> *). i -> ConduitT i o m ()
leftover ByteString
rest
    go (Lzma.DecompressStreamError LzmaRet
err) =
        IOError -> ConduitT ByteString ByteString m ()
forall e a.
(HasCallStack, Exception e) =>
e -> ConduitT ByteString ByteString m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (IOError -> ConduitT ByteString ByteString m ())
-> IOError -> ConduitT ByteString ByteString m ()
forall a b. (a -> b) -> a -> b
$ String -> IOError
userError (String -> IOError) -> String -> IOError
forall a b. (a -> b) -> a -> b
$ String
"Data.Conduit.Lzma.decompress: error: "String -> String -> String
forall a. [a] -> [a] -> [a]
++LzmaRet -> String
prettyRet LzmaRet
err


-- | Compress a 'ByteString' into a xz container stream.
compress
  :: (MonadIO m)
  => Maybe Int -- ^ Compression level from [0..9], defaults to 6.
  -> ConduitM ByteString ByteString m ()
compress :: forall (m :: * -> *).
MonadIO m =>
Maybe Int -> ConduitM ByteString ByteString m ()
compress Maybe Int
level =
   CompressParams -> ConduitM ByteString ByteString m ()
forall (m :: * -> *).
MonadIO m =>
CompressParams -> ConduitM ByteString ByteString m ()
compressWith CompressParams
Lzma.defaultCompressParams { Lzma.compressLevel = level' }
 where
   level' :: CompressionLevel
level' = case Maybe Int
level of
              Maybe Int
Nothing -> CompressionLevel
Lzma.CompressionLevel6
              Just Int
n  -> let
                            n' :: Int
n' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
forall a. Bounded a => a
minBound (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
forall a. Bounded a => a
maxBound Int
n) -- clamp to [minBound..maxBound] range
                        in Int -> CompressionLevel
forall a. Enum a => Int -> a
toEnum Int
n'

compressWith
  :: MonadIO m
  => Lzma.CompressParams
  -> ConduitM ByteString ByteString m ()
compressWith :: forall (m :: * -> *).
MonadIO m =>
CompressParams -> ConduitM ByteString ByteString m ()
compressWith CompressParams
parms =
    IO (CompressStream IO)
-> ConduitT ByteString ByteString m (CompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (CompressParams -> IO (CompressStream IO)
Lzma.compressIO CompressParams
parms) ConduitT ByteString ByteString m (CompressStream IO)
-> (CompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= CompressStream IO -> ConduitT ByteString ByteString m ()
forall {m :: * -> *}.
MonadIO m =>
CompressStream IO -> ConduitT ByteString ByteString m ()
go
  where
    go :: CompressStream IO -> ConduitT ByteString ByteString m ()
go s :: CompressStream IO
s@(Lzma.CompressInputRequired IO (CompressStream IO)
_flush ByteString -> IO (CompressStream IO)
more) = do
        Maybe ByteString
mx <- ConduitT ByteString ByteString m (Maybe ByteString)
forall (m :: * -> *) a o. Monad m => ConduitT a o m (Maybe a)
await
        case Maybe ByteString
mx of
          Just ByteString
x
            | ByteString -> Bool
B.null ByteString
x     -> CompressStream IO -> ConduitT ByteString ByteString m ()
go CompressStream IO
s -- ignore/skip empty bytestring chunks
            | Bool
otherwise    -> IO (CompressStream IO)
-> ConduitT ByteString ByteString m (CompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ByteString -> IO (CompressStream IO)
more ByteString
x) ConduitT ByteString ByteString m (CompressStream IO)
-> (CompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= CompressStream IO -> ConduitT ByteString ByteString m ()
go
          Maybe ByteString
Nothing          -> IO (CompressStream IO)
-> ConduitT ByteString ByteString m (CompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ByteString -> IO (CompressStream IO)
more ByteString
B.empty) ConduitT ByteString ByteString m (CompressStream IO)
-> (CompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= CompressStream IO -> ConduitT ByteString ByteString m ()
go
    go (Lzma.CompressOutputAvailable ByteString
output IO (CompressStream IO)
cont) = do
        ByteString -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield ByteString
output
        IO (CompressStream IO)
-> ConduitT ByteString ByteString m (CompressStream IO)
forall a. IO a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (CompressStream IO)
cont ConduitT ByteString ByteString m (CompressStream IO)
-> (CompressStream IO -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
forall a b.
ConduitT ByteString ByteString m a
-> (a -> ConduitT ByteString ByteString m b)
-> ConduitT ByteString ByteString m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= CompressStream IO -> ConduitT ByteString ByteString m ()
go
    go CompressStream IO
Lzma.CompressStreamEnd = () -> ConduitT ByteString ByteString m ()
forall a. a -> ConduitT ByteString ByteString m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()