Transcoder.hs

-- | Module      : XCodec.Transcoder
--   Description : Type-class for building generic codecs on binary data.
--   Copyright   : Zoey McBride (c) 2026
--   License     : BSD-3-Clause
--   Maintainer  : zoeymcbride@mailbox.org
--   Stability   : experimental
module XCodec.Transcoder
  ( -- * Interface for generic transcoding of binary data.
    Transcoder (..),

    -- * Transcoder data as numeric values
    BitSet,
    Octet,

    -- * Conversion of Integral values to transcoder data.
    packValueBE,
    packValueLE,

    -- * Transcoder utilities
    fromOctet,
    null,
    uncons,
    unsnoc,
    takeOctetsTop,
    dropOctetsTop,
    takeOctetsEnd,
    dropOctetsEnd,
    takeWhileEndOctets,
    takeWhileTopOctets,
    dropWhileEndOctets,
    dropWhileTopOctets,
  )
where

import Data.Bits (Bits, (.&.), (.<<.), (.>>.), (.|.))
import Data.ByteString (ByteString)
import Data.ByteString qualified as Bytes
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder qualified as BSB
import Data.ByteString.Builder qualified as Builder
import Data.ByteString.Lazy (LazyByteString)
import Data.ByteString.Lazy qualified as LBS
import Data.ByteString.Short (ShortByteString)
import Data.ByteString.Short qualified as SBS
import Data.List qualified as List
import Data.Word (Word8)
import Prelude hiding (null)

-- | Stores 8-bit values.
type Octet = Word8

-- | Stores Transcoder as numeric values.
type BitSet = Integer

-- | Class of functions for working on binary transcoder streams, that is
-- streams of bytes that we can easily turn into a bitset, and can partition
-- into smaller streams for extracting and parsing data.
class (Monoid tc) => Transcoder tc where
  {-# MINIMAL
    totalOctets,
    swapOrder,
    packOctets,
    execBuilder,
    unpackValueCPU,
    unpackValueBE,
    unpackValueLE,
    unpackBuilder,
    unpackOctets,
    pushOctetTop,
    pushOctetEnd,
    spanOctetsTop,
    spanOctetsEnd,
    splitOffset,
    replicateOctet
    #-}

  -- | Gives back the length in octets of the Transcoder data.
  totalOctets :: tc -> Int

  -- | Reverses the byte order of the transcoder data.
  swapOrder :: tc -> tc

  -- | Packs Octet values into transcoder data `tc`.
  packOctets :: [Octet] -> tc

  -- | Executes the `Builder` into transcoder data.
  execBuilder :: Builder -> tc

  -- | Unpacks transcoder data into a list of bytes.
  unpackOctets :: tc -> [Octet]

  -- | Unpacks transcoder data as BitSet in the host CPU's endianness.
  unpackValueCPU :: tc -> BitSet

  -- | Unpacks transcoder data as BitSet in Big-Endian (BE) order.
  unpackValueBE :: tc -> BitSet

  -- | Unpacks transcoder data as BitSet in Little-Endian (LE) order.
  unpackValueLE :: tc -> BitSet

  -- | Places the contents of the transcoder into a ByteString `Builder`.
  unpackBuilder :: tc -> Builder

  -- | Prepends an octet value to the top of the transcoder data.
  pushOctetTop :: tc -> Octet -> tc

  -- | Appends the octet value to the end of the transcoder data.
  pushOctetEnd :: tc -> Octet -> tc

  -- | Parititions the binary data at
  splitOffset :: Int -> tc -> (tc, tc)

  -- | Repeats an octet value in binary data.
  replicateOctet :: Int -> Octet -> tc

  -- | Splits octet values based on the first false result from the predicate.
  spanOctetsTop :: (Octet -> Bool) -> tc -> (tc, tc)

  -- | Splits the list while the element predicate holds staring from the left
  -- side.
  spanOctetsEnd :: (Octet -> Bool) -> tc -> (tc, tc)

-- | Implements binary decoder for Bytestrings, best for constant values that
-- exist for the lifetime of the program.
-- https://hackage-content.haskell.org/package/bytestring-0.12.2.0/docs/Data-ByteString.html#g:2
instance Transcoder ByteString where
  totalOctets = Bytes.length
  swapOrder = Bytes.reverse
  execBuilder = LBS.toStrict . execBuilder
  packOctets = Bytes.pack
  unpackBuilder = Builder.byteString
  unpackValueCPU = unpackValueCPU . LBS.fromStrict
  unpackValueLE = unpackValueLE . LBS.fromStrict
  unpackValueBE = unpackValueBE . LBS.fromStrict
  unpackOctets = Bytes.unpack
  pushOctetTop = flip Bytes.cons
  pushOctetEnd = Bytes.snoc
  spanOctetsTop = Bytes.span
  spanOctetsEnd = Bytes.spanEnd
  splitOffset = Bytes.splitAt
  replicateOctet k = Bytes.replicate $ fromIntegral k

-- TODO: we can improve the proformance of unpackBits by examining the current
-- bytestring, and loading a range of 1-8 bytes to OR simulaneously

-- | Implements binary decoder for LazyByteString, best for very large data sets
-- that can be loaded in and out of memory on demand.
-- https://hackage-content.haskell.org/package/bytestring-0.12.2.0/docs/Data-ByteString-Lazy.html
instance Transcoder LazyByteString where
  totalOctets = fromIntegral . LBS.length
  swapOrder = LBS.reverse
  execBuilder = BSB.toLazyByteString
  packOctets = LBS.pack
  unpackOctets = LBS.unpack
  unpackValueLE =
    -- From MSB to LSB, OR the current byte and shift the total bits.
    LBS.foldr (\byte bits -> (bits .<<. 8) .|. fromIntegral byte) 0
      . LBS.dropWhileEnd (== 0)
  unpackValueBE =
    -- From LSB to MSB, OR the current byte and shift the total bits.
    LBS.foldl' (\bits byte -> (bits .<<. 8) .|. fromIntegral byte) 0
      . LBS.dropWhile (== 0)
  unpackBuilder = Builder.lazyByteString
  pushOctetTop = flip LBS.cons
  pushOctetEnd = LBS.snoc
  spanOctetsTop = LBS.span
  spanOctetsEnd = LBS.spanEnd
  splitOffset k = LBS.splitAt $ fromIntegral k
  replicateOctet k = LBS.replicate $ fromIntegral k
  unpackValueCPU = unpackValueLE

-- | Implements binary decoder for ShortByteString, best for compact data that
-- doesn't exist long in memory. Notably, it packs better in memory than normal
-- ByteString (prone to Heap Fragmentation).
-- https://hackage-content.haskell.org/package/bytestring-0.12.2.0/docs/Data-ByteString-Short.html#g:1
instance Transcoder ShortByteString where
  totalOctets = SBS.length
  swapOrder = SBS.reverse
  execBuilder = SBS.toShort . execBuilder
  packOctets = SBS.pack
  unpackOctets = SBS.unpack
  unpackValueLE =
    -- From MSB to LSB , OR the current byte and shift the total bits.
    SBS.foldr (\byte bits -> (bits .<<. 8) .|. fromIntegral byte) 0
      . SBS.dropWhileEnd (== 0)
  unpackValueBE =
    -- From LSB to MSB, OR the current byte and shift the total bits.
    SBS.foldl' (\bits byte -> (bits .<<. 8) .|. fromIntegral byte) 0
      . SBS.dropWhile (== 0)
  unpackBuilder = Builder.shortByteString
  pushOctetTop = flip SBS.cons
  pushOctetEnd = SBS.snoc
  spanOctetsTop = SBS.span
  spanOctetsEnd = SBS.spanEnd
  splitOffset = SBS.splitAt
  replicateOctet k = SBS.replicate $ fromIntegral k
  unpackValueCPU = unpackValueLE

-- | Takes a subsection of octets from the beginning of the transcoder data.
takeOctetsTop :: (Transcoder tc) => Int -> tc -> tc
takeOctetsTop idx = fst . splitOffset idx

-- | Removes octets from the beginning of transcoder data.
dropOctetsTop :: (Transcoder tc) => Int -> tc -> tc
dropOctetsTop idx = snd . splitOffset idx

-- | Removes octets from the end of transcoder data.
takeOctetsEnd :: (Transcoder tc) => Int -> tc -> tc
takeOctetsEnd idx tc = snd $ splitOffset (totalOctets tc - idx) tc

-- | Removes octets from the end of transcoder data.
dropOctetsEnd :: (Transcoder tc) => Int -> tc -> tc
dropOctetsEnd idx tc = fst $ splitOffset (totalOctets tc - idx) tc

-- | Implements `takeWhile` for byte data in binary transcoder.
takeWhileTopOctets :: (Transcoder tc) => (Octet -> Bool) -> tc -> tc
takeWhileTopOctets f = fst . spanOctetsEnd f

-- | Implements `dropWhile` for byte data in binary transcoder.
dropWhileTopOctets :: (Transcoder tc) => (Octet -> Bool) -> tc -> tc
dropWhileTopOctets f = snd . spanOctetsEnd f

-- | Implements `takeWhileEnd` for byte data in binary transcoder.
takeWhileEndOctets :: (Transcoder tc) => (Octet -> Bool) -> tc -> tc
takeWhileEndOctets f = snd . spanOctetsEnd f

-- | Implements `dropWhileEnd` for byte data in binary transcoder.
dropWhileEndOctets :: (Transcoder tc) => (Octet -> Bool) -> tc -> tc
dropWhileEndOctets f = fst . spanOctetsEnd f

-- | Creates a binary transcoder from a single byte.
{-# INLINE fromOctet #-}
fromOctet :: (Transcoder tc) => Octet -> tc
fromOctet w8 = packOctets [w8]

-- | Returns True if the binary transcoder data is empty.
{-# INLINE null #-}
null :: (Transcoder tc) => tc -> Bool
null tc = totalOctets tc == 0

-- | Returns the first byte removed from the transcoder data, if available.
{-# INLINE uncons #-}
uncons :: (Transcoder tc) => tc -> Maybe (Octet, tc)
uncons tc =
  let (start, rest) = splitOffset 1 tc
   in case unpackOctets start of
        [top] -> Just (top, rest)
        _ -> Nothing

-- | Returns the last byte removed from the transcoder data, if available.
{-# INLINE unsnoc #-}
unsnoc :: (Transcoder tc) => tc -> Maybe (tc, Octet)
unsnoc tc =
  let offset = totalOctets tc - 1
      (rest, taken) = splitOffset offset tc
   in case unpackOctets taken of
        [final] -> Just (rest, final)
        _ -> Nothing

-- | Packs a bit string value into some bytes data `tc`.
{-# INLINE packValueOrder #-}
packValueOrder ::
  (Integral val, Bits val, Transcoder tc) =>
  ([Octet] -> [Octet]) ->
  val ->
  tc
packValueOrder setorder intdata =
  -- This code can be optimized; if we precalculate the length of the BitSet
  -- in bits we can mask from the other direction using left shifts, and
  -- remove the call to setorder, maybe we can even use a list comprehension to
  -- build in place for packOctets
  packOctets
    . replaceNull [0]
    . setorder
    . map (\(_, byte) -> fromIntegral byte)
    . takeWhile (\(rest, byte) -> rest > 0 || byte > 0)
    $ iterateInit (\(input, _) -> extractByte input) (,0) intdata
  where
    -- Gets the byte from LSB in input and shifts the value 8 bits.
    extractByte input = (input .>>. 8, input .&. 0xFF)

-- | Packs a bit string value into some bytes data `tc` in Big-Endian format.
packValueBE :: (Integral val, Bits val, Transcoder tc) => val -> tc
packValueBE = packValueOrder reverse

-- | Packs a bit string value into some bytes data `tc` in Little-Endian
-- format.
packValueLE :: (Integral val, Bits val, Transcoder tc) => val -> tc
packValueLE = packValueOrder id

-- | Return the the first param if the second param is empty.
{-# INLINE replaceNull #-}
replaceNull :: [a] -> [a] -> [a]
replaceNull replace xs
  | List.null xs = replace
  | otherwise = xs

-- | Wraps iterate to compose with an intermediary type constructor.
iterateInit :: (b -> b) -> (a -> b) -> a -> [b]
iterateInit f iterinit = drop 1 . iterate f . iterinit