{-# LANGUAGE ScopedTypeVariables #-}

-- | Unsigned LEB128 codec.
--
-- Any /getXXX/ decoder can decode bytes generated using any of the /putXXX/ 
-- encoders, provided the encoded number fits in the target type.
module Data.Binary.ULEB128
 ( -- * Put
   putNatural
 , putWord64
 , putWord32
 , putWord16
 , putWord8
 , putWord

   -- * Get
 , getNatural
 , getWord64
 , getWord32
 , getWord16
 , getWord8
 , getWord
 , getInteger
 , getInt64
 , getInt32
 , getInt16
 , getInt8
 , getInt
   
   -- * ByteString
 , putByteString
 , getByteString
   -- ** Lazy
 , putLazyByteString
 , getLazyByteString
   -- ** Short
 , putShortByteString
 , getShortByteString
 ) where

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Short as BS
import qualified Data.Binary.Get as Bin
import qualified Data.Binary.Put as Bin
import Data.Bits
import Data.Int
import Data.Word
import Numeric.Natural

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

putNatural :: Natural -> Bin.Put
putNatural = \a ->
  let w8 = fromIntegral a
  in case unsafeShiftR a 7 of
       0 -> Bin.putWord8 (w8 .&. 0x7f)
       b -> Bin.putWord8 (w8 .|. 0x80) >> putNatural b

-- TODO: The following dispatch to 'putNatural'. Make faster.

putWord8 :: Word8 -> Bin.Put
putWord8 = putNatural . fromIntegral
{-# INLINE putWord8 #-}

putWord16 :: Word16 -> Bin.Put
putWord16 = putNatural . fromIntegral
{-# INLINE putWord16 #-}

putWord32 :: Word32 -> Bin.Put
putWord32 = putNatural . fromIntegral
{-# INLINE putWord32 #-}

putWord64 :: Word64 -> Bin.Put
putWord64 = putNatural . fromIntegral
{-# INLINE putWord64 #-}

putWord :: Word -> Bin.Put
putWord = putNatural . fromIntegral
{-# INLINE putWord #-}

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

getNatural 
  :: Word  
  -- ^ /Maximum/ number of bytes to consume. If the 'Natural' number can be
  -- determined before consuming this number of bytes, it will be. If @0@, 
  -- parsing fails. 
  --
  -- Each ULEB128 byte encodes at most 7 bits of data. That is, 
  -- \(length(encoded) == \lceil\frac{length(data)}{7}\rceil\).
  -> Bin.Get Natural
getNatural mx = Bin.label "ULEB128" (go mx)
  where 
    go 0 = fail "input too big"
    go n = do
      w8 <- Bin.getWord8
      if w8 < 0x80
         then pure $! fromIntegral w8
         else do 
           a <- go (n - 1)
           pure $! unsafeShiftL a 7 .|. fromIntegral (w8 .&. 0x7f)

getInteger 
  :: Word
  -- ^ /Maximum/ number of bytes to consume. If the 'Integer' number can be
  -- determined before consuming this number of bytes, it will be. If @0@, 
  -- parsing fails. 
  --
  -- Each ULEB128 byte encodes at most 7 bits of data. That is, 
  -- \(length(encoded) == \lceil\frac{length(data)}{7}\rceil\).
  -> Bin.Get Integer
getInteger = fmap toInteger . getNatural 
{-# INLINE getInteger #-}

-- TODO: The following dispatch to 'getNatural'. Make faster.

getBoundedIntegral 
  :: forall a. (Integral a, Bounded a, FiniteBits a) => Bin.Get a
getBoundedIntegral = 
  let bitSizeA :: Word = fromIntegral (finiteBitSize (undefined :: a))
      mxA :: Word = case divMod bitSizeA 7 of (d, m) -> d + min m 1
  in do n <- getNatural mxA
        maybe (fail "overflow") pure (toIntegralSized n)
{-# INLINE getBoundedIntegral #-}

getWord8 :: Bin.Get Word8
getWord8 = getBoundedIntegral
{-# INLINE getWord8 #-}

getWord16 :: Bin.Get Word16
getWord16 = getBoundedIntegral
{-# INLINE getWord16 #-}

getWord32 :: Bin.Get Word32
getWord32 = getBoundedIntegral
{-# INLINE getWord32 #-}

getWord64 :: Bin.Get Word64
getWord64 = getBoundedIntegral
{-# INLINE getWord64 #-}

getWord :: Bin.Get Word
getWord = getBoundedIntegral
{-# INLINE getWord #-}

getInt8 :: Bin.Get Int8
getInt8 = getBoundedIntegral
{-# INLINE getInt8 #-}

getInt16 :: Bin.Get Int16
getInt16 = getBoundedIntegral
{-# INLINE getInt16 #-}

getInt32 :: Bin.Get Int32
getInt32 = getBoundedIntegral
{-# INLINE getInt32 #-}

getInt64 :: Bin.Get Int64
getInt64 = getBoundedIntegral
{-# INLINE getInt64 #-}

getInt :: Bin.Get Int
getInt = getBoundedIntegral
{-# INLINE getInt #-}

--------------------------------------------------------------------------------
        
-- | Puts a strict 'B.ByteString' with its ULEB128-encoded length as prefix.
--
-- See 'getByteString'.
putByteString :: B.ByteString -> Bin.Put
putByteString = \a -> do
  putNatural (fromIntegral (B.length a :: Int))
  Bin.putByteString a
{-# INLINE putByteString #-}

-- | Gets a strict 'B.ByteString' with its ULEB128-encoded length as prefix.
--
-- See 'putByteString'.
getByteString :: Bin.Get B.ByteString
getByteString = Bin.getByteString =<< getInt
{-# INLINE getByteString #-}

-- | Puts a lazy 'B.ByteString' with its ULEB128-encoded length as prefix.
--
-- See 'getLazyByteString'.
putLazyByteString :: BL.ByteString -> Bin.Put
putLazyByteString = \a -> do
  putNatural (fromIntegral (BL.length a :: Int64))
  Bin.putLazyByteString a
{-# INLINE putLazyByteString #-}

-- | Gets a lazy 'BL.ByteString' with its ULEB128-encoded length as prefix.
--
-- See 'putLazyByteString'.
getLazyByteString :: Bin.Get BL.ByteString
getLazyByteString = Bin.getLazyByteString =<< getInt64
{-# INLINE getLazyByteString #-}

-- | Puts a 'BS.ShortByteString' with its ULEB128-encoded length as prefix.
--
-- See 'getShortByteString'.
putShortByteString :: BS.ShortByteString -> Bin.Put
putShortByteString = \a -> do
  putNatural (fromIntegral (BS.length a :: Int))
  Bin.putShortByteString a
{-# INLINE putShortByteString #-}

-- | Gets a 'BS.ShortByteString' with its ULEB128-encoded length as prefix.
--
-- See 'putShortByteString'.
getShortByteString :: Bin.Get BS.ShortByteString
getShortByteString = fmap BS.toShort (Bin.getByteString =<< getInt)
{-# INLINE getShortByteString #-}