ULEB128.hs

{-# 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 = \a ->
  let w8 = fromIntegral a
  in case unsafeShiftR a 7 of
       0 -> Bin.putWord8 (w8 .&. 0x7f)
       b -> Bin.putWord8 (w8 .|. 0x80) >> putNatural b

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

putWord8 :: Word8 -> Bin.Put
putWord8 = putNatural . fromIntegral

putWord16 :: Word16 -> Bin.Put
putWord16 = putNatural . fromIntegral

putWord32 :: Word32 -> Bin.Put
putWord32 = putNatural . fromIntegral

putWord64 :: Word64 -> Bin.Put
putWord64 = putNatural . fromIntegral

putWord :: Word -> Bin.Put
putWord = putNatural . fromIntegral

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

getNatural :: Bin.Get Natural
getNatural = do
  w8 <- Bin.getWord8
  if w8 < 0x80
     then pure $! fromIntegral w8
     else do 
       a <- getNatural
       pure $! unsafeShiftL a 7 .|. fromIntegral (w8 .&. 0x7f)

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

getBoundedIntegral :: forall a. (Integral a, Bounded a) => String -> Bin.Get a
getBoundedIntegral label = do
    n <- getNatural
    let i = toInteger n
    if i <= maxA 
       then pure $! fromInteger i
       else fail err
  where
    err  :: String  = label <> ": overflow" 
    maxA :: Integer = toInteger (maxBound :: a) 

getWord8 :: Bin.Get Word8
getWord8 = getBoundedIntegral "Data.Binary.ULEB128.getWord8"

getWord16 :: Bin.Get Word16
getWord16 = getBoundedIntegral "Data.Binary.ULEB128.getWord16"

getWord32 :: Bin.Get Word32
getWord32 = getBoundedIntegral "Data.Binary.ULEB128.getWord32"

getWord64 :: Bin.Get Word64
getWord64 = getBoundedIntegral "Data.Binary.ULEB128.getWord64"

getWord :: Bin.Get Word
getWord = getBoundedIntegral "Data.Binary.ULEB128.getWord"