-- | Unsigned LEB128 varint encoding/decoding.
--
-- Used throughout libp2p for length-prefixed framing, protocol codes,
-- and multiaddr/multihash encoding.
--
-- Follows the multiformats unsigned-varint spec
-- (https://github.com/multiformats/unsigned-varint): varints are
-- restricted to a maximum of 9 bytes (63 bits), and non-minimal
-- (zero-padded) encodings are rejected on decode.
module LibP2P.Core.Varint
  ( encodeUvarint
  , decodeUvarint
  , maxVarintBytes
  ) 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 BL
import Data.Word (Word64)

-- | Maximum number of bytes for a valid unsigned varint.
-- The spec mandates: "Implementations MUST restrict the size of the
-- varint to a max of 9 bytes (63 bits)."
maxVarintBytes :: Int
maxVarintBytes :: Int
maxVarintBytes = Int
9

-- | Largest value representable in a spec-compliant varint (2^63 - 1).
maxVarintValue :: Word64
maxVarintValue :: Word64
maxVarintValue = (Word64
1 Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
63) Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1

-- | Encode a Word64 as an unsigned LEB128 varint.
-- Calls 'error' for values >= 2^63, which the spec makes unrepresentable
-- (go-varint's PutUvarint panics identically).
encodeUvarint :: Word64 -> ByteString
encodeUvarint :: Word64 -> ByteString
encodeUvarint Word64
n
  | Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
maxVarintValue =
      String -> ByteString
forall a. HasCallStack => String -> a
error String
"encodeUvarint: value exceeds 63 bits (unsigned-varint spec maximum)"
  | Bool
otherwise = LazyByteString -> ByteString
BL.toStrict (Builder -> LazyByteString
Builder.toLazyByteString (Word64 -> Builder
go Word64
n))
  where
    go :: Word64 -> Builder.Builder
    go :: Word64 -> Builder
go Word64
m
      | Word64
m 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
m)
      | Bool
otherwise =
          Word8 -> Builder
Builder.word8 (Word64 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64
m 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
m 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.
-- Rejects varints longer than 9 bytes and non-minimal encodings
-- (a multi-byte varint whose final byte is 0x00).
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 9 bytes / 63 bits)"
      | 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
                  -- Spec: leading zeros "must be rejected when decoding.
                  -- The only number that can end in a 0x00 is 0" — and 0
                  -- is the single-byte encoding 0x00 (bitShift == 0).
                  if Word8
byte Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0x00 Bool -> Bool -> Bool
&& Int
bitShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                    then String -> Either String (Word64, ByteString)
forall a b. a -> Either a b
Left String
"decodeUvarint: non-minimal encoding (trailing zero byte)"
                    else (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'