{-# LANGUAGE ForeignFunctionInterface #-}

-- SPDX-License-Identifier: MPL-2.0

{- | Optional zlib codecs with an explicit hard decompression budget. No shell
process, temporary files, or implicit fallback. A disabled build returns
Unsupported. gzip decoding accepts one member and rejects trailing bytes.
-}
module HMath.Interchange.Compression (Compression (..), compressionAvailable, compress, decompress) where

import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as BI
import Foreign
import Foreign.C.Types
import HMath.Interchange.Array

-- | Wire wrappers around the same DEFLATE coding.
data Compression = RawDeflate | Zlib | Gzip deriving (Compression -> Compression -> Bool
(Compression -> Compression -> Bool)
-> (Compression -> Compression -> Bool) -> Eq Compression
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Compression -> Compression -> Bool
== :: Compression -> Compression -> Bool
$c/= :: Compression -> Compression -> Bool
/= :: Compression -> Compression -> Bool
Eq, Int -> Compression -> ShowS
[Compression] -> ShowS
Compression -> String
(Int -> Compression -> ShowS)
-> (Compression -> String)
-> ([Compression] -> ShowS)
-> Show Compression
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Compression -> ShowS
showsPrec :: Int -> Compression -> ShowS
$cshow :: Compression -> String
show :: Compression -> String
$cshowList :: [Compression] -> ShowS
showList :: [Compression] -> ShowS
Show, Int -> Compression
Compression -> Int
Compression -> [Compression]
Compression -> Compression
Compression -> Compression -> [Compression]
Compression -> Compression -> Compression -> [Compression]
(Compression -> Compression)
-> (Compression -> Compression)
-> (Int -> Compression)
-> (Compression -> Int)
-> (Compression -> [Compression])
-> (Compression -> Compression -> [Compression])
-> (Compression -> Compression -> [Compression])
-> (Compression -> Compression -> Compression -> [Compression])
-> Enum Compression
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: Compression -> Compression
succ :: Compression -> Compression
$cpred :: Compression -> Compression
pred :: Compression -> Compression
$ctoEnum :: Int -> Compression
toEnum :: Int -> Compression
$cfromEnum :: Compression -> Int
fromEnum :: Compression -> Int
$cenumFrom :: Compression -> [Compression]
enumFrom :: Compression -> [Compression]
$cenumFromThen :: Compression -> Compression -> [Compression]
enumFromThen :: Compression -> Compression -> [Compression]
$cenumFromTo :: Compression -> Compression -> [Compression]
enumFromTo :: Compression -> Compression -> [Compression]
$cenumFromThenTo :: Compression -> Compression -> Compression -> [Compression]
enumFromThenTo :: Compression -> Compression -> Compression -> [Compression]
Enum)

foreign import ccall unsafe "hmath_interchange_has_zlib" c_available :: IO CInt
foreign import ccall safe "hmath_interchange_compress"
    c_compress :: CInt -> CInt -> Ptr Word8 -> CSize -> Ptr Word8 -> CSize -> Ptr CSize -> IO CInt

{- | Probe whether the optional native zlib adapter is available, returning IO Bool. False
means compressed operations cannot run; no payload is processed.
-}
compressionAvailable :: IO Bool
compressionAvailable :: IO Bool
compressionAvailable = (CInt -> CInt -> Bool
forall a. Eq a => a -> a -> Bool
/= CInt
0) (CInt -> Bool) -> IO CInt -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO CInt
c_available

{- | Given limits,RawDeflate/Zlib/Gzip wrapper,then bytes, compress within a bounded output
allocation. Return InterchangeError for unavailable provider, invalid limits,
byte-budget excess or compression failure in IO.
-}
compress :: Limits -> Compression -> B.ByteString -> IO (Either InterchangeError B.ByteString)
compress :: Limits
-> Compression
-> ByteString
-> IO (Either InterchangeError ByteString)
compress Limits
lim Compression
wrapper ByteString
input =
    -- zlib's conservative deflate bound plus wrapper slack, capped by the
    -- caller budget. This avoids allocating the entire budget for tiny data.
    let n :: Integer
n = Int -> Integer
forall a. Integral a => a -> Integer
toInteger (ByteString -> Int
B.length ByteString
input)
        bound :: Integer
bound = Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
7) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
8 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
63) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
64 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
64
        capacity :: Int
capacity = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Limits -> Int
maxBytes Limits
lim)) Integer
bound)
     in Limits
-> Bool
-> Compression
-> Int
-> ByteString
-> IO (Either InterchangeError ByteString)
codec Limits
lim Bool
False Compression
wrapper Int
capacity ByteString
input

{- | The output length may be smaller than the budget. Archives additionally
compare it to the declared member size before accepting their payload.
-}
decompress ::
    Limits -> Compression -> Int -> B.ByteString -> IO (Either InterchangeError B.ByteString)
decompress :: Limits
-> Compression
-> Int
-> ByteString
-> IO (Either InterchangeError ByteString)
decompress Limits
lim Compression
wrapper Int
cap = Limits
-> Bool
-> Compression
-> Int
-> ByteString
-> IO (Either InterchangeError ByteString)
codec Limits
lim Bool
True Compression
wrapper Int
cap

codec ::
    Limits -> Bool -> Compression -> Int -> B.ByteString -> IO (Either InterchangeError B.ByteString)
codec :: Limits
-> Bool
-> Compression
-> Int
-> ByteString
-> IO (Either InterchangeError ByteString)
codec Limits
lim Bool
decode Compression
wrapper Int
cap ByteString
input = case Either InterchangeError ()
valid of
    Left InterchangeError
err -> Either InterchangeError ByteString
-> IO (Either InterchangeError ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InterchangeError -> Either InterchangeError ByteString
forall a b. a -> Either a b
Left InterchangeError
err)
    Right () -> do
        available <- IO Bool
compressionAvailable
        if not available
            then pure (Left (Unsupported "built without interchange-zlib"))
            else B.useAsCStringLen input $ \(Ptr CChar
source, Int
len) -> (Ptr CSize -> IO (Either InterchangeError ByteString))
-> IO (Either InterchangeError ByteString)
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr CSize -> IO (Either InterchangeError ByteString))
 -> IO (Either InterchangeError ByteString))
-> (Ptr CSize -> IO (Either InterchangeError ByteString))
-> IO (Either InterchangeError ByteString)
forall a b. (a -> b) -> a -> b
$ \Ptr CSize
written -> do
                -- One spare byte distinguishes an exact zero-size decode from
                -- an output that requires bytes, without a zero-sized buffer.
                let capacity :: Int
capacity = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
cap
                (bytes, status) <- Int
-> (Ptr Word8 -> IO (Int, Int, (CInt, CSize)))
-> IO (ByteString, (CInt, CSize))
forall a.
Int -> (Ptr Word8 -> IO (Int, Int, a)) -> IO (ByteString, a)
BI.createAndTrim' Int
capacity ((Ptr Word8 -> IO (Int, Int, (CInt, CSize)))
 -> IO (ByteString, (CInt, CSize)))
-> (Ptr Word8 -> IO (Int, Int, (CInt, CSize)))
-> IO (ByteString, (CInt, CSize))
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dest -> do
                    Ptr CSize -> CSize -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr CSize
written CSize
0
                    result <-
                        CInt
-> CInt
-> Ptr Word8
-> CSize
-> Ptr Word8
-> CSize
-> Ptr CSize
-> IO CInt
c_compress
                            (if Bool
decode then CInt
1 else CInt
0)
                            (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Compression -> Int
forall a. Enum a => a -> Int
fromEnum Compression
wrapper))
                            (Ptr CChar -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
source)
                            (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
                            Ptr Word8
dest
                            (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
capacity)
                            Ptr CSize
written
                    size <- peek written
                    pure (0, if result == 0 && size <= fromIntegral cap then fromIntegral size else 0, (result, size))
                pure $ case status of
                    (CInt
0, CSize
n) | CSize
n CSize -> CSize -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap -> ByteString -> Either InterchangeError ByteString
forall a b. b -> Either a b
Right ByteString
bytes
                    (-3, CSize
_) -> InterchangeError -> Either InterchangeError ByteString
forall a b. a -> Either a b
Left (String -> InterchangeError
LimitExceeded String
"decompressed/compressed output budget")
                    (CInt
0, CSize
_) -> InterchangeError -> Either InterchangeError ByteString
forall a b. a -> Either a b
Left (String -> InterchangeError
LimitExceeded String
"decompressed output budget")
                    (CInt, CSize)
_ -> InterchangeError -> Either InterchangeError ByteString
forall a b. a -> Either a b
Left (String -> InterchangeError
InvalidData String
"invalid/truncated compressed stream or trailing bytes")
  where
    valid :: Either InterchangeError ()
valid = do
        Limits -> Either InterchangeError ()
checkLimits Limits
lim
        Bool -> InterchangeError -> Either InterchangeError ()
guardI
            ( ByteString -> Int
B.length ByteString
input Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Limits -> Int
maxBytes Limits
lim
                Bool -> Bool -> Bool
&& Int
cap Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
                Bool -> Bool -> Bool
&& Int
cap Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Limits -> Int
maxBytes Limits
lim
                Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
cap Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Word32 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word32
forall a. Bounded a => a
maxBound :: Word32)
                Bool -> Bool -> Bool
&& Int -> Integer
forall a. Integral a => a -> Integer
toInteger (ByteString -> Int
B.length ByteString
input) Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Word32 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word32
forall a. Bounded a => a
maxBound :: Word32)
            )
            (String -> InterchangeError
LimitExceeded String
"compression input/output size")