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 Control.Monad ((>=>))
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 Data.XCodec.BinaryTranscoder (BinaryTranscoder)
import Data.XCodec.BinaryTranscoder qualified as BXC
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
  -- Fail if empty bytes.
  | BXC.null bxc = Left EmptyBytes
  -- Otherwise, collect the bytes with MSB=1.
  | otherwise =
      case BXC.spanOctetsTop 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.totalOctets msb1s + 1
               in Right (BXC.splitOffset 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.unpackOctets
  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.unpackValueBE (BXC.takeOctetsEnd 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.totalOctets vardata
   in not (1 <= wholelen && wholelen <= bytesCapacity)

-- | Extracts the first minimally encoded varint from a ByteString.
extractVarInt :: (BinaryTranscoder t) => t -> Okay (VarInt, t)
extractVarInt = spanVarInt >=> spanBuild
  where
    -- Verifies the result from running spanVarInt
    spanBuild (vardata, rest)
      | aboveCapacity vardata = Left BytesOverflow
      | nonMinimal vardata = Left NotMinimal
      | otherwise = Right (buildVarInt vardata, 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 =
      let -- Expand the bits into byte-sized 7-bit groups
          grouped = SBS.pack [var7sAt bytepos | bytepos <- [1 .. bytesCapacity]]
          -- Get the minimal #bytes to represent a number
          minimal = replaceNull0 $ SBS.dropWhileEnd (== 0) grouped
          -- Set the MSB continuation bits
          setbits = SBS.unfoldr setMSBs minimal
       in -- Strictly evalate the resulting builder
          Right $! Builder.shortByteString setbits
  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