Zstd.hs

module Resource.Compressed.Zstd where

import RIO

import Data.Typeable (typeOf)
import RIO.ByteString qualified as ByteString
import RIO.FilePath (takeExtension)

import Codec.Compression.Zstd qualified as Zstd

-- * Compressed container

newtype Compressed a = Compressed { getCompressed :: a }

compressBytes :: ByteString -> Compressed ByteString
compressBytes = Compressed . Zstd.compress Zstd.maxCLevel

decompressBytes :: Compressed ByteString -> Either CompressedError ByteString
decompressBytes (Compressed bytes) =
  case Zstd.decompress bytes of
    Zstd.Decompress buf ->
      Right buf
    Zstd.Skip ->
      Right mempty
    Zstd.Error str ->
      Left $ ZstdError (fromString str)

instance Typeable a => Show (Compressed a) where
  show (Compressed x) = "Compressed " <> show (typeOf x)

data CompressedError
  = ZstdError Text
  | EmptyFile FilePath
  deriving (Eq, Show)

instance Exception CompressedError

-- * Loading files

fromFileWith :: MonadIO m => (ByteString -> m b) -> (FilePath -> m b) -> FilePath -> m b
fromFileWith withBS withFilePath filePath
  | elem (takeExtension filePath) compressedExts =
      loadCompressed withBS filePath
  | otherwise =
      withFilePath filePath

loadCompressed :: MonadIO m => (ByteString -> m b) -> FilePath -> m b
loadCompressed withBS filePath = do
  bytes <- ByteString.readFile filePath
  case decompressBytes (Compressed bytes) of
    Right buf ->
      if ByteString.null buf then
        throwIO $ EmptyFile filePath
      else
        withBS buf
    Left err ->
      throwIO err

compressedExts :: [FilePath]
compressedExts =
  [ ".zst"
  , ".zstd"
  ]