packages feed

haskell-opentimestamps-0.5.5.0: src/OpenTimestamps/VarInt.hs

{- | Variable-length integer encoding for OpenTimestamps.

This module provides functions for reading and writing variable-length
integers and byte arrays as used in the OpenTimestamps protocol.
-}
module OpenTimestamps.VarInt
  ( getVarBytes
  , putVarBytes
  , getVarInt
  , putVarInt
  ) where

import Data.Binary.Get
  ( Get
  , getByteString
  , getWord8
  )
import Data.Binary.Put (Put, putByteString, putWord8)
import Data.Bits (shiftL, shiftR, (.&.), (.|.))
import qualified Data.ByteString as BS

-- | Read variable-length bytes with length limits.
getVarBytes :: Int -> Int -> Get BS.ByteString
getVarBytes maxLen minLen = do
  len <- getVarInt
  if len > maxLen
    then fail $ "varbytes max length exceeded; " ++ show len ++ " > " ++ show maxLen
    else
      if len < minLen
        then fail $ "varbytes min length not met; " ++ show len ++ " < " ++ show minLen
        else getByteString len

-- | Read a variable-length integer.
getVarInt :: Get Int
getVarInt = getVarInt' 0 0
  where
    -- \| Helper function to recursively read variable-length integer.
    getVarInt' acc shift = do
      byte <- getWord8
      let val = fromIntegral (byte .&. 0x7f)
      let newAcc = acc .|. (val `shiftL` shift)
      if (byte .&. 0x80) == 0
        then pure newAcc
        else getVarInt' newAcc (shift + 7)

-- | Write variable-length bytes with length prefix.
putVarBytes :: BS.ByteString -> Put
putVarBytes bs = do
  putVarInt (BS.length bs)
  putByteString bs

-- | Write a variable-length integer.
putVarInt :: Int -> Put
putVarInt n
  | n == 0 = putWord8 0x00
  | otherwise = putVarInt' n
  where
    -- \| Helper function to recursively write variable-length integer.
    putVarInt' i = do
      let byte = fromIntegral (i .&. 0x7f)
      let rest = i `shiftR` 7
      if rest > 0
        then putWord8 (byte .|. 0x80) >> putVarInt' rest
        else putWord8 byte