sak-0.1.0.1: src/Compression.hs

module Compression ( Compression (Lzip)
                   , CompressionLevel (..)
                   , detectCompression
                   , toCompressor
                   , toFileCompressor
                   , toFileDecompressor
                   , uncompressedExt
                   , check
                   -- * Utilities
                   , fileSize
                   ) where

import qualified Codec.Compression.BZip      as BZip
import qualified Codec.Compression.GZip      as GZip
import qualified Codec.Compression.Lzma      as Lzma
import qualified Codec.Compression.Zlib      as Zlib
import qualified Codec.Compression.Zstd.Lazy as Zstd
import qualified Codec.Lz4                   as Lz4
import qualified Codec.Lzip                  as Lzip
import qualified Data.ByteString.Lazy        as BSL
import           Data.List                   (isSuffixOf)
import           System.FilePath             (dropExtension, (-<.>))
import           System.IO                   (IOMode (ReadMode), hFileSize, withFile)

data CompressionLevel = Best
                      | Fastest
                      | Custom !Int

data Compression = Lzma
                 | Lzip
                 | BZip
                 | GZip
                 | Z
                 | Zstd
                 | Lz4
                 | None

uncompressedExt :: FilePath -> FilePath
uncompressedExt fp | ".tlz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".txz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".tlz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".tgz" `isSuffixOf` fp    = fp -<.> ".tar"
                   | ".cpgz" `isSuffixOf` fp   = fp -<.> ".cpio"
                   | ".cpbz2" `isSuffixOf` fp  = fp -<.> ".cpio"
                   | ".xcfgz" `isSuffixOf` fp  = fp -<.> ".xcf"
                   | ".xcfbz2" `isSuffixOf` fp = fp -<.> ".xcf"
                   | ".tbz2" `isSuffixOf` fp   = fp -<.> "tar"
                   | ".tbz" `isSuffixOf` fp    = fp -<.> "tar"
                   | otherwise                 = dropExtension fp

detectCompression :: FilePath -> Compression
detectCompression fp
    | ".xz" `isSuffixOf` fp     = Lzma
    | ".txz" `isSuffixOf` fp    = Lzma
    | ".lz" `isSuffixOf` fp     = Lzip
    | ".tlz" `isSuffixOf` fp    = Lzip
    | ".tbz2" `isSuffixOf` fp   = BZip
    | ".tbz" `isSuffixOf` fp    = BZip
    | ".bz2" `isSuffixOf` fp    = BZip
    | ".gz" `isSuffixOf` fp     = GZip
    | ".tgz" `isSuffixOf` fp    = GZip
    | ".Z" `isSuffixOf` fp      = Z
    | ".zst" `isSuffixOf` fp    = Zstd
    | ".lz4" `isSuffixOf` fp    = Lz4
    | ".xcfgz" `isSuffixOf` fp  = GZip
    | ".xcfbz2" `isSuffixOf` fp = BZip
    | ".cpgz" `isSuffixOf` fp   = GZip
    | ".cpbz2" `isSuffixOf` fp  = BZip
    | otherwise                 = None -- error "Suffix not supported or invalid"

toFileDecompressor :: Compression -> FilePath -> IO BSL.ByteString
toFileDecompressor Lz4 = lz4DecompressFile
toFileDecompressor x   = fmap (toDecompressor x) . BSL.readFile

toDecompressor :: Compression -> BSL.ByteString -> BSL.ByteString
toDecompressor Lzma = Lzma.decompress
toDecompressor Lzip = Lzip.decompress
toDecompressor BZip = BZip.decompress
toDecompressor GZip = GZip.decompress
toDecompressor Z    = Zlib.decompress
toDecompressor Zstd = Zstd.decompress
toDecompressor Lz4  = Lz4.decompress
toDecompressor None = id

check :: Compression -> BSL.ByteString -> IO ()
check = (forceBSL .) . toDecompressor

levelGuard :: (Int, Int) -> Int -> Int
levelGuard (min', max') lvl | lvl < min' || lvl > max' =
    error ("Invalid compression level. Compression must be between " ++ show min' ++ " and " ++ show max')
                            | otherwise = lvl

toFileCompressor :: Compression -> CompressionLevel -> FilePath -> IO BSL.ByteString
toFileCompressor Lzip lvl = Lzip.compressFileLevel (toEnum $ toInt Lzip lvl)
toFileCompressor x lvl    = fmap (toCompressor x lvl Nothing) . BSL.readFile

gzipCompression :: Int -> GZip.CompressParams
gzipCompression lvl = GZip.defaultCompressParams { GZip.compressLevel = GZip.compressionLevel lvl }

zlibCompression :: Int -> Zlib.CompressParams
zlibCompression lvl = Zlib.defaultCompressParams { Zlib.compressLevel = Zlib.compressionLevel lvl }

lzmaCompression :: Int -> Lzma.CompressParams
lzmaCompression lvl = Lzma.defaultCompressParams { Lzma.compressLevel = toEnum lvl }

toInt :: Compression -> CompressionLevel -> Int
toInt Zstd Best       = Zstd.maxCLevel
toInt Zstd Fastest    = 1
toInt Zstd (Custom i) = levelGuard (1, Zstd.maxCLevel) i
toInt Lzip Best       = 9
toInt Lzip Fastest    = 0
toInt Lzip (Custom i) = levelGuard (0, 9) i
toInt Lzma Best       = 9
toInt Lzma Fastest    = 0
toInt Lzma (Custom i) = levelGuard (0, 9) i
toInt BZip Best       = 9
toInt BZip Fastest    = 1
toInt BZip (Custom 6) = 7
toInt BZip (Custom i) = i
toInt GZip Best       = 9
toInt GZip Fastest    = 0
toInt GZip (Custom i) = i
toInt Z Best          = 9
toInt Z Fastest       = 0
toInt Z (Custom i)    = i
toInt None _          = error "Internal error."

toCompressor :: Compression -> CompressionLevel -> Maybe Int -> BSL.ByteString -> BSL.ByteString
toCompressor Lzma lvl _         = Lzma.compressWith (lzmaCompression $ toInt Lzma lvl)
toCompressor Lzip lvl (Just sz) = flip (Lzip.compressWithSz (toEnum $ toInt Lzip lvl)) sz
toCompressor Lzip _ _           = error "Internal error."
toCompressor BZip lvl _         = BZip.compressWith (fromIntegral $ toInt BZip lvl) 30
toCompressor GZip lvl _         = GZip.compressWith (gzipCompression $ toInt GZip lvl)
toCompressor Z    lvl _         = Zlib.compressWith (zlibCompression $ toInt Z lvl)
toCompressor Zstd lvl _         = Zstd.compress (toInt Zstd lvl)
toCompressor Lz4  _ _           = Lz4.compress -- TODO: compression levels upstream
toCompressor None _ _           = id

fileSize :: FilePath -> IO Integer
fileSize fp = withFile fp ReadMode hFileSize

lz4DecompressFile :: FilePath -> IO BSL.ByteString
lz4DecompressFile fp = do
    fSz <- fileSize fp
    let f = if fSz <= 32 * 1024
        then Lz4.decompressBufSz (32 * 1024)
        else Lz4.decompressBufSz (128 * 1024)
    f <$> BSL.readFile fp

forceBSL :: BSL.ByteString -> IO ()
forceBSL bsl = last (BSL.toChunks bsl) `seq` mempty