content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE NoMonomorphismRestriction #-}
  
  -- | Strict Decoder Primitives
  module Flat.Decoder.Prim
    ( dBool,
      dWord8,
      dBE8,
      dBE16,
      dBE32,
      dBE64,
      dBEBits8,
      dBEBits16,
      dBEBits32,
      dBEBits64,
      dropBits,
      dFloat,
      dDouble,
      getChunksInfo,
      dByteString_,
      dLazyByteString_,
      dByteArray_,
      ConsState (..),
      consOpen,
      consClose,
      consBool,
      consBits,
      sizeOf,
      binOf,
    )
  where
  
  import Control.Monad (when)
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy as L
  import Data.FloatCast (wordToDouble, wordToFloat)
  import Data.Word (Word16, Word32, Word64, Word8)
  import Flat.Decoder.Types
    ( Get (Get, runGet),
      GetResult (..),
      S (..),
      badEncoding,
      badOp,
      notEnoughSpace,
    )
  import Flat.Endian (toBE16, toBE32, toBE64)
  import Flat.Memory
    ( ByteArray,
      chunksToByteArray,
      chunksToByteString,
      minusPtr,
      peekByteString,
    )
  import Foreign
    ( Bits (unsafeShiftL, unsafeShiftR, (.&.), (.|.)),
      FiniteBits (finiteBitSize),
      Ptr,
      Storable (peek),
      castPtr,
      plusPtr,
      ptrToIntPtr,
    )
  
  -- $setup
  -- >>> :set -XBinaryLiterals
  -- >>> import Data.Word
  -- >>> import Data.Int
  -- >>> import Flat.Run
  -- >>> import Flat.Bits
  -- >>> import Text.PrettyPrint.HughesPJClass (Pretty (pPrint))
  
  -- | A special state, optimised for constructor decoding.
  --
  -- It consists of:
  --
  -- * The bits to parse, the top bit being the first to parse (could use a Word16 instead, no difference in performance)
  --
  -- * The number of decoded bits
  --
  -- Supports up to 512 constructors (9 bits).
  data ConsState
    = ConsState {-# UNPACK #-} !Word !Int
  
  -- | Switch to constructor decoding
  --  {-# INLINE consOpen  #-}
  consOpen :: Get ConsState
  consOpen = Get $ \endPtr s -> do
    let u = usedBits s
    let d = ptrToIntPtr endPtr - ptrToIntPtr (currPtr s)
    w <-
      if d > 1
        then do
          -- two different bytes
          w16 :: Word16 <- toBE16 <$> peek (castPtr $ currPtr s)
          return $ fromIntegral w16 `unsafeShiftL` (u + (wordSize - 16))
        else
          if d == 1
            then do
              -- single last byte left
              w8 :: Word8 <- peek (currPtr s)
              return $ fromIntegral w8 `unsafeShiftL` (u + (wordSize - 8))
            else notEnoughSpace endPtr s
    return $ GetResult s (ConsState w 0)
  
  -- | Switch back to normal decoding
  --  {-# NOINLINE consClose  #-}
  consClose :: Int -> Get ()
  consClose n = Get $ \endPtr s -> do
    let u' = n + usedBits s
    if u' < 8
      then return $ GetResult (s {usedBits = u'}) ()
      else
        if currPtr s >= endPtr
          then notEnoughSpace endPtr s
          else return $ GetResult (s {currPtr = currPtr s `plusPtr` 1, usedBits = u' - 8}) ()
  
  {- ensureBits endPtr s n = when ((endPtr `minusPtr` currPtr s) * 8 - usedBits s < n) $ notEnoughSpace endPtr s
  dropBits8 s n =
    let u' = n+usedBits s
    in if u' < 8
        then s {usedBits=u'}
        else s {currPtr=currPtr s `plusPtr` 1,usedBits=u'-8}
  -}
  
  -- ensureBits endPtr s n
  -- return $ GetResult (dropBits8 s n) ()
  
  -- | Decode a single bit
  consBool :: ConsState -> (ConsState, Bool)
  consBool cs = (0 /=) <$> consBits cs 1
  
  -- consBool (ConsState w usedBits) = (ConsState (w `unsafeShiftL` 1) (1+usedBits),0 /= 32768 .&. w)
  
  -- | Decode from 1 to 3 bits
  --
  --  It could read more bits that are available, but it doesn't matter, errors will be checked in consClose.
  consBits :: ConsState -> Int -> (ConsState, Word)
  consBits cs 3 = consBits_ cs 3 7
  consBits cs 2 = consBits_ cs 2 3
  consBits cs 1 = consBits_ cs 1 1
  consBits _ _ = error "unsupported"
  
  consBits_ :: ConsState -> Int -> Word -> (ConsState, Word)
  
  -- Different decoding primitives
  -- All with equivalent performance
  -- #define CONS_ROT
  -- #define CONS_SHL
  #define CONS_STA
  
  #ifdef CONS_ROT
  consBits_ (ConsState w usedBits) numBits mask =
    let usedBits' = numBits+usedBits
        w' = w `rotateL` numBits -- compiles to an or+shiftl+shiftr
    in (ConsState w' usedBits',w' .&. mask)
  #endif
  
  #ifdef CONS_SHL
  consBits_ (ConsState w usedBits) numBits mask =
    let usedBits' = numBits+usedBits
        w' = w `unsafeShiftL` numBits
    in (ConsState w' usedBits', (w `unsafeShiftR` (wordSize - numBits)) .&. mask)
  #endif
  
  #ifdef CONS_STA
  consBits_ (ConsState w usedBits) numBits mask =
    let usedBits' = numBits+usedBits
    in (ConsState w usedBits', (w `unsafeShiftR` (wordSize - usedBits')) .&. mask)
  #endif
  
  wordSize :: Int
  wordSize = finiteBitSize (0 :: Word)
  
  {-# INLINE ensureBits #-}
  
  -- | Ensure that the specified number of bits is available
  ensureBits :: Ptr Word8 -> S -> Int -> IO ()
  ensureBits endPtr s n = when ((endPtr `minusPtr` currPtr s) * 8 - usedBits s < n) $ notEnoughSpace endPtr s
  
  {-# INLINE dropBits #-}
  
  -- | Drop the specified number of bits
  dropBits :: Int -> Get ()
  dropBits n
    | n > 0 = Get $ \endPtr s -> do
        ensureBits endPtr s n
        return $ GetResult (dropBits_ s n) ()
    | n == 0 = return ()
    | otherwise = error $ unwords ["dropBits", show n]
  
  {-# INLINE dropBits_ #-}
  dropBits_ :: S -> Int -> S
  dropBits_ s n =
    let (bytes, bits) = (n + usedBits s) `divMod` 8
     in -- let
        --   n' = n+usedBits s
        --   bytes = n' `unsafeShiftR` 3
        --   bits = n' .|. 7
        S {currPtr = currPtr s `plusPtr` bytes, usedBits = bits}
  
  {-# INLINE dBool #-}
  
  -- Inlining dBool massively increases compilation time but decreases run time by a third
  -- TODO: test dBool inlining for ghc >= 8.8.4
  
  -- | Decode a boolean
  dBool :: Get Bool
  dBool = Get $ \endPtr s ->
    if currPtr s >= endPtr
      then notEnoughSpace endPtr s
      else do
        !w <- peek (currPtr s)
        let !b = 0 /= (w .&. (128 `unsafeShiftR` usedBits s))
        let !s' =
              if usedBits s == 7
                then s {currPtr = currPtr s `plusPtr` 1, usedBits = 0}
                else s {usedBits = usedBits s + 1}
        return $ GetResult s' b
  
  {-# INLINE dBEBits8 #-}
  
  -- | Return the n most significant bits (up to maximum of 8)
  --
  -- The bits are returned right shifted:
  --
  -- >>> unflatWith (dBEBits8 3) [0b11100001::Word8] == Right 0b00000111
  -- True
  --
  -- >>> unflatWith (dBEBits8 9) [0b11100001::Word8,0b11111111]
  -- Left (BadOp "read8: cannot read 9 bits")
  dBEBits8 :: Int -> Get Word8
  dBEBits8 n = Get $ \endPtr s -> do
    ensureBits endPtr s n
    take8 s n
  
  {-# INLINE dBEBits16 #-}
  
  -- | Return the n most significant bits (up to maximum of 16)
  --
  -- The bits are returned right shifted:
  --
  -- >>> pPrint . asBits <$> unflatWith (dBEBits16 11) [0b10110111::Word8,0b11100001]
  -- Right 00000101 10111111
  --
  -- If more than 16 bits are requested, only the last 16 are returned:
  --
  -- >>> pPrint . asBits <$> unflatWith (dBEBits16 19) [0b00000000::Word8,0b11111111,0b11100001]
  -- Right 00000111 11111111
  dBEBits16 :: Int -> Get Word16
  dBEBits16 n = Get $ \endPtr s -> do
    ensureBits endPtr s n
    takeN n s
  
  {-# INLINE dBEBits32 #-}
  
  -- | Return the n most significant bits (up to maximum of 32)
  --  The bits are returned right shifted.
  dBEBits32 :: Int -> Get Word32
  dBEBits32 n = Get $ \endPtr s -> do
    ensureBits endPtr s n
    takeN n s
  
  {-# INLINE dBEBits64 #-}
  
  -- | Return the n most significant bits (up to maximum of 64)
  --  The bits are returned right shifted.
  dBEBits64 :: Int -> Get Word64
  dBEBits64 n = Get $ \endPtr s -> do
    ensureBits endPtr s n
    takeN n s
  
  -- {-# INLINE take8 #-}
  -- take8 :: Int -> S -> IO (GetResult Word8)
  -- take8 n s
  --   | n == 0 = return $ GetResult s 0
  
  --   -- all bits in the same byte
  --   | n <= 8 - usedBits s = do
  --       w <- peek (currPtr s)
  --       let (bytes,bits) = (n+usedBits s) `divMod` 8
  --       return $ GetResult (S {currPtr=currPtr s `plusPtr` bytes,usedBits=bits}) ((w `unsafeShiftL` usedBits s) `unsafeShiftR` (8 - n))
  
  --   -- two different bytes
  --   | n <= 8 = do
  --       w::Word16 <- toBE16 <$> peek (castPtr $ currPtr s)
  --       return $ GetResult (S {currPtr=currPtr s `plusPtr` 1,usedBits=(usedBits s + n) `mod` 8}) (fromIntegral $ (w `unsafeShiftL` usedBits s) `unsafeShiftR` (16 - n))
  
  --   | otherwise = error $ unwords ["take8: cannot take",show n,"bits"]
  
  {-# INLINE take8 #-}
  take8 :: S -> Int -> IO (GetResult Word8)
  -- take8 s n = GetResult (dropBits_ s n) <$> read8 s n
  take8 s n = GetResult (dropBits8 s n) <$> read8 s n
    where
      -- {-# INLINE read8 #-}
      read8 :: S -> Int -> IO Word8
      read8 s n
        | n >= 0 && n <= 8 =
            if n <= 8 - usedBits s
              then do
                -- all bits in the same byte
                w <- peek (currPtr s)
                return $ (w `unsafeShiftL` usedBits s) `unsafeShiftR` (8 - n)
              else do
                -- two different bytes
                w :: Word16 <- toBE16 <$> peek (castPtr $ currPtr s)
                return $ fromIntegral $ (w `unsafeShiftL` usedBits s) `unsafeShiftR` (16 - n)
        | otherwise = badOp $ unwords ["read8: cannot read", show n, "bits"]
      -- {-# INLINE dropBits8 #-}
      -- -- Assume n <= 8
      dropBits8 :: S -> Int -> S
      dropBits8 s n =
        let u' = n + usedBits s
         in if u' < 8
              then s {usedBits = u'}
              else s {currPtr = currPtr s `plusPtr` 1, usedBits = u' - 8}
  
  {-# INLINE takeN #-}
  takeN :: (Num a, Bits a) => Int -> S -> IO (GetResult a)
  takeN n s = read s 0 (n - (n `min` 8)) n
    where
      read s r sh n
        | n <= 0 = return $ GetResult s r
        | otherwise = do
            let m = n `min` 8
            GetResult s' b <- take8 s m
            read s' (r .|. (fromIntegral b `unsafeShiftL` sh)) ((sh - 8) `max` 0) (n - 8)
  
  -- takeN n = Get $ \endPtr s -> do
  --   ensureBits endPtr s n
  --   let (bytes,bits) = (n+usedBits s) `divMod` 8
  --   r <- case bytes of
  --     0 -> do
  --       w <- peek (currPtr s)
  --       return . fromIntegral $ ((w `unsafeShiftL` usedBits s) `unsafeShiftR` (8 - n))
  --     1 -> do
  --       w::Word16 <- toBE16 <$> peek (castPtr $ currPtr s)
  --       return $ fromIntegral $ (w `unsafeShiftL` usedBits s) `unsafeShiftR` (16 - n)
  --     2 -> do
  --       let r = 0
  --       w1 <- fromIntegral <$> r8 s
  --       w2 <- fromIntegral <$> r16 s
  --       w1
  --   return $ GetResult (S {currPtr=currPtr s `plusPtr` bytes,usedBits=bits}) r
  
  -- r8 s = peek (currPtr s)
  -- r16 s = toBE16 <$> peek (castPtr $ currPtr s)
  
  -- | Return the 8 most significant bits (same as dBE8)
  dWord8 :: Get Word8
  dWord8 = dBE8
  
  {-# INLINE dBE8 #-}
  
  -- | Return the 8 most significant bits
  dBE8 :: Get Word8
  dBE8 = Get $ \endPtr s -> do
    ensureBits endPtr s 8
    !w1 <- peek (currPtr s)
    !w <-
      if usedBits s == 0
        then return w1
        else do
          !w2 <- peek (currPtr s `plusPtr` 1)
          return $ (w1 `unsafeShiftL` usedBits s) .|. (w2 `unsafeShiftR` (8 - usedBits s))
    return $ GetResult (s {currPtr = currPtr s `plusPtr` 1}) w
  
  {-# INLINE dBE16 #-}
  
  -- | Return the 16 most significant bits
  dBE16 :: Get Word16
  dBE16 = Get $ \endPtr s -> do
    ensureBits endPtr s 16
    !w1 <- toBE16 <$> peek (castPtr $ currPtr s)
    !w <-
      if usedBits s == 0
        then return w1
        else do
          !(w2 :: Word8) <- peek (currPtr s `plusPtr` 2)
          return $ w1 `unsafeShiftL` usedBits s .|. fromIntegral (w2 `unsafeShiftR` (8 - usedBits s))
    return $ GetResult (s {currPtr = currPtr s `plusPtr` 2}) w
  
  {-# INLINE dBE32 #-}
  
  -- | Return the 32 most significant bits
  dBE32 :: Get Word32
  dBE32 = Get $ \endPtr s -> do
    ensureBits endPtr s 32
    !w1 <- toBE32 <$> peek (castPtr $ currPtr s)
    !w <-
      if usedBits s == 0
        then return w1
        else do
          !(w2 :: Word8) <- peek (currPtr s `plusPtr` 4)
          return $ w1 `unsafeShiftL` usedBits s .|. fromIntegral (w2 `unsafeShiftR` (8 - usedBits s))
    return $ GetResult (s {currPtr = currPtr s `plusPtr` 4}) w
  
  {-# INLINE dBE64 #-}
  
  -- | Return the 64 most significant bits
  dBE64 :: Get Word64
  dBE64 = Get $ \endPtr s -> do
    ensureBits endPtr s 64
    -- !w1 <- toBE64 <$> peek (castPtr $ currPtr s)
    !w1 <- toBE64 <$> peek64 (castPtr $ currPtr s)
    !w <-
      if usedBits s == 0
        then return w1
        else do
          !(w2 :: Word8) <- peek (currPtr s `plusPtr` 8)
          return $ w1 `unsafeShiftL` usedBits s .|. fromIntegral (w2 `unsafeShiftR` (8 - usedBits s))
    return $ GetResult (s {currPtr = currPtr s `plusPtr` 8}) w
    where
      -- {-# INLINE peek64 #-}
      peek64 :: Ptr Word64 -> IO Word64
      peek64 = peek
  
  -- peek64 ptr = fix64 <$> peek ptr
  
  {-# INLINE dFloat #-}
  
  -- | Decode a Float
  dFloat :: Get Float
  dFloat = wordToFloat <$> dBE32
  
  {-# INLINE dDouble #-}
  
  -- | Decode a Double
  dDouble :: Get Double
  dDouble = wordToDouble <$> dBE64
  
  -- | Decode a Lazy ByteString
  dLazyByteString_ :: Get L.ByteString
  dLazyByteString_ = L.fromStrict <$> dByteString_
  
  -- | Decode a ByteString
  dByteString_ :: Get B.ByteString
  dByteString_ = chunksToByteString <$> getChunksInfo
  
  -- | Decode a ByteArray and its length
  dByteArray_ :: Get (ByteArray, Int)
  dByteArray_ = chunksToByteArray <$> getChunksInfo
  
  -- | Decode an Array (a list of chunks up to 255 bytes long) returning the pointer to the first data byte and a list of chunk sizes
  getChunksInfo :: Get (Ptr Word8, [Int])
  getChunksInfo = Get $ \endPtr s -> do
    let getChunks srcPtr l = do
          ensureBits endPtr s 8
          !n <- fromIntegral <$> peek srcPtr
          if n == 0
            then return (srcPtr `plusPtr` 1, l [])
            else do
              ensureBits endPtr s ((n + 1) * 8)
              getChunks (srcPtr `plusPtr` (n + 1)) (l . (n :)) -- ETA: stack overflow (missing tail call optimisation)
    when (usedBits s /= 0) $ badEncoding endPtr s "usedBits /= 0"
    (currPtr', ns) <- getChunks (currPtr s) id
    return $ GetResult (s {currPtr = currPtr'}) (currPtr s `plusPtr` 1, ns)
  
  -- | Given a value's decoder, returns the size in bits of the encoded value
  --
  -- @since 0.6
  sizeOf :: Get a -> Get Int
  sizeOf g =
    Get $ \end s -> do
      GetResult s' _ <- runGet g end s
      return $ GetResult s' $ (currPtr s' `minusPtr` currPtr s) * 8 - usedBits s + usedBits s'
  
  -- | Given a value's decoder, returns the value's bit encoding.
  --
  -- The encoding starts at the returned bit position in the return bytestring's first byte
  -- and ends in an unspecified bit position in its final byte
  --
  -- @since 0.6
  binOf :: Get a -> Get (B.ByteString, Int)
  binOf g =
    Get $ \end s -> do
      GetResult s' _ <- runGet g end s
      return $ GetResult s' (peekByteString (currPtr s) (currPtr s' `minusPtr` currPtr s + if usedBits s' == 0 then 0 else 1), usedBits s)
  