{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Signed 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.SLEB128
 ( -- * Put
   putInteger
 , putInt64
 , putInt32
 , putInt16
 , putInt8
 , putInt
 , putNatural
 , putWord64
 , putWord32
 , putWord16
 , putWord8
 , putWord
   -- * Get
 , getInteger
 , getInt64
 , getInt32
 , getInt16
 , getInt8
 , getInt
 , getNatural
 , getWord64
 , getWord32
 , getWord16
 , getWord8
 , getWord
 ) where

import Control.Monad
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

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

putInteger :: Integer -> Bin.Put
putInteger :: Integer -> Put
putInteger = \Integer
a ->  do
  let w8 :: Word8
w8 = Integer -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer
a Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
0x7f) :: Word8
      b :: Integer
b = Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
unsafeShiftR Integer
a Int
7
      w8s :: Word8
w8s = Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x40 
  if (Word8
w8s Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0 Bool -> Bool -> Bool
&& Integer
b Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0) Bool -> Bool -> Bool
|| (Word8
w8s Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0 Bool -> Bool -> Bool
&& Integer
b Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== -Integer
1)
     then Word8 -> Put
Bin.putWord8 Word8
w8
     else do Word8 -> Put
Bin.putWord8 (Word8 -> Put) -> Word8 -> Put
forall a b. (a -> b) -> a -> b
$! Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
0x80
             Integer -> Put
putInteger Integer
b

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

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

putInt8 :: Int8 -> Bin.Put
putInt8 :: Int8 -> Put
putInt8 = Integer -> Put
putInteger (Integer -> Put) -> (Int8 -> Integer) -> Int8 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int8 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putInt8 #-}

putInt16 :: Int16 -> Bin.Put
putInt16 :: Int16 -> Put
putInt16 = Integer -> Put
putInteger (Integer -> Put) -> (Int16 -> Integer) -> Int16 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putInt16 #-}

putInt32 :: Int32 -> Bin.Put
putInt32 :: Int32 -> Put
putInt32 = Integer -> Put
putInteger (Integer -> Put) -> (Int32 -> Integer) -> Int32 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putInt32 #-}

putInt64 :: Int64 -> Bin.Put
putInt64 :: Int64 -> Put
putInt64 = Integer -> Put
putInteger (Integer -> Put) -> (Int64 -> Integer) -> Int64 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putInt64 #-}

putInt :: Int -> Bin.Put
putInt :: Int -> Put
putInt = Integer -> Put
putInteger (Integer -> Put) -> (Int -> Integer) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putInt #-}

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

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

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

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

putWord :: Word -> Bin.Put
putWord :: Word -> Put
putWord = Integer -> Put
putInteger (Integer -> Put) -> (Word -> Integer) -> Word -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE putWord #-}

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

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 Word
mx = String -> Get Integer -> Get Integer
forall a. String -> Get a -> Get a
Bin.label String
"SLEB128" (Word -> Int -> Integer -> Get Integer
f Word
mx Int
0 Integer
0)
  where
    f :: Word -> Int -> Integer -> Bin.Get Integer
    f :: Word -> Int -> Integer -> Get Integer
f Word
0 Int
_  Integer
_  = String -> Get Integer
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"input too big"
    f Word
n !Int
p !Integer
a = do
        Word8
w8 <- Get Word8
Bin.getWord8 
        let Integer
b :: Integer = Integer
a Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.|. Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
unsafeShiftL (Word8 -> Integer
forall a. Integral a => a -> Integer
toInteger (Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x7f)) Int
p
        case Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x80 of
          Word8
0 -> Integer -> Get Integer
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Integer -> Get Integer) -> Integer -> Get Integer
forall a b. (a -> b) -> a -> b
$! case Word8
w8 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x40 of 
                         Word8
0 -> Integer
b 
                         Word8
_ -> Integer
b Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Bits a => Int -> a
bit (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7)
          Word8
_ -> Word -> Int -> Integer -> Get Integer
f (Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Integer
b

getNatural
  :: 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 Natural
getNatural :: Word -> Get Natural
getNatural Word
mx = do
  Integer
i <- Word -> Get Integer
getInteger Word
mx
  Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$ String -> Get () -> Get ()
forall a. String -> Get a -> Get a
Bin.label String
"SLEB128" (String -> Get ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"underflow")
  Natural -> Get Natural
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
i)
{-# INLINE getNatural #-} 

-- TODO: The following dispatch to 'getInteger'. 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 Integer
i <- Word -> Get Integer
getInteger 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
"underflow or overflow") a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Integer -> Maybe a
forall a b.
(Integral a, Integral b, Bits a, Bits b) =>
a -> Maybe b
toIntegralSized Integer
i)
{-# INLINE getBoundedIntegral #-}

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 #-}

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 #-}