{-# 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 = \a ->  do
  let w8 = fromIntegral (a .&. 0x7f) :: Word8
      b = unsafeShiftR a 7
      w8s = w8 .&. 0x40 
  if (w8s == 0 && b == 0) || (w8s /= 0 && b == -1)
     then Bin.putWord8 w8
     else do Bin.putWord8 $! w8 .|. 0x80
             putInteger b

putNatural :: Natural -> Bin.Put
putNatural = putInteger . fromIntegral
{-# INLINE putNatural #-}

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

putInt8 :: Int8 -> Bin.Put
putInt8 = putInteger . fromIntegral
{-# INLINE putInt8 #-}

putInt16 :: Int16 -> Bin.Put
putInt16 = putInteger . fromIntegral
{-# INLINE putInt16 #-}

putInt32 :: Int32 -> Bin.Put
putInt32 = putInteger . fromIntegral
{-# INLINE putInt32 #-}

putInt64 :: Int64 -> Bin.Put
putInt64 = putInteger . fromIntegral
{-# INLINE putInt64 #-}

putInt :: Int -> Bin.Put
putInt = putInteger . fromIntegral
{-# INLINE putInt #-}

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

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

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

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

putWord :: Word -> Bin.Put
putWord = putInteger . 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 mx = Bin.label "SLEB128" (f mx 0 0)
  where
    f :: Word -> Int -> Integer -> Bin.Get Integer
    f 0 _  _  = fail "input too big"
    f n !p !a = do
        w8 <- Bin.getWord8 
        let b :: Integer = a .|. unsafeShiftL (toInteger (w8 .&. 0x7f)) p
        case w8 .&. 0x80 of
          0 -> pure $! case w8 .&. 0x40 of 
                         0 -> b 
                         _ -> b - bit (p + 7)
          _ -> f (n - 1) (p + 7) 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 mx = do
  i <- getInteger mx
  when (i < 0) $ Bin.label "SLEB128" (fail "underflow")
  pure (fromInteger i)
{-# INLINE getNatural #-} 

-- TODO: The following dispatch to 'getInteger'. 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 i <- getInteger mxA
        maybe (fail "underflow or overflow") pure (toIntegralSized i)
{-# INLINE getBoundedIntegral #-}

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

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