SLEB128.hs

{-# 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
   -- * Get
 , getInteger
 , getInt64
 , getInt32
 , getInt16
 , getInt8
 , getInt
 ) 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

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

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

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

putInt8 :: Int8 -> Bin.Put
putInt8 = putInteger . fromIntegral

putInt16 :: Int16 -> Bin.Put
putInt16 = putInteger . fromIntegral

putInt32 :: Int32 -> Bin.Put
putInt32 = putInteger . fromIntegral

putInt64 :: Int64 -> Bin.Put
putInt64 = putInteger . fromIntegral

putInt :: Int -> Bin.Put
putInt = putInteger . fromIntegral

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

getInteger :: Bin.Get Integer
getInteger = f 0 0
  where
    f :: Int -> Integer -> Bin.Get Integer
    f !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 (p + 7) b

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

getBoundedIntegral :: forall a. (Integral a, Bounded a) => String -> Bin.Get a
getBoundedIntegral label = do
    i <- getInteger
    when (i < minA) (fail erru)
    when (i > maxA) (fail erro)
    pure $! fromInteger i
  where
    erru :: String  = label <> ": underflow"
    erro :: String  = label <> ": overflow" 
    minA :: Integer = toInteger (minBound :: a) 
    maxA :: Integer = toInteger (maxBound :: a) 

getInt8 :: Bin.Get Int8
getInt8 = getBoundedIntegral "Data.Binary.SLEB128.getInt8"

getInt16 :: Bin.Get Int16
getInt16 = getBoundedIntegral "Data.Binary.SLEB128.getInt16"

getInt32 :: Bin.Get Int32
getInt32 = getBoundedIntegral "Data.Binary.SLEB128.getInt32"

getInt64 :: Bin.Get Int64
getInt64 = getBoundedIntegral "Data.Binary.SLEB128.getInt64"

getInt :: Bin.Get Int
getInt = getBoundedIntegral "Data.Binary.SLEB128.getInt"