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 { forall a. Compressed a -> a
getCompressed :: a }

compressBytes :: ByteString -> Compressed ByteString
compressBytes :: ByteString -> Compressed ByteString
compressBytes = ByteString -> Compressed ByteString
forall a. a -> Compressed a
Compressed (ByteString -> Compressed ByteString)
-> (ByteString -> ByteString)
-> ByteString
-> Compressed ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
Zstd.compress Int
Zstd.maxCLevel

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

instance Typeable a => Show (Compressed a) where
  show :: Compressed a -> String
show (Compressed a
x) = String
"Compressed " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> TypeRep -> String
forall a. Show a => a -> String
show (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
x)

data CompressedError
  = ZstdError Text
  | EmptyFile FilePath
  deriving (CompressedError -> CompressedError -> Bool
(CompressedError -> CompressedError -> Bool)
-> (CompressedError -> CompressedError -> Bool)
-> Eq CompressedError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CompressedError -> CompressedError -> Bool
$c/= :: CompressedError -> CompressedError -> Bool
== :: CompressedError -> CompressedError -> Bool
$c== :: CompressedError -> CompressedError -> Bool
Eq, Int -> CompressedError -> ShowS
[CompressedError] -> ShowS
CompressedError -> String
(Int -> CompressedError -> ShowS)
-> (CompressedError -> String)
-> ([CompressedError] -> ShowS)
-> Show CompressedError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CompressedError] -> ShowS
$cshowList :: [CompressedError] -> ShowS
show :: CompressedError -> String
$cshow :: CompressedError -> String
showsPrec :: Int -> CompressedError -> ShowS
$cshowsPrec :: Int -> CompressedError -> ShowS
Show)

instance Exception CompressedError

-- * Loading files

fromFileWith :: MonadIO m => (ByteString -> m b) -> (FilePath -> m b) -> FilePath -> m b
fromFileWith :: forall (m :: * -> *) b.
MonadIO m =>
(ByteString -> m b) -> (String -> m b) -> String -> m b
fromFileWith ByteString -> m b
withBS String -> m b
withFilePath String
filePath
  | String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem (ShowS
takeExtension String
filePath) [String]
compressedExts =
      (ByteString -> m b) -> String -> m b
forall (m :: * -> *) b.
MonadIO m =>
(ByteString -> m b) -> String -> m b
loadCompressed ByteString -> m b
withBS String
filePath
  | Bool
otherwise =
      String -> m b
withFilePath String
filePath

loadCompressed :: MonadIO m => (ByteString -> m b) -> FilePath -> m b
loadCompressed :: forall (m :: * -> *) b.
MonadIO m =>
(ByteString -> m b) -> String -> m b
loadCompressed ByteString -> m b
withBS String
filePath = do
  ByteString
bytes <- String -> m ByteString
forall (m :: * -> *). MonadIO m => String -> m ByteString
ByteString.readFile String
filePath
  case Compressed ByteString -> Either CompressedError ByteString
decompressBytes (ByteString -> Compressed ByteString
forall a. a -> Compressed a
Compressed ByteString
bytes) of
    Right ByteString
buf ->
      if ByteString -> Bool
ByteString.null ByteString
buf then
        CompressedError -> m b
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (CompressedError -> m b) -> CompressedError -> m b
forall a b. (a -> b) -> a -> b
$ String -> CompressedError
EmptyFile String
filePath
      else
        ByteString -> m b
withBS ByteString
buf
    Left CompressedError
err ->
      CompressedError -> m b
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO CompressedError
err

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