{-# 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 :: Natural -> Put
putNatural = \Natural
a ->
  let w8 :: Word8
w8 = Natural -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
a
  in case Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
unsafeShiftR Natural
a Int
7 of
       Natural
0 -> Word8 -> Put
Bin.putWord8 (Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x7f)
       Natural
b -> Word8 -> Put
Bin.putWord8 (Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
0x80) Put -> Put -> Put
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Natural -> Put
putNatural Natural
b

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

putWord8 :: Word8 -> Bin.Put
putWord8 :: Word8 -> Put
putWord8 = Natural -> Put
putNatural (Natural -> Put) -> (Word8 -> Natural) -> Word8 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word8 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putWord8 #-}

putWord16 :: Word16 -> Bin.Put
putWord16 :: Word16 -> Put
putWord16 = Natural -> Put
putNatural (Natural -> Put) -> (Word16 -> Natural) -> Word16 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putWord16 #-}

putWord32 :: Word32 -> Bin.Put
putWord32 :: Word32 -> Put
putWord32 = Natural -> Put
putNatural (Natural -> Put) -> (Word32 -> Natural) -> Word32 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putWord32 #-}

putWord64 :: Word64 -> Bin.Put
putWord64 :: Word64 -> Put
putWord64 = Natural -> Put
putNatural (Natural -> Put) -> (Word64 -> Natural) -> Word64 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putWord64 #-}

putWord :: Word -> Bin.Put
putWord :: Word -> Put
putWord = Natural -> Put
putNatural (Natural -> Put) -> (Word -> Natural) -> Word -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Natural
forall a b. (Integral a, Num b) => a -> b
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 :: Word -> Get Natural
getNatural Word
mx = String -> Get Natural -> Get Natural
forall a. String -> Get a -> Get a
Bin.label String
"ULEB128" (Word -> Get Natural
forall t a. (Num t, Num a, Bits a, Eq t) => t -> Get a
go Word
mx)
  where 
    go :: t -> Get a
go t
0 = String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"input too big"
    go t
n = do
      Word8
w8 <- Get Word8
Bin.getWord8
      if Word8
w8 Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
0x80
         then a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Get a) -> a -> Get a
forall a b. (a -> b) -> a -> b
$! Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w8
         else do 
           a
a <- t -> Get a
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1)
           a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Get a) -> a -> Get a
forall a b. (a -> b) -> a -> b
$! a -> Int -> a
forall a. Bits a => a -> Int -> a
unsafeShiftL a
a Int
7 a -> a -> a
forall a. Bits a => a -> a -> a
.|. Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
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 :: Word -> Get Integer
getInteger = (Natural -> Integer) -> Get Natural -> Get Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Get Natural -> Get Integer)
-> (Word -> Get Natural) -> Word -> Get Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Get Natural
getNatural 
{-# INLINE getInteger #-}

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

getBoundedIntegral 
  :: forall a. (Integral a, Bounded a, FiniteBits a) => Bin.Get a
getBoundedIntegral :: Get a
getBoundedIntegral = 
  let Word
bitSizeA :: Word = Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (a
forall a. HasCallStack => a
undefined :: a))
      Word
mxA :: Word = case Word -> Word -> (Word, Word)
forall a. Integral a => a -> a -> (a, a)
divMod Word
bitSizeA Word
7 of (Word
d, Word
m) -> Word
d Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word -> Word -> Word
forall a. Ord a => a -> a -> a
min Word
m Word
1
  in do Natural
n <- Word -> Get Natural
getNatural Word
mxA
        Get a -> (a -> Get a) -> Maybe a -> Get a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"overflow") a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> Maybe a
forall a b.
(Integral a, Integral b, Bits a, Bits b) =>
a -> Maybe b
toIntegralSized Natural
n)
{-# INLINE getBoundedIntegral #-}

getWord8 :: Bin.Get Word8
getWord8 :: Get Word8
getWord8 = Get Word8
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getWord8 #-}

getWord16 :: Bin.Get Word16
getWord16 :: Get Word16
getWord16 = Get Word16
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getWord16 #-}

getWord32 :: Bin.Get Word32
getWord32 :: Get Word32
getWord32 = Get Word32
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getWord32 #-}

getWord64 :: Bin.Get Word64
getWord64 :: Get Word64
getWord64 = Get Word64
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getWord64 #-}

getWord :: Bin.Get Word
getWord :: Get Word
getWord = Get Word
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getWord #-}

getInt8 :: Bin.Get Int8
getInt8 :: Get Int8
getInt8 = Get Int8
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getInt8 #-}

getInt16 :: Bin.Get Int16
getInt16 :: Get Int16
getInt16 = Get Int16
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getInt16 #-}

getInt32 :: Bin.Get Int32
getInt32 :: Get Int32
getInt32 = Get Int32
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getInt32 #-}

getInt64 :: Bin.Get Int64
getInt64 :: Get Int64
getInt64 = Get Int64
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getInt64 #-}

getInt :: Bin.Get Int
getInt :: Get Int
getInt = Get Int
forall a. (Integral a, Bounded a, FiniteBits a) => Get a
getBoundedIntegral
{-# INLINE getInt #-}

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

-- | Gets a strict 'B.ByteString' with its ULEB128-encoded length as prefix.
--
-- See 'putByteString'.
getByteString :: Bin.Get B.ByteString
getByteString :: Get ByteString
getByteString = Int -> Get ByteString
Bin.getByteString (Int -> Get ByteString) -> Get Int -> Get ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get Int
getInt
{-# INLINE getByteString #-}

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

-- | Gets a lazy 'BL.ByteString' with its ULEB128-encoded length as prefix.
--
-- See 'putLazyByteString'.
getLazyByteString :: Bin.Get BL.ByteString
getLazyByteString :: Get ByteString
getLazyByteString = Int64 -> Get ByteString
Bin.getLazyByteString (Int64 -> Get ByteString) -> Get Int64 -> Get ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get Int64
getInt64
{-# INLINE getLazyByteString #-}

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

-- | Gets a 'BS.ShortByteString' with its ULEB128-encoded length as prefix.
--
-- See 'putShortByteString'.
getShortByteString :: Bin.Get BS.ShortByteString
getShortByteString :: Get ShortByteString
getShortByteString = (ByteString -> ShortByteString)
-> Get ByteString -> Get ShortByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ShortByteString
BS.toShort (Int -> Get ByteString
Bin.getByteString (Int -> Get ByteString) -> Get Int -> Get ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get Int
getInt)
{-# INLINE getShortByteString #-}