{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
-- | This module implement helper functions to read & write data

-- at bits level.

module Codec.Picture.BitWriter( BoolReader
                              , emptyBoolState
                              , BoolState
                              , byteAlignJpg
                              , getNextBitsLSBFirst
                              , getNextBitsMSBFirst 
                              , getNextBitJpg
                              , getNextIntJpg
                              , setDecodedString
                              , setDecodedStringMSB
                              , setDecodedStringJpg
                              , runBoolReader

                              , BoolWriteStateRef 
                              , newWriteStateRef
                              , finalizeBoolWriter
                              , finalizeBoolWriterGif
                              , writeBits'
                              , writeBitsGif

                              , initBoolState 
                              , initBoolStateJpg
                              , execBoolReader
                              , runBoolReaderWith
                              ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative( (<*>), (<$>) )
#endif

import Data.STRef
import Control.Monad( when )
import Control.Monad.ST( ST )
import qualified Control.Monad.Trans.State.Strict as S
import Data.Int ( Int32 )
import Data.Word( Word8, Word32 )
import Data.Bits( (.&.), (.|.), unsafeShiftR, unsafeShiftL )

import Codec.Picture.VectorByteConversion( blitVector )
import qualified Data.Vector.Storable.Mutable as M
import qualified Data.Vector.Storable as VS
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L


--------------------------------------------------

----            Reader

--------------------------------------------------

-- | Current bit index, current value, string

data BoolState = BoolState {-# UNPACK #-} !Int
                           {-# UNPACK #-} !Word8
                           !B.ByteString

emptyBoolState :: BoolState
emptyBoolState :: BoolState
emptyBoolState = Int -> Word8 -> ByteString -> BoolState
BoolState (-Int
1) Word8
0 ByteString
B.empty

-- | Type used to read bits

type BoolReader s a = S.StateT BoolState (ST s) a

runBoolReader :: BoolReader s a -> ST s a
runBoolReader :: forall s a. BoolReader s a -> ST s a
runBoolReader BoolReader s a
action = BoolReader s a -> BoolState -> ST s a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
S.evalStateT BoolReader s a
action (BoolState -> ST s a) -> BoolState -> ST s a
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState Int
0 Word8
0 ByteString
B.empty

runBoolReaderWith :: BoolState -> BoolReader s a -> ST s (a, BoolState)
runBoolReaderWith :: forall s a. BoolState -> BoolReader s a -> ST s (a, BoolState)
runBoolReaderWith BoolState
st BoolReader s a
action = BoolReader s a -> BoolState -> ST s (a, BoolState)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
S.runStateT BoolReader s a
action BoolState
st

execBoolReader :: BoolState -> BoolReader s a -> ST s BoolState
execBoolReader :: forall s a. BoolState -> BoolReader s a -> ST s BoolState
execBoolReader BoolState
st BoolReader s a
reader = BoolReader s a -> BoolState -> ST s BoolState
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
S.execStateT BoolReader s a
reader BoolState
st

initBoolState :: B.ByteString -> BoolState
initBoolState :: ByteString -> BoolState
initBoolState ByteString
str = case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
str of
     Maybe (Word8, ByteString)
Nothing -> Int -> Word8 -> ByteString -> BoolState
BoolState Int
0 Word8
0 ByteString
B.empty
     Just (Word8
v, ByteString
rest) -> Int -> Word8 -> ByteString -> BoolState
BoolState Int
0 Word8
v ByteString
rest

initBoolStateJpg :: B.ByteString -> BoolState
initBoolStateJpg :: ByteString -> BoolState
initBoolStateJpg ByteString
str = 
   case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
str of
     Maybe (Word8, ByteString)
Nothing -> Int -> Word8 -> ByteString -> BoolState
BoolState Int
0 Word8
0 ByteString
B.empty
     Just (Word8
0xFF, ByteString
rest) -> case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
rest of
            Maybe (Word8, ByteString)
Nothing                  -> Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
0 ByteString
B.empty
            Just (Word8
0x00, ByteString
afterMarker) -> Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
0xFF ByteString
afterMarker
            Just (Word8
_   , ByteString
afterMarker) -> ByteString -> BoolState
initBoolStateJpg ByteString
afterMarker
     Just (Word8
v, ByteString
rest) -> Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
v ByteString
rest

-- | Bitify a list of things to decode.

setDecodedString :: B.ByteString -> BoolReader s ()
setDecodedString :: forall s. ByteString -> BoolReader s ()
setDecodedString ByteString
str = case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
str of
     Maybe (Word8, ByteString)
Nothing        -> BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState      Int
0 Word8
0 ByteString
B.empty
     Just (Word8
v, ByteString
rest) -> BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState      Int
0 Word8
v    ByteString
rest

-- | Drop all bit until the bit of indice 0, usefull to parse restart

-- marker, as they are byte aligned, but Huffman might not.

byteAlignJpg :: BoolReader s ()
byteAlignJpg :: forall s. BoolReader s ()
byteAlignJpg = do
  BoolState idx _ chain <- StateT BoolState (ST s) BoolState
forall (m :: * -> *) s. Monad m => StateT s m s
S.get
  when (idx /= 7) (setDecodedStringJpg chain)

getNextBitJpg :: BoolReader s Bool
{-# INLINE getNextBitJpg #-}
getNextBitJpg :: forall s. BoolReader s Bool
getNextBitJpg = do
    BoolState idx v chain <- StateT BoolState (ST s) BoolState
forall (m :: * -> *) s. Monad m => StateT s m s
S.get
    let val = (Word8
v Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
idx)) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0
    if idx == 0
      then setDecodedStringJpg chain
      else S.put $ BoolState (idx - 1) v chain
    return val

getNextIntJpg :: Int -> BoolReader s Int32
{-# INLINE getNextIntJpg #-}
getNextIntJpg :: forall s. Int -> BoolReader s Int32
getNextIntJpg = Int32 -> Int -> StateT BoolState (ST s) Int32
forall {t} {s}.
(Bits t, Num t) =>
t -> Int -> StateT BoolState (ST s) t
go Int32
0 where
  go :: t -> Int -> StateT BoolState (ST s) t
go !t
acc !Int
0 = t -> StateT BoolState (ST s) t
forall a. a -> StateT BoolState (ST s) a
forall (m :: * -> *) a. Monad m => a -> m a
return t
acc
  go !t
acc !Int
n = do
    BoolState idx v chain <- StateT BoolState (ST s) BoolState
forall (m :: * -> *) s. Monad m => StateT s m s
S.get
    let !leftBits = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx
    if n >= leftBits then do
      setDecodedStringJpg chain
      let !remaining = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftBits
          !mask = (t
1 t -> Int -> t
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
leftBits) t -> t -> t
forall a. Num a => a -> a -> a
- t
1
          !finalV = Word8 -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
v t -> t -> t
forall a. Bits a => a -> a -> a
.&. t
mask
          !theseBits = t
finalV t -> Int -> t
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
remaining
      go (acc .|. theseBits) remaining
    else do
      let !remaining = Int
leftBits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n
          !mask = (t
1 t -> Int -> t
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
n) t -> t -> t
forall a. Num a => a -> a -> a
- t
1
          !finalV = Word8 -> t
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
v t -> Int -> t
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
remaining
      S.put $ BoolState (fromIntegral remaining - 1) v chain
      return $ (finalV .&. mask) .|. acc


setDecodedStringMSB :: B.ByteString -> BoolReader s ()
setDecodedStringMSB :: forall s. ByteString -> BoolReader s ()
setDecodedStringMSB ByteString
str = case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
str of
  Maybe (Word8, ByteString)
Nothing        -> BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState      Int
8 Word8
0 ByteString
B.empty
  Just (Word8
v, ByteString
rest) -> BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState      Int
8 Word8
v    ByteString
rest


{-# INLINE getNextBitsMSBFirst #-}
getNextBitsMSBFirst :: Int -> BoolReader s Word32
getNextBitsMSBFirst :: forall s. Int -> BoolReader s Word32
getNextBitsMSBFirst Int
requested = Word32 -> Int -> BoolReader s Word32
forall s. Word32 -> Int -> BoolReader s Word32
go Word32
0 Int
requested where
  go :: Word32 -> Int -> BoolReader s Word32
  go :: forall s. Word32 -> Int -> BoolReader s Word32
go !Word32
acc !Int
0 = Word32 -> StateT BoolState (ST s) Word32
forall a. a -> StateT BoolState (ST s) a
forall (m :: * -> *) a. Monad m => a -> m a
return Word32
acc
  go !Word32
acc !Int
n = do
    BoolState idx v chain <- StateT BoolState (ST s) BoolState
forall (m :: * -> *) s. Monad m => StateT s m s
S.get
    let !leftBits = Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
idx
    if n >= leftBits then do
      setDecodedStringMSB chain
      let !theseBits = Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
v Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftBits)
      go (acc .|. theseBits) (n - leftBits)
    else do
      let !remaining = Int
leftBits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n
          !mask = (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
remaining) Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
1
      S.put $ BoolState (fromIntegral remaining) (v .&. mask) chain
      return $ (fromIntegral v `unsafeShiftR` remaining) .|. acc

{-# INLINE getNextBitsLSBFirst #-}
getNextBitsLSBFirst :: Int -> BoolReader s Word32
getNextBitsLSBFirst :: forall s. Int -> BoolReader s Word32
getNextBitsLSBFirst Int
count = Word32 -> Int -> StateT BoolState (ST s) Word32
forall {t} {s}.
(Bits t, Num t) =>
t -> Int -> StateT BoolState (ST s) t
aux Word32
0 Int
count
  where aux :: t -> Int -> StateT BoolState (ST s) t
aux t
acc Int
0 = t -> StateT BoolState (ST s) t
forall a. a -> StateT BoolState (ST s) a
forall (m :: * -> *) a. Monad m => a -> m a
return t
acc
        aux t
acc Int
n = do
            bit <- BoolReader s Bool
forall s. BoolReader s Bool
getNextBit
            let nextVal | Bool
bit = t
acc t -> t -> t
forall a. Bits a => a -> a -> a
.|. (t
1 t -> Int -> t
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n))
                        | Bool
otherwise = t
acc
            aux nextVal (n - 1)

{-# INLINE getNextBit #-}
getNextBit :: BoolReader s Bool
getNextBit :: forall s. BoolReader s Bool
getNextBit = do
    BoolState idx v chain <- StateT BoolState (ST s) BoolState
forall (m :: * -> *) s. Monad m => StateT s m s
S.get
    let val = (Word8
v Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. (Word8
1 Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
idx)) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0
    if idx == 7
      then setDecodedString chain
      else S.put $ BoolState (idx + 1) v chain
    return val

-- | Bitify a list of things to decode. Handle Jpeg escape

-- code (0xFF 0x00), thus should be only used in JPEG decoding.

setDecodedStringJpg :: B.ByteString -> BoolReader s ()
setDecodedStringJpg :: forall s. ByteString -> BoolReader s ()
setDecodedStringJpg ByteString
str = case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
str of
     Maybe (Word8, ByteString)
Nothing        -> BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
0 ByteString
B.empty
     Just (Word8
0xFF, ByteString
rest) -> case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
rest of
            Maybe (Word8, ByteString)
Nothing                  -> BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
0 ByteString
B.empty
            Just (Word8
0x00, ByteString
afterMarker) -> -- trace "00" $ 

                BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
0xFF ByteString
afterMarker
            Just (Word8
_   , ByteString
afterMarker) -> ByteString -> BoolReader s ()
forall s. ByteString -> BoolReader s ()
setDecodedStringJpg ByteString
afterMarker
     Just (Word8
v, ByteString
rest) ->
        BoolState -> BoolReader s ()
forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put (BoolState -> BoolReader s ()) -> BoolState -> BoolReader s ()
forall a b. (a -> b) -> a -> b
$ Int -> Word8 -> ByteString -> BoolState
BoolState Int
7 Word8
v ByteString
rest

--------------------------------------------------

----            Writer

--------------------------------------------------

defaultBufferSize :: Int
defaultBufferSize :: Int
defaultBufferSize = Int
256 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024

data BoolWriteStateRef s = BoolWriteStateRef
        { forall s. BoolWriteStateRef s -> STRef s (MVector s Word8)
bwsCurrBuffer   :: STRef s (M.MVector s Word8)
        , forall s. BoolWriteStateRef s -> STRef s [ByteString]
bwsBufferList   :: STRef s [B.ByteString]
        , forall s. BoolWriteStateRef s -> STRef s Int
bwsWrittenWords :: STRef s Int
        , forall s. BoolWriteStateRef s -> STRef s Word8
bwsBitAcc       :: STRef s Word8
        , forall s. BoolWriteStateRef s -> STRef s Int
bwsBitReaded    :: STRef s Int
        }

newWriteStateRef :: ST s (BoolWriteStateRef s)
newWriteStateRef :: forall s. ST s (BoolWriteStateRef s)
newWriteStateRef = do
    origMv <- Int -> ST s (MVector (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
M.new Int
defaultBufferSize
    BoolWriteStateRef <$> newSTRef origMv
                      <*> newSTRef []
                      <*> newSTRef 0
                      <*> newSTRef 0
                      <*> newSTRef 0

finalizeBoolWriter :: BoolWriteStateRef s -> ST s L.ByteString
finalizeBoolWriter :: forall s. BoolWriteStateRef s -> ST s ByteString
finalizeBoolWriter BoolWriteStateRef s
st = do
    BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
flushLeftBits' BoolWriteStateRef s
st
    BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
forceBufferFlushing' BoolWriteStateRef s
st
    [ByteString] -> ByteString
L.fromChunks ([ByteString] -> ByteString)
-> ST s [ByteString] -> ST s ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STRef s [ByteString] -> ST s [ByteString]
forall s a. STRef s a -> ST s a
readSTRef (BoolWriteStateRef s -> STRef s [ByteString]
forall s. BoolWriteStateRef s -> STRef s [ByteString]
bwsBufferList BoolWriteStateRef s
st)

forceBufferFlushing' :: BoolWriteStateRef s -> ST s ()
forceBufferFlushing' :: forall s. BoolWriteStateRef s -> ST s ()
forceBufferFlushing' (BoolWriteStateRef { bwsCurrBuffer :: forall s. BoolWriteStateRef s -> STRef s (MVector s Word8)
bwsCurrBuffer = STRef s (MVector s Word8)
vecRef
                                        , bwsWrittenWords :: forall s. BoolWriteStateRef s -> STRef s Int
bwsWrittenWords = STRef s Int
countRef
                                        , bwsBufferList :: forall s. BoolWriteStateRef s -> STRef s [ByteString]
bwsBufferList = STRef s [ByteString]
lstRef
                                        }) = do
    vec <- STRef s (MVector s Word8) -> ST s (MVector s Word8)
forall s a. STRef s a -> ST s a
readSTRef STRef s (MVector s Word8)
vecRef
    count <- readSTRef countRef
    lst <- readSTRef lstRef

    nmv <- M.new defaultBufferSize
    str <- byteStringFromVector vec count

    writeSTRef vecRef nmv
    writeSTRef lstRef $ lst ++ [str]
    writeSTRef countRef 0

flushCurrentBuffer' :: BoolWriteStateRef s -> ST s ()
flushCurrentBuffer' :: forall s. BoolWriteStateRef s -> ST s ()
flushCurrentBuffer' BoolWriteStateRef s
st = do
    count <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ BoolWriteStateRef s -> STRef s Int
forall s. BoolWriteStateRef s -> STRef s Int
bwsWrittenWords BoolWriteStateRef s
st
    when (count >= defaultBufferSize)
         (forceBufferFlushing' st)

byteStringFromVector :: M.MVector s Word8 -> Int -> ST s B.ByteString
byteStringFromVector :: forall s. MVector s Word8 -> Int -> ST s ByteString
byteStringFromVector MVector s Word8
vec Int
size = do
    frozen <- MVector (PrimState (ST s)) Word8 -> ST s (Vector Word8)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VS.unsafeFreeze MVector s Word8
MVector (PrimState (ST s)) Word8
vec
    return $ blitVector frozen 0 size

setBitCount' :: BoolWriteStateRef s -> Word8 -> Int -> ST s ()
{-# INLINE setBitCount' #-}
setBitCount' :: forall s. BoolWriteStateRef s -> Word8 -> Int -> ST s ()
setBitCount' BoolWriteStateRef s
st Word8
acc Int
count = do
    STRef s Word8 -> Word8 -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef (BoolWriteStateRef s -> STRef s Word8
forall s. BoolWriteStateRef s -> STRef s Word8
bwsBitAcc BoolWriteStateRef s
st) Word8
acc
    STRef s Int -> Int -> ST s ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef (BoolWriteStateRef s -> STRef s Int
forall s. BoolWriteStateRef s -> STRef s Int
bwsBitReaded BoolWriteStateRef s
st) Int
count

resetBitCount' :: BoolWriteStateRef s -> ST s ()
{-# INLINE resetBitCount' #-}
resetBitCount' :: forall s. BoolWriteStateRef s -> ST s ()
resetBitCount' BoolWriteStateRef s
st = BoolWriteStateRef s -> Word8 -> Int -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> Int -> ST s ()
setBitCount' BoolWriteStateRef s
st Word8
0 Int
0

pushByte' :: BoolWriteStateRef s -> Word8 -> ST s ()
{-# INLINE pushByte' #-}
pushByte' :: forall s. BoolWriteStateRef s -> Word8 -> ST s ()
pushByte' BoolWriteStateRef s
st Word8
v = do
    BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
flushCurrentBuffer' BoolWriteStateRef s
st
    idx <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef (BoolWriteStateRef s -> STRef s Int
forall s. BoolWriteStateRef s -> STRef s Int
bwsWrittenWords BoolWriteStateRef s
st)
    vec <- readSTRef (bwsCurrBuffer st)
    M.write vec idx v
    writeSTRef (bwsWrittenWords st) $ idx + 1

flushLeftBits' :: BoolWriteStateRef s -> ST s ()
flushLeftBits' :: forall s. BoolWriteStateRef s -> ST s ()
flushLeftBits' BoolWriteStateRef s
st = do
    currCount <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ BoolWriteStateRef s -> STRef s Int
forall s. BoolWriteStateRef s -> STRef s Int
bwsBitReaded BoolWriteStateRef s
st
    when (currCount > 0) $ do
      currWord <- readSTRef $ bwsBitAcc st
      pushByte' st $ currWord `unsafeShiftL` (8 - currCount)

-- | Append some data bits to a Put monad.

writeBits' :: BoolWriteStateRef s
           -> Word32     -- ^ The real data to be stored. Actual data should be in the LSB

           -> Int        -- ^ Number of bit to write from 1 to 32

           -> ST s ()
{-# INLINE writeBits' #-}
writeBits' :: forall s. BoolWriteStateRef s -> Word32 -> Int -> ST s ()
writeBits' BoolWriteStateRef s
st Word32
d Int
c = do
    currWord <- STRef s Word8 -> ST s Word8
forall s a. STRef s a -> ST s a
readSTRef (STRef s Word8 -> ST s Word8) -> STRef s Word8 -> ST s Word8
forall a b. (a -> b) -> a -> b
$ BoolWriteStateRef s -> STRef s Word8
forall s. BoolWriteStateRef s -> STRef s Word8
bwsBitAcc BoolWriteStateRef s
st
    currCount <- readSTRef $  bwsBitReaded st
    serialize d c currWord currCount
  where dumpByte :: Word8 -> ST s ()
dumpByte Word8
0xFF = BoolWriteStateRef s -> Word8 -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> ST s ()
pushByte' BoolWriteStateRef s
st Word8
0xFF ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> BoolWriteStateRef s -> Word8 -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> ST s ()
pushByte' BoolWriteStateRef s
st Word8
0x00
        dumpByte    Word8
i = BoolWriteStateRef s -> Word8 -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> ST s ()
pushByte' BoolWriteStateRef s
st Word8
i

        serialize :: Word32 -> Int -> Word8 -> Int -> ST s ()
serialize Word32
bitData Int
bitCount Word8
currentWord Int
count
            | Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
8 = do
                     BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
resetBitCount' BoolWriteStateRef s
st
                     Word8 -> ST s ()
dumpByte (Word8 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word8) -> Word8 -> Word8
forall a b. (a -> b) -> a -> b
$ (Word8
currentWord Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
bitCount) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|.
                                                Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
cleanData)

            | Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
8 =
                let newVal :: Word8
newVal = Word8
currentWord Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
bitCount
                in BoolWriteStateRef s -> Word8 -> Int -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> Int -> ST s ()
setBitCount' BoolWriteStateRef s
st (Word8
newVal Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
cleanData) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bitCount

            | Bool
otherwise =
                let leftBitCount :: Int
leftBitCount = Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
count :: Int
                    highPart :: Word32
highPart = Word32
cleanData Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftR` (Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftBitCount) :: Word32
                    prevPart :: Word32
prevPart = Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
currentWord Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
leftBitCount :: Word32

                    nextMask :: Word32
nextMask = (Word32
1 Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftBitCount)) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1 :: Word32
                    newData :: Word32
newData = Word32
cleanData Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
nextMask :: Word32
                    newCount :: Int
newCount = Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftBitCount :: Int

                    toWrite :: Word8
toWrite = Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Word8) -> Word32 -> Word8
forall a b. (a -> b) -> a -> b
$ Word32
prevPart Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32
highPart :: Word8
                in Word8 -> ST s ()
dumpByte Word8
toWrite ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Int -> Word8 -> Int -> ST s ()
serialize Word32
newData Int
newCount Word8
0 Int
0

              where cleanMask :: Word32
cleanMask = (Word32
1 Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
bitCount) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1 :: Word32
                    cleanData :: Word32
cleanData = Word32
bitData Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
cleanMask     :: Word32

-- | Append some data bits to a Put monad.

writeBitsGif :: BoolWriteStateRef s
             -> Word32     -- ^ The real data to be stored. Actual data should be in the LSB

             -> Int        -- ^ Number of bit to write from 1 to 32

             -> ST s ()
{-# INLINE writeBitsGif #-}
writeBitsGif :: forall s. BoolWriteStateRef s -> Word32 -> Int -> ST s ()
writeBitsGif BoolWriteStateRef s
st Word32
d Int
c = do
    currWord <- STRef s Word8 -> ST s Word8
forall s a. STRef s a -> ST s a
readSTRef (STRef s Word8 -> ST s Word8) -> STRef s Word8 -> ST s Word8
forall a b. (a -> b) -> a -> b
$ BoolWriteStateRef s -> STRef s Word8
forall s. BoolWriteStateRef s -> STRef s Word8
bwsBitAcc BoolWriteStateRef s
st
    currCount <- readSTRef $  bwsBitReaded st
    serialize d c currWord currCount
  where dumpByte :: Word8 -> ST s ()
dumpByte = BoolWriteStateRef s -> Word8 -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> ST s ()
pushByte' BoolWriteStateRef s
st

        serialize :: Word32 -> Int -> Word8 -> Int -> ST s ()
serialize Word32
bitData Int
bitCount Word8
currentWord Int
count
            | Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
8 = do
                     BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
resetBitCount' BoolWriteStateRef s
st
                     Word8 -> ST s ()
dumpByte (Word8 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word8) -> Word8 -> Word8
forall a b. (a -> b) -> a -> b
$ Word8
currentWord  Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|.
                                                (Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
cleanData Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
count))

            | Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
8 =
                let newVal :: Word8
newVal = Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
cleanData Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
count
                in BoolWriteStateRef s -> Word8 -> Int -> ST s ()
forall s. BoolWriteStateRef s -> Word8 -> Int -> ST s ()
setBitCount' BoolWriteStateRef s
st (Word8
newVal Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
currentWord) (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bitCount

            | Bool
otherwise =
                let leftBitCount :: Int
leftBitCount = Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
count :: Int
                    newData :: Word32
newData = Word32
cleanData Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
leftBitCount :: Word32
                    newCount :: Int
newCount = Int
bitCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
leftBitCount :: Int
                    toWrite :: Word8
toWrite = Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Word8) -> Word32 -> Word8
forall a b. (a -> b) -> a -> b
$ Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
currentWord 
                                            Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word32
cleanData Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
count) :: Word8
                in Word8 -> ST s ()
dumpByte Word8
toWrite ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Int -> Word8 -> Int -> ST s ()
serialize Word32
newData Int
newCount Word8
0 Int
0

              where cleanMask :: Word32
cleanMask = (Word32
1 Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
bitCount) Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1 :: Word32
                    cleanData :: Word32
cleanData = Word32
bitData Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
cleanMask     :: Word32

finalizeBoolWriterGif :: BoolWriteStateRef s -> ST s L.ByteString
finalizeBoolWriterGif :: forall s. BoolWriteStateRef s -> ST s ByteString
finalizeBoolWriterGif BoolWriteStateRef s
st = do
    BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
flushLeftBitsGif BoolWriteStateRef s
st
    BoolWriteStateRef s -> ST s ()
forall s. BoolWriteStateRef s -> ST s ()
forceBufferFlushing' BoolWriteStateRef s
st
    [ByteString] -> ByteString
L.fromChunks ([ByteString] -> ByteString)
-> ST s [ByteString] -> ST s ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STRef s [ByteString] -> ST s [ByteString]
forall s a. STRef s a -> ST s a
readSTRef (BoolWriteStateRef s -> STRef s [ByteString]
forall s. BoolWriteStateRef s -> STRef s [ByteString]
bwsBufferList BoolWriteStateRef s
st)

flushLeftBitsGif :: BoolWriteStateRef s -> ST s ()
flushLeftBitsGif :: forall s. BoolWriteStateRef s -> ST s ()
flushLeftBitsGif BoolWriteStateRef s
st = do
    currCount <- STRef s Int -> ST s Int
forall s a. STRef s a -> ST s a
readSTRef (STRef s Int -> ST s Int) -> STRef s Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ BoolWriteStateRef s -> STRef s Int
forall s. BoolWriteStateRef s -> STRef s Int
bwsBitReaded BoolWriteStateRef s
st
    when (currCount > 0) $ do
      currWord <- readSTRef $ bwsBitAcc st
      pushByte' st currWord

{-# ANN module "HLint: ignore Reduce duplication" #-}