{-# 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
 ) where

import qualified Data.Binary.Get as Bin
import qualified Data.Binary.Put as Bin
import Data.Bits
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

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

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

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

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

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

getNatural :: Bin.Get Natural
getNatural :: Get Natural
getNatural = do
  Word8
w8 <- Get Word8
Bin.getWord8
  if Word8
w8 Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
< Word8
0x80
     then Natural -> Get Natural
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> Get Natural) -> Natural -> Get Natural
forall a b. (a -> b) -> a -> b
$! Word8 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
w8
     else do 
       Natural
a <- Get Natural
getNatural
       Natural -> Get Natural
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> Get Natural) -> Natural -> Get Natural
forall a b. (a -> b) -> a -> b
$! Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
unsafeShiftL Natural
a Int
7 Natural -> Natural -> Natural
forall a. Bits a => a -> a -> a
.|. Word8 -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x7f)

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

getBoundedIntegral :: forall a. (Integral a, Bounded a) => String -> Bin.Get a
getBoundedIntegral :: String -> Get a
getBoundedIntegral String
label = do
    Natural
n <- Get Natural
getNatural
    let i :: Integer
i = Natural -> Integer
forall a. Integral a => a -> Integer
toInteger Natural
n
    if Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
maxA 
       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
$! Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
i
       else String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
err
  where
    String
err  :: String  = String
label String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": overflow" 
    Integer
maxA :: Integer = a -> Integer
forall a. Integral a => a -> Integer
toInteger (a
forall a. Bounded a => a
maxBound :: a) 

getWord8 :: Bin.Get Word8
getWord8 :: Get Word8
getWord8 = String -> Get Word8
forall a. (Integral a, Bounded a) => String -> Get a
getBoundedIntegral String
"Data.Binary.ULEB128.getWord8"

getWord16 :: Bin.Get Word16
getWord16 :: Get Word16
getWord16 = String -> Get Word16
forall a. (Integral a, Bounded a) => String -> Get a
getBoundedIntegral String
"Data.Binary.ULEB128.getWord16"

getWord32 :: Bin.Get Word32
getWord32 :: Get Word32
getWord32 = String -> Get Word32
forall a. (Integral a, Bounded a) => String -> Get a
getBoundedIntegral String
"Data.Binary.ULEB128.getWord32"

getWord64 :: Bin.Get Word64
getWord64 :: Get Word64
getWord64 = String -> Get Word64
forall a. (Integral a, Bounded a) => String -> Get a
getBoundedIntegral String
"Data.Binary.ULEB128.getWord64"

getWord :: Bin.Get Word
getWord :: Get Word
getWord = String -> Get Word
forall a. (Integral a, Bounded a) => String -> Get a
getBoundedIntegral String
"Data.Binary.ULEB128.getWord"