{-# LANGUAGE OverloadedStrings #-}

-- https://github.com/kspalaiologos/bzip3/blob/master/doc/bzip3_format.md
module Codec.Bz3.Binary ( Chunk (..)
                        , getFileH
                        , getFrameH
                        , getChunk
                        , putFileH
                        , putChunk
                        ) where

import           Control.Monad   (unless, when)
import           Data.Binary.Get (Get, getByteString, getWord32le, getWord8,
                                  skip)
import           Data.Binary.Put (Put, putByteString, putWord32le)
import           Data.Bits       (popCount, (.&.))
import           Data.Word       (Word32)

data Chunk = Chunk { Chunk -> Word32
compressedSz, Chunk -> Word32
origSz :: !Word32 }

putFileH :: Word32 -> Put
putFileH :: Word32 -> Put
putFileH Word32
maxSz = do
    ByteString -> Put
putByteString ByteString
"BZ3v1"
    Word32 -> Put
putWord32le Word32
maxSz

-- max block sz
getFileH, getFrameH :: Get Word32
getFileH :: Get Word32
getFileH = do {ByteString
sig <- Int -> Get ByteString
getByteString Int
5; Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString
sigByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
==ByteString
"BZ3v1") (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$ String -> Get ()
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"bad signature"; Get Word32
getWord32le}

getFrameH :: Get Word32
getFrameH = Get Word32
getFileH Get Word32 -> Get Word32 -> Get Word32
forall a b. Get a -> Get b -> Get a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Get Word32
getWord32le

putChunk :: Chunk -> Put
putChunk :: Chunk -> Put
putChunk (Chunk Word32
csz Word32
osz) = Word32 -> Put
putWord32le Word32
csz Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Word32 -> Put
putWord32le Word32
osz

getChunk :: Get Chunk
getChunk :: Get Chunk
getChunk = do
    Word32
csz <- Get Word32
getWord32le; Word32
osz <- Get Word32
getWord32le
    if Word32
oszWord32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<Word32
64
        then Word32 -> Get ()
getSmallBlock Word32
csz
        else Word32 -> Get ()
getRegularBlock Word32
csz
    Chunk -> Get Chunk
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Chunk -> Get Chunk) -> Chunk -> Get Chunk
forall a b. (a -> b) -> a -> b
$ Word32 -> Word32 -> Chunk
Chunk Word32
csz Word32
osz

getSmallBlock :: Word32 -> Get ()
getSmallBlock :: Word32 -> Get ()
getSmallBlock Word32
sz = do
    Int -> Get ()
skip Int
4; Word32
lit <- Get Word32
getWord32le
    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word32
litWord32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
==Word32
0xffffffff) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$ String -> Get ()
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Small block not expected to have bwtIx"
    Int -> Get ()
skip (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
szInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
8)

getRegularBlock :: Word32 -> Get ()
getRegularBlock :: Word32 -> Get ()
getRegularBlock Word32
sz = do
    Int -> Get ()
skip Int
8
    Word8
model <- Get Word8
getWord8
    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word8
model Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x2 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0) (Int -> Get ()
skip Int
4)
    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word8
model Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x4 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0) (Int -> Get ()
skip Int
4)
    Int -> Get ()
skip (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Word8 -> Int
forall a. Bits a => a -> Int
popCount Word8
modelInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
9))