-- | Unsigned LEB128 varint encoding/decoding.
--
-- Used throughout libp2p for length-prefixed framing, protocol codes,
-- and multiaddr/multihash encoding.
module Network.LibP2P.Core.Varint
  ( encodeUvarint
  , decodeUvarint
  ) where

import Data.Bits (Bits (..))
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as Builder
import qualified Data.ByteString.Lazy as LBS
import Data.Word (Word64)

-- | Maximum number of bytes for a valid unsigned varint (ceil(64/7) = 10).
maxVarintBytes :: Int
maxVarintBytes :: Int
maxVarintBytes = Int
10

-- | Encode a Word64 as an unsigned LEB128 varint.
encodeUvarint :: Word64 -> ByteString
encodeUvarint :: Word64 -> ByteString
encodeUvarint = LazyByteString -> ByteString
LBS.toStrict (LazyByteString -> ByteString)
-> (Word64 -> LazyByteString) -> Word64 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
Builder.toLazyByteString (Builder -> LazyByteString)
-> (Word64 -> Builder) -> Word64 -> LazyByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Builder
go
  where
    go :: Word64 -> Builder.Builder
    go :: Word64 -> Builder
go Word64
n
      | Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
0x80 = Word8 -> Builder
Builder.word8 (Word64 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n)
      | Bool
otherwise =
          Word8 -> Builder
Builder.word8 (Word64 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
n Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
0x7f) Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
0x80)
            Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
go (Word64
n Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftR` Int
7)

-- | Decode an unsigned LEB128 varint from a ByteString.
-- Returns the decoded value and remaining bytes, or an error message.
decodeUvarint :: ByteString -> Either String (Word64, ByteString)
decodeUvarint :: ByteString -> Either String (Word64, ByteString)
decodeUvarint ByteString
bs
  | ByteString -> Bool
BS.null ByteString
bs = String -> Either String (Word64, ByteString)
forall a b. a -> Either a b
Left String
"decodeUvarint: empty input"
  | Bool
otherwise = ByteString -> Int -> Word64 -> Either String (Word64, ByteString)
go ByteString
bs Int
0 Word64
0
  where
    go :: ByteString -> Int -> Word64 -> Either String (Word64, ByteString)
    go :: ByteString -> Int -> Word64 -> Either String (Word64, ByteString)
go ByteString
input Int
bitShift Word64
acc
      | Int
bitShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxVarintBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
7 =
          String -> Either String (Word64, ByteString)
forall a b. a -> Either a b
Left String
"decodeUvarint: varint too long (exceeds 10 bytes)"
      | ByteString -> Bool
BS.null ByteString
input =
          String -> Either String (Word64, ByteString)
forall a b. a -> Either a b
Left String
"decodeUvarint: unexpected end of input"
      | Bool
otherwise =
          let byte :: Word8
byte = HasCallStack => ByteString -> Word8
ByteString -> Word8
BS.head ByteString
input
              rest :: ByteString
rest = HasCallStack => ByteString -> ByteString
ByteString -> ByteString
BS.tail ByteString
input
              val :: Word64
val = Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8
byte Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x7f) :: Word64
              acc' :: Word64
acc' = Word64
acc Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word64
val Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
bitShift)
           in if Word8
byte Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8
0x80 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0
                then (Word64, ByteString) -> Either String (Word64, ByteString)
forall a b. b -> Either a b
Right (Word64
acc', ByteString
rest)
                else ByteString -> Int -> Word64 -> Either String (Word64, ByteString)
go ByteString
rest (Int
bitShift Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7) Word64
acc'