module Dahdit.Binary
  ( Binary (..)
  , byteSizeFoldable
  )
where

import Dahdit.Free (Get, Put)
import Dahdit.Funs
  ( getDoubleBE
  , getDoubleLE
  , getFloatBE
  , getFloatLE
  , getInt16BE
  , getInt16LE
  , getInt24BE
  , getInt24LE
  , getInt32BE
  , getInt32LE
  , getInt64BE
  , getInt64LE
  , getInt8
  , getList
  , getSeq
  , getWord16BE
  , getWord16LE
  , getWord24BE
  , getWord24LE
  , getWord32BE
  , getWord32LE
  , getWord64BE
  , getWord64LE
  , getWord8
  , putDoubleBE
  , putDoubleLE
  , putFloatBE
  , putFloatLE
  , putInt16BE
  , putInt16LE
  , putInt24BE
  , putInt24LE
  , putInt32BE
  , putInt32LE
  , putInt64BE
  , putInt64LE
  , putInt8
  , putList
  , putSeq
  , putWord16BE
  , putWord16LE
  , putWord24BE
  , putWord24LE
  , putWord32BE
  , putWord32LE
  , putWord64BE
  , putWord64LE
  , putWord8
  )
import Dahdit.Nums
  ( DoubleBE (..)
  , DoubleLE (..)
  , FloatBE (..)
  , FloatLE (..)
  , Int16BE (..)
  , Int16LE (..)
  , Int24BE (..)
  , Int24LE (..)
  , Int32BE (..)
  , Int32LE (..)
  , Int64BE (..)
  , Int64LE (..)
  , Word16BE (..)
  , Word16LE (..)
  , Word24BE (..)
  , Word24LE (..)
  , Word32BE (..)
  , Word32LE (..)
  , Word64BE (..)
  , Word64LE (..)
  )
import Dahdit.Run (runCount)
import Dahdit.Sizes (ByteCount (..), ElemCount (..))
import Data.ByteString.Internal (c2w, w2c)
import Data.Coerce (coerce)
import Data.Foldable (foldMap')
import Data.Int (Int16, Int32, Int64, Int8)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Semigroup (Sum (..))
import Data.Sequence (Seq (..))
import qualified Data.Sequence as Seq
import Data.Set (Set)
import qualified Data.Set as Set
import Data.ShortWord (Int24, Word24)
import Data.Word (Word16, Word32, Word64, Word8)

class Binary a where
  byteSize :: a -> ByteCount
  byteSize = Put -> ByteCount
runCount forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Binary a => a -> Put
put
  get :: Get a
  put :: a -> Put

-- Basic types

instance Binary () where
  byteSize :: () -> ByteCount
byteSize ()
_ = ByteCount
0
  get :: Get ()
get = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  put :: () -> Put
put ()
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

instance Binary Word8 where
  byteSize :: Word8 -> ByteCount
byteSize Word8
_ = ByteCount
1
  get :: Get Word8
get = Get Word8
getWord8
  put :: Word8 -> Put
put = Word8 -> Put
putWord8

instance Binary Int8 where
  byteSize :: Int8 -> ByteCount
byteSize Int8
_ = ByteCount
1
  get :: Get Int8
get = Get Int8
getInt8
  put :: Int8 -> Put
put = Int8 -> Put
putInt8

instance Binary Word16LE where
  byteSize :: Word16LE -> ByteCount
byteSize Word16LE
_ = ByteCount
2
  get :: Get Word16LE
get = Get Word16LE
getWord16LE
  put :: Word16LE -> Put
put = Word16LE -> Put
putWord16LE

instance Binary Int16LE where
  byteSize :: Int16LE -> ByteCount
byteSize Int16LE
_ = ByteCount
2
  get :: Get Int16LE
get = Get Int16LE
getInt16LE
  put :: Int16LE -> Put
put = Int16LE -> Put
putInt16LE

instance Binary Word24LE where
  byteSize :: Word24LE -> ByteCount
byteSize Word24LE
_ = ByteCount
3
  get :: Get Word24LE
get = Get Word24LE
getWord24LE
  put :: Word24LE -> Put
put = Word24LE -> Put
putWord24LE

instance Binary Int24LE where
  byteSize :: Int24LE -> ByteCount
byteSize Int24LE
_ = ByteCount
3
  get :: Get Int24LE
get = Get Int24LE
getInt24LE
  put :: Int24LE -> Put
put = Int24LE -> Put
putInt24LE

instance Binary Word32LE where
  byteSize :: Word32LE -> ByteCount
byteSize Word32LE
_ = ByteCount
4
  get :: Get Word32LE
get = Get Word32LE
getWord32LE
  put :: Word32LE -> Put
put = Word32LE -> Put
putWord32LE

instance Binary Int32LE where
  byteSize :: Int32LE -> ByteCount
byteSize Int32LE
_ = ByteCount
4
  get :: Get Int32LE
get = Get Int32LE
getInt32LE
  put :: Int32LE -> Put
put = Int32LE -> Put
putInt32LE

instance Binary Word64LE where
  byteSize :: Word64LE -> ByteCount
byteSize Word64LE
_ = ByteCount
8
  get :: Get Word64LE
get = Get Word64LE
getWord64LE
  put :: Word64LE -> Put
put = Word64LE -> Put
putWord64LE

instance Binary Int64LE where
  byteSize :: Int64LE -> ByteCount
byteSize Int64LE
_ = ByteCount
8
  get :: Get Int64LE
get = Get Int64LE
getInt64LE
  put :: Int64LE -> Put
put = Int64LE -> Put
putInt64LE

instance Binary FloatLE where
  byteSize :: FloatLE -> ByteCount
byteSize FloatLE
_ = ByteCount
4
  get :: Get FloatLE
get = Get FloatLE
getFloatLE
  put :: FloatLE -> Put
put = FloatLE -> Put
putFloatLE

instance Binary DoubleLE where
  byteSize :: DoubleLE -> ByteCount
byteSize DoubleLE
_ = ByteCount
8
  get :: Get DoubleLE
get = Get DoubleLE
getDoubleLE
  put :: DoubleLE -> Put
put = DoubleLE -> Put
putDoubleLE

instance Binary Word16BE where
  byteSize :: Word16BE -> ByteCount
byteSize Word16BE
_ = ByteCount
2
  get :: Get Word16BE
get = Get Word16BE
getWord16BE
  put :: Word16BE -> Put
put = Word16BE -> Put
putWord16BE

instance Binary Int16BE where
  byteSize :: Int16BE -> ByteCount
byteSize Int16BE
_ = ByteCount
2
  get :: Get Int16BE
get = Get Int16BE
getInt16BE
  put :: Int16BE -> Put
put = Int16BE -> Put
putInt16BE

instance Binary Word24BE where
  byteSize :: Word24BE -> ByteCount
byteSize Word24BE
_ = ByteCount
3
  get :: Get Word24BE
get = Get Word24BE
getWord24BE
  put :: Word24BE -> Put
put = Word24BE -> Put
putWord24BE

instance Binary Int24BE where
  byteSize :: Int24BE -> ByteCount
byteSize Int24BE
_ = ByteCount
3
  get :: Get Int24BE
get = Get Int24BE
getInt24BE
  put :: Int24BE -> Put
put = Int24BE -> Put
putInt24BE

instance Binary Word32BE where
  byteSize :: Word32BE -> ByteCount
byteSize Word32BE
_ = ByteCount
4
  get :: Get Word32BE
get = Get Word32BE
getWord32BE
  put :: Word32BE -> Put
put = Word32BE -> Put
putWord32BE

instance Binary Int32BE where
  byteSize :: Int32BE -> ByteCount
byteSize Int32BE
_ = ByteCount
4
  get :: Get Int32BE
get = Get Int32BE
getInt32BE
  put :: Int32BE -> Put
put = Int32BE -> Put
putInt32BE

instance Binary Word64BE where
  byteSize :: Word64BE -> ByteCount
byteSize Word64BE
_ = ByteCount
8
  get :: Get Word64BE
get = Get Word64BE
getWord64BE
  put :: Word64BE -> Put
put = Word64BE -> Put
putWord64BE

instance Binary Int64BE where
  byteSize :: Int64BE -> ByteCount
byteSize Int64BE
_ = ByteCount
8
  get :: Get Int64BE
get = Get Int64BE
getInt64BE
  put :: Int64BE -> Put
put = Int64BE -> Put
putInt64BE

instance Binary FloatBE where
  byteSize :: FloatBE -> ByteCount
byteSize FloatBE
_ = ByteCount
4
  get :: Get FloatBE
get = Get FloatBE
getFloatBE
  put :: FloatBE -> Put
put = FloatBE -> Put
putFloatBE

instance Binary DoubleBE where
  byteSize :: DoubleBE -> ByteCount
byteSize DoubleBE
_ = ByteCount
8
  get :: Get DoubleBE
get = Get DoubleBE
getDoubleBE
  put :: DoubleBE -> Put
put = DoubleBE -> Put
putDoubleBE

deriving via Word16LE instance Binary Word16

deriving via Int16LE instance Binary Int16

deriving via Word24LE instance Binary Word24

deriving via Int24LE instance Binary Int24

deriving via Word32LE instance Binary Word32

deriving via Int32LE instance Binary Int32

deriving via Word64LE instance Binary Word64

deriving via Int64LE instance Binary Int64

deriving via FloatLE instance Binary Float

deriving via DoubleLE instance Binary Double

instance Binary Bool where
  byteSize :: Bool -> ByteCount
byteSize Bool
_ = ByteCount
1
  get :: Get Bool
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Eq a => a -> a -> Bool
/= Word8
0) Get Word8
getWord8
  put :: Bool -> Put
put Bool
b = Word8 -> Put
putWord8 (if Bool
b then Word8
1 else Word8
0)

instance Binary Char where
  byteSize :: Char -> ByteCount
byteSize Char
_ = ByteCount
1
  get :: Get Char
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Word8 -> Char
w2c Get Word8
getWord8
  put :: Char -> Put
put = Word8 -> Put
putWord8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Word8
c2w

instance Binary Int where
  byteSize :: Int -> ByteCount
byteSize Int
_ = ByteCount
8
  get :: Get Int
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (Integral a, Num b) => a -> b
fromIntegral Get Int64LE
getInt64LE
  put :: Int -> Put
put = Int64LE -> Put
putInt64LE forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral

instance (Binary a) => Binary [a] where
  byteSize :: [a] -> ByteCount
byteSize [a]
as = ByteCount
8 forall a. Num a => a -> a -> a
+ forall (f :: * -> *) a. (Foldable f, Binary a) => f a -> ByteCount
byteSizeFoldable [a]
as
  get :: Get [a]
get = do
    Int
ec <- forall a. Binary a => Get a
get @Int
    forall a. ElemCount -> Get a -> Get [a]
getList (coerce :: forall a b. Coercible a b => a -> b
coerce Int
ec) forall a. Binary a => Get a
get
  put :: [a] -> Put
put [a]
s = forall a. Binary a => a -> Put
put @Int (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s) forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. (a -> Put) -> [a] -> Put
putList forall a. Binary a => a -> Put
put [a]
s

instance (Binary a) => Binary (Seq a) where
  byteSize :: Seq a -> ByteCount
byteSize Seq a
as = ByteCount
8 forall a. Num a => a -> a -> a
+ forall (f :: * -> *) a. (Foldable f, Binary a) => f a -> ByteCount
byteSizeFoldable Seq a
as
  get :: Get (Seq a)
get = do
    Int
ec <- forall a. Binary a => Get a
get @Int
    forall a. ElemCount -> Get a -> Get (Seq a)
getSeq (coerce :: forall a b. Coercible a b => a -> b
coerce Int
ec) forall a. Binary a => Get a
get
  put :: Seq a -> Put
put Seq a
s = forall a. Binary a => a -> Put
put @Int (forall a. Seq a -> Int
Seq.length Seq a
s) forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. (a -> Put) -> Seq a -> Put
putSeq forall a. Binary a => a -> Put
put Seq a
s

instance (Binary a) => Binary (Set a) where
  byteSize :: Set a -> ByteCount
byteSize = forall a. Binary a => a -> ByteCount
byteSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
Set.toAscList
  get :: Get (Set a)
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. [a] -> Set a
Set.fromDistinctAscList forall a. Binary a => Get a
get
  put :: Set a -> Put
put = forall a. Binary a => a -> Put
put forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
Set.toAscList

instance (Binary k, Binary v) => Binary (Map k v) where
  byteSize :: Map k v -> ByteCount
byteSize = forall a. Binary a => a -> ByteCount
byteSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toAscList
  get :: Get (Map k v)
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList forall a. Binary a => Get a
get
  put :: Map k v -> Put
put = forall a. Binary a => a -> Put
put forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toAscList

instance Binary IntSet where
  byteSize :: IntSet -> ByteCount
byteSize = forall a. Binary a => a -> ByteCount
byteSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> [Int]
IntSet.toAscList
  get :: Get IntSet
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Int] -> IntSet
IntSet.fromDistinctAscList forall a. Binary a => Get a
get
  put :: IntSet -> Put
put = forall a. Binary a => a -> Put
put forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> [Int]
IntSet.toAscList

instance (Binary v) => Binary (IntMap v) where
  byteSize :: IntMap v -> ByteCount
byteSize = forall a. Binary a => a -> ByteCount
byteSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntMap a -> [(Int, a)]
IntMap.toAscList
  get :: Get (IntMap v)
get = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. [(Int, a)] -> IntMap a
IntMap.fromDistinctAscList forall a. Binary a => Get a
get
  put :: IntMap v -> Put
put = forall a. Binary a => a -> Put
put forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntMap a -> [(Int, a)]
IntMap.toAscList

instance (Binary a) => Binary (Maybe a) where
  byteSize :: Maybe a -> ByteCount
byteSize = \case
    Maybe a
Nothing -> ByteCount
1
    Just a
a -> ByteCount
1 forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize a
a
  get :: Get (Maybe a)
get = do
    Int
tag <- forall a. Binary a => Get a
get @Int
    case Int
tag of
      Int
0 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
      Int
1 -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. a -> Maybe a
Just forall a. Binary a => Get a
get
      Int
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unknown encoding for constructor"
  put :: Maybe a -> Put
put = \case
    Maybe a
Nothing -> forall a. Binary a => a -> Put
put @Int Int
0
    Just a
a -> forall a. Binary a => a -> Put
put @Int Int
1 forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put a
a

instance (Binary b, Binary a) => Binary (Either b a) where
  byteSize :: Either b a -> ByteCount
byteSize = \case
    Left b
b -> ByteCount
1 forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize b
b
    Right a
a -> ByteCount
1 forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize a
a
  get :: Get (Either b a)
get = do
    Int
tag <- forall a. Binary a => Get a
get @Int
    case Int
tag of
      Int
0 -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. a -> Either a b
Left forall a. Binary a => Get a
get
      Int
1 -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. b -> Either a b
Right forall a. Binary a => Get a
get
      Int
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unknown encoding for constructor"
  put :: Either b a -> Put
put = \case
    Left b
b -> forall a. Binary a => a -> Put
put @Int Int
0 forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put b
b
    Right a
a -> forall a. Binary a => a -> Put
put @Int Int
1 forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put a
a

instance (Binary a, Binary b) => Binary (a, b) where
  byteSize :: (a, b) -> ByteCount
byteSize (a
a, b
b) = forall a. Binary a => a -> ByteCount
byteSize a
a forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize b
b
  get :: Get (a, b)
get = do
    a
a <- forall a. Binary a => Get a
get
    b
b <- forall a. Binary a => Get a
get
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, b
b)
  put :: (a, b) -> Put
put (a
a, b
b) = forall a. Binary a => a -> Put
put a
a forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put b
b

instance (Binary a, Binary b, Binary c) => Binary (a, b, c) where
  byteSize :: (a, b, c) -> ByteCount
byteSize (a
a, b
b, c
c) = forall a. Binary a => a -> ByteCount
byteSize a
a forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize b
b forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize c
c
  get :: Get (a, b, c)
get = do
    a
a <- forall a. Binary a => Get a
get
    b
b <- forall a. Binary a => Get a
get
    c
c <- forall a. Binary a => Get a
get
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, b
b, c
c)
  put :: (a, b, c) -> Put
put (a
a, b
b, c
c) = forall a. Binary a => a -> Put
put a
a forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put b
b forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put c
c

instance (Binary a, Binary b, Binary c, Binary d) => Binary (a, b, c, d) where
  byteSize :: (a, b, c, d) -> ByteCount
byteSize (a
a, b
b, c
c, d
d) = forall a. Binary a => a -> ByteCount
byteSize a
a forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize b
b forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize c
c forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize d
d
  get :: Get (a, b, c, d)
get = do
    a
a <- forall a. Binary a => Get a
get
    b
b <- forall a. Binary a => Get a
get
    c
c <- forall a. Binary a => Get a
get
    d
d <- forall a. Binary a => Get a
get
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, b
b, c
c, d
d)
  put :: (a, b, c, d) -> Put
put (a
a, b
b, c
c, d
d) = forall a. Binary a => a -> Put
put a
a forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put b
b forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put c
c forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put d
d

instance (Binary a, Binary b, Binary c, Binary d, Binary e) => Binary (a, b, c, d, e) where
  byteSize :: (a, b, c, d, e) -> ByteCount
byteSize (a
a, b
b, c
c, d
d, e
e) = forall a. Binary a => a -> ByteCount
byteSize a
a forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize b
b forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize c
c forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize d
d forall a. Num a => a -> a -> a
+ forall a. Binary a => a -> ByteCount
byteSize e
e
  get :: Get (a, b, c, d, e)
get = do
    a
a <- forall a. Binary a => Get a
get
    b
b <- forall a. Binary a => Get a
get
    c
c <- forall a. Binary a => Get a
get
    d
d <- forall a. Binary a => Get a
get
    e
e <- forall a. Binary a => Get a
get
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, b
b, c
c, d
d, e
e)
  put :: (a, b, c, d, e) -> Put
put (a
a, b
b, c
c, d
d, e
e) = forall a. Binary a => a -> Put
put a
a forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put b
b forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put c
c forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put d
d forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> forall a. Binary a => a -> Put
put e
e

byteSizeFoldable :: (Foldable f, Binary a) => f a -> ByteCount
byteSizeFoldable :: forall (f :: * -> *) a. (Foldable f, Binary a) => f a -> ByteCount
byteSizeFoldable = forall a. Sum a -> a
getSum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' (forall a. a -> Sum a
Sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Binary a => a -> ByteCount
byteSize)