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.BaseSystem.BinaryTranscoder (BinaryTranscoder)
import Data.BaseSystem.BinaryTranscoder qualified as BXC
import Data.Bits ((.&.), (.<<.), (.>>.), (.|.))
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder qualified as Builder
import Data.ByteString.Short qualified as SBS
import Data.Word (Word16, Word64, Word8)
import MultiFormats.VarInt.Errors (Error (..), Okay)

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

-- | Byte with most significant bit (MSB) set to 1.
msb1 :: Word8
msb1 = 1 .<<. 7

-- | Byte with least significant bit (LSB) set to 1.
lsb1 :: Word8
lsb1 = 1

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

-- | The value that extracts all bits but the MSB in a bitwise AND.
mask7lsbs :: Word8
mask7lsbs = msb1 - 1

-- | The maximum value the spec allows is 2^63-1, so a 64-bit integer with bit
-- missing to account for the final byte's terminating MSB bit.
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 :: (BinaryTranscoder t) => t -> Okay (t, t)
spanVarInt bxc =
  -- Collect the bytes with MSB=1.
  case BXC.spanBytes isMSB1 bxc 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.
      | BXC.null msb0s -> Left MissingTermination
      -- The loop invariant ensures &msb1s + 1 has MSB=0, so compute this offset
      -- and split the data there.
      | otherwise ->
          let nulloffset = BXC.lengthBytes msb1s + 1
           in Right (BXC.splitAtByte nulloffset bxc)
  where
    -- Checks if the most-significant bit (MSB) is set to 1
    isMSB1 byte = (byte .&. msb1) /= 0

-- | Concats the lower 7 bits from bytes into single int value.
buildVarInt :: (BinaryTranscoder t) => t -> VarInt
buildVarInt = foldr build 0 . BXC.unpackBytes
  where
    -- Shifts the accumulator 7 bits (1 byte minus MSB value) and ORs the
    -- current byte into place, and remove the MSB set bit.
    build byte acc = acc .<<. 7 .|. fromIntegral (byte .&. mask7lsbs)

-- | 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!
{-# INLINE nonMinimal #-}
nonMinimal :: (BinaryTranscoder bxc) => bxc -> Bool
nonMinimal vardata = last16bits == invalidsuffix
  where
    -- Puts the last two bytes of the transcoder into a Word16.
    last16bits = fromIntegral $ BXC.unpackValue (BXC.takeBytesEnd 2 vardata)
    -- This value is the same to the value from bytes 0x7F, so it's not minimal
    invalidsuffix = (fromIntegral msb1lsb1 :: Word16) .<<. 8

-- | Checks if a VarInt encoded ByteString overflows a Word64.
{-# INLINE aboveCapacity #-}
aboveCapacity :: (BinaryTranscoder t) => t -> Bool
aboveCapacity vardata =
  let wholelen = BXC.lengthBytes vardata
   in not (1 <= wholelen && wholelen <= bytesCapacity)

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

-- | Puts a parsed VarInt (Word64 value) into a Builder of ShortByteString with
-- MSB bits set, and a max size of 9 bytes.
serializeVarInt :: VarInt -> Okay Builder
serializeVarInt varint
  | varint > maxVarInt = Left AboveMaxValue
  | otherwise =
      Right
        -- Put the ShortByteString into a ByteString Builder.
        . Builder.shortByteString
        -- Set the continuation bits in the bytes.
        . SBS.unfoldr setMSBs
        -- Replace empty bytestrings with a single zero byte.
        . replaceNull0
        -- Drop bytes until a minimal value is found.
        . SBS.dropWhileEnd (== 0)
        -- Pack the bytes from varint 7-bit groups.
        $ SBS.pack [var7sAt bytepos | bytepos <- [1 .. bytesCapacity]]
  where
    -- Extracts 1-indexed groups of 7 bits from varint starting from the LSB.
    {-# INLINEABLE var7sAt #-}
    var7sAt idx = fromIntegral (varint .>>. (7 * (idx - 1))) .&. mask7lsbs
    -- Sets MSB=1 for leading bytes in unfoldr.
    {-# INLINEABLE setMSBs #-}
    setMSBs bytes = do
      (top, rest) <- SBS.uncons bytes
      Just $
        if SBS.null rest
          then (top, SBS.empty)
          else (top .|. msb1, rest)
    -- Replaces a null ByteString with a single zero.
    {-# INLINEABLE replaceNull0 #-}
    replaceNull0 bytes =
      if SBS.null bytes
        then SBS.singleton 0
        else bytes