sak-0.1.2.1: src/Compression.hs

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

import qualified Codec.Compression.Brotli    as Br
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
                      | Default
                      | Custom !Int

data Compression = Lzma
                 | Lzip
                 | BZip
                 | GZip
                 | Zstd
                 | Brotli
                 | Lz4
                 | Z
                 | None
                 deriving (Enum)

ext :: Compression -> String
ext Lzma   = ".xz"
ext Lzip   = ".lz"
ext BZip   = ".bz2"
ext GZip   = ".gz"
ext Zstd   = ".zst"
ext Lz4    = ".lz4"
ext Z      = ".Z"
ext Brotli = ".br"
ext 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"
                   | ".tbr" `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
    | ".br" `isSuffixOf` fp     = Brotli
    | 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 Brotli = Br.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 }

brotliCompression :: Int -> Br.CompressParams
brotliCompression lvl = Br.defaultCompressParams { Br.compressLevel = toEnum lvl }

toInt :: Compression -> CompressionLevel -> Int
toInt Brotli Fastest    = fromEnum (minBound :: Br.CompressionLevel)
toInt Brotli (Custom i) = levelGuard (0,11) i
toInt Brotli _          = fromEnum (maxBound :: Br.CompressionLevel)
toInt Zstd Best         = Zstd.maxCLevel
toInt Zstd Fastest      = 1
toInt Zstd (Custom i)   = levelGuard (1, Zstd.maxCLevel) i
toInt Zstd Default      = 3
toInt Lzip Best         = fromEnum (maxBound :: Lzip.CompressionLevel)
toInt Lzip Fastest      = fromEnum (minBound :: Lzip.CompressionLevel)
toInt Lzip (Custom i)   = levelGuard (0, 9) i
toInt Lzip Default      = 6
toInt Lzma Best         = 9
toInt Lzma Fastest      = 0
toInt Lzma (Custom i)   = levelGuard (0, 9) i
toInt Lzma Default      = 6
toInt BZip Best         = 9
toInt BZip Fastest      = 1
toInt BZip (Custom i)   = i
toInt BZip Default      = 7
toInt GZip Best         = 9
toInt GZip Fastest      = 0
toInt GZip (Custom i)   = i
toInt GZip Default      = 6
toInt Z Best            = 9
toInt Z Fastest         = 0
toInt Z (Custom i)      = i
toInt Z Default         = 6
toInt Lz4 Best          = Lz4.lZ4HCClevelMax
toInt Lz4 Fastest       = 0
toInt Lz4 (Custom i)    = levelGuard (0, Lz4.lZ4HCClevelMax) i
toInt Lz4 Default       = 0 -- 1?
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  lvl _         = Lz4.compressSz (toInt Lz4 lvl)
toCompressor Brotli lvl _       = Br.compressWith (brotliCompression $ toInt Brotli lvl)
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