Integer.hs

{-# LANGUAGE BinaryLiterals #-}

module Network.QUIC.Types.Integer (
    encodeInt
  , encodeInt8
  , encodeInt'
  , encodeInt'2
  , encodeInt'4
  , decodeInt
  , decodeInt'
  ) where

import Data.ByteString.Internal (unsafeCreate)
import Foreign.Ptr
import Foreign.Storable
import System.IO.Unsafe (unsafeDupablePerformIO)

import Network.QUIC.Imports

-- $setup
-- >>> :set -XOverloadedStrings
-- >>> import Network.QUIC.Utils

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

-- |
-- >>> enc16 $ encodeInt 151288809941952652
-- "c2197c5eff14e88c"
-- >>> enc16 $ encodeInt 494878333
-- "9d7f3e7d"
-- >>> enc16 $ encodeInt 15293
-- "7bbd"
-- >>> enc16 $ encodeInt 37
-- "25"
encodeInt :: Int64  -> ByteString
encodeInt i = unsafeCreate n $ go tag n i'
  where
    (tag,n,i') = tagLen i
{-# NOINLINE encodeInt #-}

encodeInt8 :: Int64  -> ByteString
encodeInt8 i = unsafeCreate n $ go tag n i
  where
    n = 8
    tag = 0b11000000
{-# NOINLINE encodeInt8 #-}

encodeInt' :: WriteBuffer -> Int64 -> IO ()
encodeInt' wbuf i = go' tag n i' wbuf
  where
    (tag,n,i') = tagLen i

encodeInt'2 :: WriteBuffer -> Int64 -> IO ()
encodeInt'2 wbuf i = go' tag n i' wbuf
  where
    tag = 0b01000000
    n = 2
    i' = i .<<. 48

encodeInt'4 :: WriteBuffer -> Int64 -> IO ()
encodeInt'4 wbuf i = go' tag n i' wbuf
  where
    tag = 0b10000000
    n = 4
    i' = i .<<. 32

tagLen :: Int64 -> (Word8, Int, Int64)
tagLen i | i <=         63 = (0b00000000, 1, i .<<. 56)
         | i <=      16383 = (0b01000000, 2, i .<<. 48)
         | i <= 1073741823 = (0b10000000, 4, i .<<. 32)
         | otherwise       = (0b11000000, 8, i)
{-# INLINE tagLen #-}

msb8 :: Int64 -> Word8
msb8 i = fromIntegral (i .>>. 56)
{-# INLINE msb8 #-}

go :: Word8 -> Int -> Int64 -> Ptr Word8 -> IO ()
go tag n0 i0 p0 = do
    poke p0 (tag .|. msb8 i0)
    let n' = n0 - 1
        i' = i0 .<<. 8
        p' = p0 `plusPtr` 1
    loop n' i' p'
  where
    loop 0 _ _ = return ()
    loop n i p = do
        poke p $ msb8 i
        let n' = n - 1
            i' = i .<<. 8
            p' = p `plusPtr` 1
        loop n' i' p'
{-# INLINE go #-}

go' :: Word8 -> Int -> Int64 -> WriteBuffer -> IO ()
go' tag n0 i0 wbuf = do
    write8 wbuf (tag .|. msb8 i0)
    let n' = n0 - 1
        i' = i0 .<<. 8
    loop n' i'
  where
    loop 0 _ = return ()
    loop n i = do
        write8 wbuf $ msb8 i
        let n' = n - 1
            i' = i .<<. 8
        loop n' i'
{-# INLINE go' #-}

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

-- |
-- >>> decodeInt (dec16 "c2197c5eff14e88c")
-- 151288809941952652
-- >>> decodeInt (dec16 "9d7f3e7d")
-- 494878333
-- >>> decodeInt (dec16 "7bbd")
-- 15293
-- >>> decodeInt (dec16 "25")
-- 37
decodeInt :: ByteString -> Int64
decodeInt bs = unsafeDupablePerformIO $ withReadBuffer bs decodeInt'
{-# NOINLINE decodeInt #-}

decodeInt' :: ReadBuffer -> IO Int64
decodeInt' rbuf = do
    b0 <- read8 rbuf
    let flag = b0 .>>. 6
        b1 = fromIntegral (b0 .&. 0b00111111)
    case flag of
      0 -> return b1
      1 -> loop b1 1
      2 -> loop b1 3
      _ -> loop b1 7
  where
    loop :: Int64 -> Int -> IO Int64
    loop r 0 = return r
    loop r n = do
        b <- fromIntegral <$> read8 rbuf
        loop (r*256 + b) (n - 1)