content: {-# 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 (-1) 0 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 action = S.evalStateT action $ BoolState 0 0 B.empty
  
  runBoolReaderWith :: BoolState -> BoolReader s a -> ST s (a, BoolState)
  runBoolReaderWith st action = S.runStateT action st
  
  execBoolReader :: BoolState -> BoolReader s a -> ST s BoolState
  execBoolReader st reader = S.execStateT reader st
  
  initBoolState :: B.ByteString -> BoolState
  initBoolState str = case B.uncons str of
       Nothing -> BoolState 0 0 B.empty
       Just (v, rest) -> BoolState 0 v rest
  
  initBoolStateJpg :: B.ByteString -> BoolState
  initBoolStateJpg str = 
     case B.uncons str of
       Nothing -> BoolState 0 0 B.empty
       Just (0xFF, rest) -> case B.uncons rest of
              Nothing                  -> BoolState 7 0 B.empty
              Just (0x00, afterMarker) -> BoolState 7 0xFF afterMarker
              Just (_   , afterMarker) -> initBoolStateJpg afterMarker
       Just (v, rest) -> BoolState 7 v rest
  
  -- | Bitify a list of things to decode.
  setDecodedString :: B.ByteString -> BoolReader s ()
  setDecodedString str = case B.uncons str of
       Nothing        -> S.put $ BoolState      0 0 B.empty
       Just (v, rest) -> S.put $ BoolState      0 v    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 = do
    BoolState idx _ chain <- S.get
    when (idx /= 7) (setDecodedStringJpg chain)
  
  getNextBitJpg :: BoolReader s Bool
  {-# INLINE getNextBitJpg #-}
  getNextBitJpg = do
      BoolState idx v chain <- S.get
      let val = (v .&. (1 `unsafeShiftL` idx)) /= 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 = go 0 where
    go !acc !0 = return acc
    go !acc !n = do
      BoolState idx v chain <- S.get
      let !leftBits = 1 + fromIntegral idx
      if n >= leftBits then do
        setDecodedStringJpg chain
        let !remaining = n - leftBits
            !mask = (1 `unsafeShiftL` leftBits) - 1
            !finalV = fromIntegral v .&. mask
            !theseBits = finalV `unsafeShiftL` remaining
        go (acc .|. theseBits) remaining
      else do
        let !remaining = leftBits - n
            !mask = (1 `unsafeShiftL` n) - 1
            !finalV = fromIntegral v `unsafeShiftR` remaining
        S.put $ BoolState (fromIntegral remaining - 1) v chain
        return $ (finalV .&. mask) .|. acc
  
  
  setDecodedStringMSB :: B.ByteString -> BoolReader s ()
  setDecodedStringMSB str = case B.uncons str of
    Nothing        -> S.put $ BoolState      8 0 B.empty
    Just (v, rest) -> S.put $ BoolState      8 v    rest
  
  
  {-# INLINE getNextBitsMSBFirst #-}
  getNextBitsMSBFirst :: Int -> BoolReader s Word32
  getNextBitsMSBFirst requested = go 0 requested where
    go :: Word32 -> Int -> BoolReader s Word32
    go !acc !0 = return acc
    go !acc !n = do
      BoolState idx v chain <- S.get
      let !leftBits = fromIntegral idx
      if n >= leftBits then do
        setDecodedStringMSB chain
        let !theseBits = fromIntegral v `unsafeShiftL` (n - leftBits)
        go (acc .|. theseBits) (n - leftBits)
      else do
        let !remaining = leftBits - n
            !mask = (1 `unsafeShiftL` remaining) - 1
        S.put $ BoolState (fromIntegral remaining) (v .&. mask) chain
        return $ (fromIntegral v `unsafeShiftR` remaining) .|. acc
  
  {-# INLINE getNextBitsLSBFirst #-}
  getNextBitsLSBFirst :: Int -> BoolReader s Word32
  getNextBitsLSBFirst count = aux 0 count
    where aux acc 0 = return acc
          aux acc n = do
              bit <- getNextBit
              let nextVal | bit = acc .|. (1 `unsafeShiftL` (count - n))
                          | otherwise = acc
              aux nextVal (n - 1)
  
  {-# INLINE getNextBit #-}
  getNextBit :: BoolReader s Bool
  getNextBit = do
      BoolState idx v chain <- S.get
      let val = (v .&. (1 `unsafeShiftL` idx)) /= 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 str = case B.uncons str of
       Nothing        -> S.put $ BoolState 7 0 B.empty
       Just (0xFF, rest) -> case B.uncons rest of
              Nothing                  -> S.put $ BoolState 7 0 B.empty
              Just (0x00, afterMarker) -> -- trace "00" $ 
                  S.put $ BoolState 7 0xFF afterMarker
              Just (_   , afterMarker) -> setDecodedStringJpg afterMarker
       Just (v, rest) ->
          S.put $ BoolState 7 v rest
  
  --------------------------------------------------
  ----            Writer
  --------------------------------------------------
  defaultBufferSize :: Int
  defaultBufferSize = 256 * 1024
  
  data BoolWriteStateRef s = BoolWriteStateRef
          { bwsCurrBuffer   :: STRef s (M.MVector s Word8)
          , bwsBufferList   :: STRef s [B.ByteString]
          , bwsWrittenWords :: STRef s Int
          , bwsBitAcc       :: STRef s Word8
          , bwsBitReaded    :: STRef s Int
          }
  
  newWriteStateRef :: ST s (BoolWriteStateRef s)
  newWriteStateRef = do
      origMv <- M.new defaultBufferSize
      BoolWriteStateRef <$> newSTRef origMv
                        <*> newSTRef []
                        <*> newSTRef 0
                        <*> newSTRef 0
                        <*> newSTRef 0
  
  finalizeBoolWriter :: BoolWriteStateRef s -> ST s L.ByteString
  finalizeBoolWriter st = do
      flushLeftBits' st
      forceBufferFlushing' st
      L.fromChunks <$> readSTRef (bwsBufferList st)
  
  forceBufferFlushing' :: BoolWriteStateRef s -> ST s ()
  forceBufferFlushing' (BoolWriteStateRef { bwsCurrBuffer = vecRef
                                          , bwsWrittenWords = countRef
                                          , bwsBufferList = lstRef
                                          }) = do
      vec <- readSTRef 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' st = do
      count <- readSTRef $ bwsWrittenWords st
      when (count >= defaultBufferSize)
           (forceBufferFlushing' st)
  
  byteStringFromVector :: M.MVector s Word8 -> Int -> ST s B.ByteString
  byteStringFromVector vec size = do
      frozen <- VS.unsafeFreeze vec
      return $ blitVector frozen 0 size
  
  setBitCount' :: BoolWriteStateRef s -> Word8 -> Int -> ST s ()
  {-# INLINE setBitCount' #-}
  setBitCount' st acc count = do
      writeSTRef (bwsBitAcc st) acc
      writeSTRef (bwsBitReaded st) count
  
  resetBitCount' :: BoolWriteStateRef s -> ST s ()
  {-# INLINE resetBitCount' #-}
  resetBitCount' st = setBitCount' st 0 0
  
  pushByte' :: BoolWriteStateRef s -> Word8 -> ST s ()
  {-# INLINE pushByte' #-}
  pushByte' st v = do
      flushCurrentBuffer' st
      idx <- readSTRef (bwsWrittenWords st)
      vec <- readSTRef (bwsCurrBuffer st)
      M.write vec idx v
      writeSTRef (bwsWrittenWords st) $ idx + 1
  
  flushLeftBits' :: BoolWriteStateRef s -> ST s ()
  flushLeftBits' st = do
      currCount <- readSTRef $ bwsBitReaded 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' st d c = do
      currWord <- readSTRef $ bwsBitAcc st
      currCount <- readSTRef $  bwsBitReaded st
      serialize d c currWord currCount
    where dumpByte 0xFF = pushByte' st 0xFF >> pushByte' st 0x00
          dumpByte    i = pushByte' st i
  
          serialize bitData bitCount currentWord count
              | bitCount + count == 8 = do
                       resetBitCount' st
                       dumpByte (fromIntegral $ (currentWord `unsafeShiftL` bitCount) .|.
                                                  fromIntegral cleanData)
  
              | bitCount + count < 8 =
                  let newVal = currentWord `unsafeShiftL` bitCount
                  in setBitCount' st (newVal .|. fromIntegral cleanData) $ count + bitCount
  
              | otherwise =
                  let leftBitCount = 8 - count :: Int
                      highPart = cleanData `unsafeShiftR` (bitCount - leftBitCount) :: Word32
                      prevPart = fromIntegral currentWord `unsafeShiftL` leftBitCount :: Word32
  
                      nextMask = (1 `unsafeShiftL` (bitCount - leftBitCount)) - 1 :: Word32
                      newData = cleanData .&. nextMask :: Word32
                      newCount = bitCount - leftBitCount :: Int
  
                      toWrite = fromIntegral $ prevPart .|. highPart :: Word8
                  in dumpByte toWrite >> serialize newData newCount 0 0
  
                where cleanMask = (1 `unsafeShiftL` bitCount) - 1 :: Word32
                      cleanData = bitData .&. 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 st d c = do
      currWord <- readSTRef $ bwsBitAcc st
      currCount <- readSTRef $  bwsBitReaded st
      serialize d c currWord currCount
    where dumpByte = pushByte' st
  
          serialize bitData bitCount currentWord count
              | bitCount + count == 8 = do
                       resetBitCount' st
                       dumpByte (fromIntegral $ currentWord  .|.
                                                  (fromIntegral cleanData `unsafeShiftL` count))
  
              | bitCount + count < 8 =
                  let newVal = fromIntegral cleanData `unsafeShiftL` count
                  in setBitCount' st (newVal .|. currentWord) $ count + bitCount
  
              | otherwise =
                  let leftBitCount = 8 - count :: Int
                      newData = cleanData `unsafeShiftR` leftBitCount :: Word32
                      newCount = bitCount - leftBitCount :: Int
                      toWrite = fromIntegral $ fromIntegral currentWord 
                                              .|. (cleanData `unsafeShiftL` count) :: Word8
                  in dumpByte toWrite >> serialize newData newCount 0 0
  
                where cleanMask = (1 `unsafeShiftL` bitCount) - 1 :: Word32
                      cleanData = bitData .&. cleanMask     :: Word32
  
  finalizeBoolWriterGif :: BoolWriteStateRef s -> ST s L.ByteString
  finalizeBoolWriterGif st = do
      flushLeftBitsGif st
      forceBufferFlushing' st
      L.fromChunks <$> readSTRef (bwsBufferList st)
  
  flushLeftBitsGif :: BoolWriteStateRef s -> ST s ()
  flushLeftBitsGif st = do
      currCount <- readSTRef $ bwsBitReaded st
      when (currCount > 0) $ do
        currWord <- readSTRef $ bwsBitAcc st
        pushByte' st currWord
  
  {-# ANN module "HLint: ignore Reduce duplication" #-}
  
  