module Codec.Bz3 ( Bz3Error (..)
                 , decompressFile
                 , compressFile
                 ) where

import           Codec.Bz3.Binary
import           Codec.Bz3.Foreign
import qualified Control.Monad.ST.Lazy        as LazyST
import qualified Control.Monad.ST.Lazy.Unsafe as LazyST
import           Data.Bifunctor               (bimap)
import           Data.Binary.Get              (runGetOrFail)
import           Data.Binary.Put              (runPut)
import qualified Data.ByteString              as BS
import           Data.ByteString.Internal     as BS
import qualified Data.ByteString.Lazy         as BSL
import qualified Data.ByteString.Unsafe       as BS
import           Data.Int                     (Int32)
import           Foreign.ForeignPtr           (ForeignPtr, castForeignPtr,
                                               mallocForeignPtrBytes,
                                               newForeignPtr, withForeignPtr)
import           Foreign.Marshal.Array        (copyArray)
import           Foreign.Ptr                  (castPtr)

newBz3StForeign :: Int32 -> IO (ForeignPtr Bz3St)
newBz3StForeign :: Int32 -> IO (ForeignPtr Bz3St)
newBz3StForeign Int32
oSz = ForeignPtr () -> ForeignPtr Bz3St
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr (ForeignPtr () -> ForeignPtr Bz3St)
-> IO (ForeignPtr ()) -> IO (ForeignPtr Bz3St)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (FinalizerPtr () -> Ptr () -> IO (ForeignPtr ())
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr ()
bz3Free (Ptr () -> IO (ForeignPtr ()))
-> (Ptr Bz3St -> Ptr ()) -> Ptr Bz3St -> IO (ForeignPtr ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Bz3St -> Ptr ()
forall a b. Ptr a -> Ptr b
castPtr (Ptr Bz3St -> IO (ForeignPtr ()))
-> IO (Ptr Bz3St) -> IO (ForeignPtr ())
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int32 -> IO (Ptr Bz3St)
bz3New Int32
oSz)

decompressFile :: BSL.ByteString -> BSL.ByteString
decompressFile :: ByteString -> ByteString
decompressFile ByteString
contents = (forall s. ST s ByteString) -> ByteString
forall a. (forall s. ST s a) -> a
LazyST.runST ((forall s. ST s ByteString) -> ByteString)
-> (forall s. ST s ByteString) -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> ByteString
BSL.fromChunks ([ByteString] -> ByteString)
-> ST s [ByteString] -> ST s ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
    let Right (ByteString
bs, ByteOffset
_, Word32
bSz) = Get Word32
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Word32)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get Word32
getFileH ByteString
contents
    ForeignPtr Bz3St
st <- IO (ForeignPtr Bz3St) -> ST s (ForeignPtr Bz3St)
forall a s. IO a -> ST s a
LazyST.unsafeIOToST (IO (ForeignPtr Bz3St) -> ST s (ForeignPtr Bz3St))
-> IO (ForeignPtr Bz3St) -> ST s (ForeignPtr Bz3St)
forall a b. (a -> b) -> a -> b
$ Int32 -> IO (ForeignPtr Bz3St)
newBz3StForeign (Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
bSz)
    ForeignPtr Bz3St -> ByteString -> ST s [ByteString]
forall {s}. ForeignPtr Bz3St -> ByteString -> ST s [ByteString]
loop ForeignPtr Bz3St
st ByteString
bs
  where
    loop :: ForeignPtr Bz3St -> ByteString -> ST s [ByteString]
loop ForeignPtr Bz3St
st ByteString
bs | ByteString -> Bool
BSL.null ByteString
bs = [ByteString] -> ST s [ByteString]
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
               | Bool
otherwise = do {(ByteString
dc, ByteString
rest) <- ForeignPtr Bz3St -> ByteString -> ST s (ByteString, ByteString)
forall s.
ForeignPtr Bz3St -> ByteString -> ST s (ByteString, ByteString)
decNext ForeignPtr Bz3St
st ByteString
bs; (ByteString
dcByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:) ([ByteString] -> [ByteString])
-> ST s [ByteString] -> ST s [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr Bz3St -> ByteString -> ST s [ByteString]
loop ForeignPtr Bz3St
st ByteString
rest}

-- will fail if bz3st cannot fit exotic big input
compressFile :: BSL.ByteString -> BSL.ByteString
compressFile :: ByteString -> ByteString
compressFile ByteString
d = (forall s. ST s ByteString) -> ByteString
forall a. (forall s. ST s a) -> a
LazyST.runST ((forall s. ST s ByteString) -> ByteString)
-> (forall s. ST s ByteString) -> ByteString
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Bz3St
st <- IO (ForeignPtr Bz3St) -> ST s (ForeignPtr Bz3St)
forall a s. IO a -> ST s a
LazyST.unsafeIOToST (IO (ForeignPtr Bz3St) -> ST s (ForeignPtr Bz3St))
-> IO (ForeignPtr Bz3St) -> ST s (ForeignPtr Bz3St)
forall a b. (a -> b) -> a -> b
$ Int32 -> IO (ForeignPtr Bz3St)
newBz3StForeign Int32
forall a. Integral a => a
bSz
    (Word32
sz,[ByteString]
bb) <- ForeignPtr Bz3St -> [ByteString] -> ST s (Word32, [ByteString])
forall {s}.
ForeignPtr Bz3St -> [ByteString] -> ST s (Word32, [ByteString])
loop ForeignPtr Bz3St
st (ByteString -> [ByteString]
BSL.toChunks ByteString
d)
    let fileH :: ByteString
fileH = Put -> ByteString
runPut (Word32 -> Put
putFileH Word32
sz)
    ByteString -> ST s ByteString
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
fileHByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>[ByteString] -> ByteString
BSL.fromChunks [ByteString]
bb)
  where
    loop :: ForeignPtr Bz3St -> [ByteString] -> ST s (Word32, [ByteString])
loop ForeignPtr Bz3St
_ []     = (Word32, [ByteString]) -> ST s (Word32, [ByteString])
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word32
0, [])
    loop ForeignPtr Bz3St
s (ByteString
b:[ByteString]
bs) | ByteString -> Int
BS.length ByteString
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
feed = let (ByteString
next,ByteString
b') = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
feed ByteString
b in ForeignPtr Bz3St -> [ByteString] -> ST s (Word32, [ByteString])
loop ForeignPtr Bz3St
s (ByteString
nextByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:ByteString
b'ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:[ByteString]
bs)
    loop ForeignPtr Bz3St
s (ByteString
b:[ByteString]
bs) = do
        (Int32
csz, ByteString
e) <- if Word32
oszWord32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<Word32
64
            then (Int32, ByteString) -> ST s (Int32, ByteString)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int32
0xffffffff, ByteString
b)
            else ForeignPtr Bz3St -> ByteString -> ST s (Int32, ByteString)
forall s.
ForeignPtr Bz3St -> ByteString -> ST s (Int32, ByteString)
encN ForeignPtr Bz3St
s ByteString
b
        let chunk :: ByteString
chunk=ByteString -> ByteString
BSL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Put -> ByteString
runPut (Chunk -> Put
putChunk (Word32 -> Word32 -> Chunk
Chunk (Int32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
csz) Word32
osz))
        (Word32 -> Word32)
-> ([ByteString] -> [ByteString])
-> (Word32, [ByteString])
-> (Word32, [ByteString])
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Word32 -> Word32 -> Word32
forall a. Ord a => a -> a -> a
max Word32
osz) ((ByteString
chunkByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:)([ByteString] -> [ByteString])
-> ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(ByteString
eByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:)) ((Word32, [ByteString]) -> (Word32, [ByteString]))
-> ST s (Word32, [ByteString]) -> ST s (Word32, [ByteString])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr Bz3St -> [ByteString] -> ST s (Word32, [ByteString])
loop ForeignPtr Bz3St
s [ByteString]
bs
      where
        osz :: Word32
osz=Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
b)

    feed :: Int
feed=CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CSize -> CSize
bz3Bound CSize
forall a. Integral a => a
bSz)

    bSz :: Integral a => a
    bSz :: forall a. Integral a => a
bSz=a
16a -> a -> a
forall a. Num a => a -> a -> a
*a
1024a -> a -> a
forall a. Num a => a -> a -> a
*a
1024

encN :: ForeignPtr Bz3St -> BS.ByteString -> LazyST.ST s (Int32, BS.ByteString)
encN :: forall s.
ForeignPtr Bz3St -> ByteString -> ST s (Int32, ByteString)
encN ForeignPtr Bz3St
st ByteString
inp = IO (Int32, ByteString) -> ST s (Int32, ByteString)
forall a s. IO a -> ST s a
LazyST.unsafeIOToST (IO (Int32, ByteString) -> ST s (Int32, ByteString))
-> IO (Int32, ByteString) -> ST s (Int32, ByteString)
forall a b. (a -> b) -> a -> b
$
    ByteString
-> (CStringLen -> IO (Int32, ByteString)) -> IO (Int32, ByteString)
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BS.unsafeUseAsCStringLen ByteString
inp ((CStringLen -> IO (Int32, ByteString)) -> IO (Int32, ByteString))
-> (CStringLen -> IO (Int32, ByteString)) -> IO (Int32, ByteString)
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
d,Int
sz) -> do
        let bufSz :: CSize
bufSz = CSize -> CSize
bz3Bound (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz)
        ForeignPtr CChar
buf <- Int -> IO (ForeignPtr CChar)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes (CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
bufSz)
        Int32
enc <- ForeignPtr CChar -> (Ptr CChar -> IO Int32) -> IO Int32
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr CChar
buf ((Ptr CChar -> IO Int32) -> IO Int32)
-> (Ptr CChar -> IO Int32) -> IO Int32
forall a b. (a -> b) -> a -> b
$ \Ptr CChar
bb -> do
            Ptr CChar -> Ptr CChar -> Int -> IO ()
forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr CChar
bb Ptr CChar
d Int
sz
            ForeignPtr Bz3St -> Ptr UInt8 -> Int32 -> IO Int32
bz3EncodeBlock ForeignPtr Bz3St
st (Ptr CChar -> Ptr UInt8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
bb) (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sz)
        (Int32, ByteString) -> IO (Int32, ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int32
enc, ForeignPtr Word8 -> Int -> ByteString
BS.BS (ForeignPtr CChar -> ForeignPtr Word8
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr CChar
buf) (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
enc))

decNext :: ForeignPtr Bz3St -> BSL.ByteString -> LazyST.ST s (BS.ByteString, BSL.ByteString)
decNext :: forall s.
ForeignPtr Bz3St -> ByteString -> ST s (ByteString, ByteString)
decNext ForeignPtr Bz3St
st ByteString
inp = IO (ByteString, ByteString) -> ST s (ByteString, ByteString)
forall a s. IO a -> ST s a
LazyST.unsafeIOToST (IO (ByteString, ByteString) -> ST s (ByteString, ByteString))
-> IO (ByteString, ByteString) -> ST s (ByteString, ByteString)
forall a b. (a -> b) -> a -> b
$ do
    let (ByteString
next, ByteOffset
off, Chunk Word32
csz Word32
osz) = Either
  (ByteString, ByteOffset, String) (ByteString, ByteOffset, Chunk)
-> (ByteString, ByteOffset, Chunk)
forall {a} {b} {b}. Either (a, b, String) b -> b
y (Get Chunk
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Chunk)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get Chunk
getChunk ByteString
inp)
    if Word32
oszWord32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
>=Word32
64
        then do
            let bufSz :: CSize
bufSz = CSize -> CSize
bz3Bound (Word32 -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
osz)
                csz32 :: Int32
csz32 = Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
csz; osz32 :: Int32
osz32 = Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
osz
            ForeignPtr UInt8
buf <- Int -> IO (ForeignPtr UInt8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes (CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
bufSz)
            let bb :: ByteString
bb=ByteString -> ByteString
BSL.toStrict (ByteOffset -> ByteString -> ByteString
BSL.drop ByteOffset
8 (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteOffset -> ByteString -> ByteString
BSL.take ByteOffset
off ByteString
inp)
            Either Bz3Error Int32
res <- ForeignPtr UInt8
-> (Ptr UInt8 -> IO (Either Bz3Error Int32))
-> IO (Either Bz3Error Int32)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr UInt8
buf ((Ptr UInt8 -> IO (Either Bz3Error Int32))
 -> IO (Either Bz3Error Int32))
-> (Ptr UInt8 -> IO (Either Bz3Error Int32))
-> IO (Either Bz3Error Int32)
forall a b. (a -> b) -> a -> b
$ \Ptr UInt8
b -> do
                ByteString -> (CStringLen -> IO ()) -> IO ()
forall a. ByteString -> (CStringLen -> IO a) -> IO a
BS.unsafeUseAsCStringLen ByteString
bb ((CStringLen -> IO ()) -> IO ()) -> (CStringLen -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
p,Int
isz) -> Ptr UInt8 -> Ptr UInt8 -> Int -> IO ()
forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr UInt8
b (Ptr CChar -> Ptr UInt8
forall a b. Ptr a -> Ptr b
castPtr Ptr CChar
p) Int
isz
                ForeignPtr Bz3St
-> Ptr UInt8
-> CSize
-> Int32
-> Int32
-> IO (Either Bz3Error Int32)
bz3DecodeBlock ForeignPtr Bz3St
st Ptr UInt8
b CSize
bufSz Int32
csz32 Int32
osz32
            case Either Bz3Error Int32
res of
                Right{} -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                Left Bz3Error
e | Bz3Error
eBz3Error -> Bz3Error -> Bool
forall a. Eq a => a -> a -> Bool
==Bz3Error
Bz3Ok -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                       | Bool
otherwise -> String -> IO ()
forall a. HasCallStack => String -> a
error (String -> IO ()) -> IO String -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ForeignPtr Bz3St -> IO String
bz3Strerror ForeignPtr Bz3St
st
            (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ForeignPtr Word8 -> Int -> ByteString
BS.BS (ForeignPtr UInt8 -> ForeignPtr Word8
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr UInt8
buf) (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
osz), ByteString
next)
        else let bb :: ByteString
bb=ByteString -> ByteString
BSL.toStrict (ByteOffset -> ByteString -> ByteString
BSL.drop ByteOffset
8 (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteOffset -> ByteString -> ByteString
BSL.take ByteOffset
off ByteString
inp) 
             in (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
bb, ByteString
next)
  where
    y :: Either (a, b, String) b -> b
y (Right b
x) = b
x; y (Left (a
_, b
_, String
e)) = String -> b
forall a. HasCallStack => String -> a
error String
e