VarInt.hs

-- | Module      : MultiFormats.VarInt
--   Description : Implements multiformat's varints
--   Copyright   : Zoey McBride (c) 2026
--   License     : AGPL-3.0-or-later
--   Maintainer  : zoeymcbride@mailbox.org
--   Stability   : experimental
--
-- Implements serialization and extraction for multiformat's unsigned binary
-- varint spec. See: https://github.com/multiformats/unsigned-varint
module MultiFormats.VarInt
  ( VarInt,
    extractVarInt,
    serializeVarInt,
    maxVarInt,
  )
where

import Data.Bits ((.&.), (.<<.), (.>>.), (.|.))
import Data.ByteString (ByteString)
import Data.ByteString qualified as Bytes
import Data.Word (Word64, Word8)
import MultiFormats.VarInt.Errors (Error (..), Okay)

-- | VarInt spec uses unsigned 64 bit integers.
type VarInt = Word64

-- | Byte with MSB=1.
msb1 :: Word8
msb1 = 1 .<<. 7

-- | Byte with LSB=1.
lsb1 :: Word8
lsb1 = 1

-- | The value w/ MSB=1 and LSB=1.
msb1lsb1 :: Word8
msb1lsb1 = msb1 .|. lsb1

-- | The value that extracts all bits but the MSB.
masklower :: Word8
masklower = msb1 - 1

-- | The maximum value the spec allows (2^63-1).
maxVarInt :: VarInt
maxVarInt = (1 .<<. 63) - 1

-- | The maximum # of bytes a VarInt can occupy when encoded.
bytesCapacity :: Int
bytesCapacity = 9

-- | Splits ByteString into it's first occurring VarInt and the remaining
-- ByteString after.
spanVarInt :: ByteString -> Maybe (ByteString, ByteString)
spanVarInt bytes =
  -- Collect the bytes with MSB=1.
  case Bytes.span isMSB1 bytes of
    -- The varint is the leading bytes with the MSB=1 and a final terminating
    -- byte with MSB=0.
    (msb1s, msb0s)
      -- Verify if a final byte exists.
      | Bytes.null msb0s -> Nothing
      -- The loop invariant ensures msb + 1 has MSB=0. Add the last byte with
      -- splitAt since it is (a lot) more efficent than append.
      | otherwise -> Just (Bytes.splitAt (Bytes.length msb1s + 1) bytes)
  where
    isMSB1 byte = (byte .&. msb1) /= 0

-- | Concats the lower 7 bits from bytes into single int value.
buildVarInt :: ByteString -> VarInt
buildVarInt = Bytes.foldr (\byte int -> advance int .|. lower byte) 0
  where
    advance int = int .<<. 7
    lower byte = fromIntegral $ byte .&. masklower

-- | Checks edge case where input is invalid from not being a minimal encoding
-- of varints, (ie, we can fit 2 bytes of 7 bits into 1 byte of 7 bits). The
-- only case of this is when a final 0x8100 should be a 0x79... I think!
nonMinimal :: ByteString -> Bool
nonMinimal bytes = Bytes.unpack (Bytes.takeEnd 2 bytes) == [msb1lsb1, 0]

-- | Checks if a VarInt encoded ByteString overflows a Word64.
aboveCapacity :: ByteString -> Bool
aboveCapacity varbytes =
  let wholelen = Bytes.length varbytes
   in not (1 <= wholelen && wholelen <= bytesCapacity)

-- | Extracts the first minimally encoded varint from a ByteString.
extractVarInt :: ByteString -> Okay (VarInt, ByteString)
extractVarInt bytes
  | Bytes.null bytes = Left EmptyBytes
  | otherwise =
      case spanVarInt bytes of
        Nothing -> Left MissingTermination
        Just (varbytes, rest)
          | aboveCapacity varbytes -> Left BytesOverflow
          | nonMinimal varbytes -> Left NotMinimal
          | otherwise -> Right (buildVarInt varbytes, rest)

-- | Puts a parsed VarInt (Word64, without MSB indicators) into a ByteString
-- with MSB indicators, and a max size of 9 bytes.
serializeVarInt :: VarInt -> Okay ByteString
serializeVarInt varint
  | varint > maxVarInt = Left AboveMaxValue
  | otherwise =
      Right
        . zeroReplaceNull
        . Bytes.unfoldr setMSBs
        . Bytes.dropWhileEnd (== 0)
        $ Bytes.pack [varintGroup i varint | i <- [1 .. bytesCapacity]]
  where
    -- Extracts Nth group of 7 bits from VarInt.
    varintGroup n var = fromIntegral (var .>>. (7 * (n - 1))) .&. masklower
    -- Sets MSB=1 for leading bytes in unfoldr.
    setMSBs bytes =
      Bytes.uncons bytes >>= \(byte, rest) ->
        if Bytes.null rest
          then Just (byte, mempty)
          else Just (byte .|. msb1, rest)
    -- Replaces a null ByteString with a single zero.
    zeroReplaceNull bytes =
      if Bytes.null bytes
        then Bytes.singleton 0
        else bytes