{-# LANGUAGE UndecidableInstances #-} -- required below GHC 9.6

module Binrep.Put where

import Binrep.BLen ( BLen(blen) )
import Binrep.CBLen ( IsCBLen(CBLen), cblen )
import Bytezap.Poke
import Raehik.Compat.Data.Primitive.Types ( Prim', sizeOf )
import Binrep.Util.ByteOrder
import Raehik.Compat.Data.Primitive.Types.Endian ( ByteSwap )
import Binrep.Common.Via.Prim ( ViaPrim(..) )

import Data.ByteString qualified as B

import Binrep.Common.Class.TypeErrors ( ENoSum, ENoEmpty )
import GHC.TypeLits ( TypeError, KnownNat )

import GHC.Generics
import Generic.Data.Function.FoldMap
import Generic.Type.Assert

import Control.Monad.ST ( RealWorld )

import Binrep.Put.Struct ( PutC(putC) )

import Refined
import Refined.Unsafe

import Data.Word
import Data.Int
import Data.Void
import Data.Functor.Identity
import Binrep.Common.Via.Generically.NonSum

type Putter = Poke RealWorld
class Put a where put :: a -> Putter

runPut :: (BLen a, Put a) => a -> B.ByteString
runPut :: forall a. (BLen a, Put a) => a -> ByteString
runPut a
a = Int -> Poke RealWorld -> ByteString
unsafeRunPokeBS (a -> Int
forall a. BLen a => a -> Int
blen a
a) (a -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put a
a)

-- | Serialize generically using generic 'foldMap'.
instance GenericFoldMap Put where
    type GenericFoldMapM Put = Putter
    type GenericFoldMapC Put a = Put a
    genericFoldMapF :: forall a. GenericFoldMapC Put a => a -> GenericFoldMapM Put
genericFoldMapF = a -> Poke RealWorld
a -> GenericFoldMapM Put
forall a. Put a => a -> Poke RealWorld
put

-- | Serialize a term of the non-sum type @a@ via its 'Generic' instance.
putGenericNonSum
    :: forall a
    .  ( Generic a, GFoldMapNonSum Put (Rep a)
       , GAssertNotVoid a, GAssertNotSum a
    ) => a -> Putter
putGenericNonSum :: forall a.
(Generic a, GFoldMapNonSum Put (Rep a), GAssertNotVoid a,
 GAssertNotSum a) =>
a -> Poke RealWorld
putGenericNonSum = forall {k} (tag :: k) a.
(Generic a, GFoldMapNonSum tag (Rep a)) =>
a -> GenericFoldMapM tag
forall (tag :: Type -> Constraint) a.
(Generic a, GFoldMapNonSum tag (Rep a)) =>
a -> GenericFoldMapM tag
genericFoldMapNonSum @Put

instance
  ( Generic a, GFoldMapNonSum Put (Rep a)
  , GAssertNotVoid a, GAssertNotSum a
  ) => Put (GenericallyNonSum a) where
    put :: GenericallyNonSum a -> Poke RealWorld
put = a -> Poke RealWorld
forall a.
(Generic a, GFoldMapNonSum Put (Rep a), GAssertNotVoid a,
 GAssertNotSum a) =>
a -> Poke RealWorld
putGenericNonSum (a -> Poke RealWorld)
-> (GenericallyNonSum a -> a)
-> GenericallyNonSum a
-> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenericallyNonSum a -> a
forall a. GenericallyNonSum a -> a
unGenericallyNonSum

-- | Serialize a term of the sum type @a@ via its 'Generic' instance.
--
-- You must provide a serializer for @a@'s constructors. This is regrettably
-- inefficient due to having to use 'String's. Alas. Do write your own instance
-- if you want better performance!
putGenericSum
    :: forall a
    .  ( Generic a, GFoldMapSum Put (Rep a)
       , GAssertNotVoid a, GAssertSum a
    ) => (String -> Putter) -> a -> Putter
putGenericSum :: forall a.
(Generic a, GFoldMapSum Put (Rep a), GAssertNotVoid a,
 GAssertSum a) =>
(String -> Poke RealWorld) -> a -> Poke RealWorld
putGenericSum = forall {k} (tag :: k) a.
(Generic a, GFoldMapSum tag (Rep a)) =>
(String -> GenericFoldMapM tag) -> a -> GenericFoldMapM tag
forall (tag :: Type -> Constraint) a.
(Generic a, GFoldMapSum tag (Rep a)) =>
(String -> GenericFoldMapM tag) -> a -> GenericFoldMapM tag
genericFoldMapSum @Put

newtype ViaPutC a = ViaPutC { forall a. ViaPutC a -> a
unViaPutC :: a }
instance (PutC a, KnownNat (CBLen a)) => Put (ViaPutC a) where
    {-# INLINE put #-}
    put :: ViaPutC a -> Poke RealWorld
put = Int -> Poke RealWorld -> Poke RealWorld
forall s. Int -> Poke s -> Poke s
fromStructPoke (forall a. KnownNat (CBLen a) => Int
forall {k} (a :: k). KnownNat (CBLen a) => Int
cblen @a) (Poke RealWorld -> Poke RealWorld)
-> (ViaPutC a -> Poke RealWorld) -> ViaPutC a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Poke RealWorld
forall a. PutC a => a -> Poke RealWorld
putC (a -> Poke RealWorld)
-> (ViaPutC a -> a) -> ViaPutC a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ViaPutC a -> a
forall a. ViaPutC a -> a
unViaPutC

-- use ViaPutC over this, but should be semantically identical
instance Prim' a => Put (ViaPrim a) where
    put :: ViaPrim a -> Poke RealWorld
put = Int -> Poke RealWorld -> Poke RealWorld
forall s. Int -> Poke s -> Poke s
fromStructPoke (a -> Int
forall a. Prim a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) (Poke RealWorld -> Poke RealWorld)
-> (ViaPrim a -> Poke RealWorld) -> ViaPrim a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ViaPrim a -> Poke RealWorld
forall a. PutC a => a -> Poke RealWorld
putC
    {-# INLINE put #-}

instance TypeError ENoEmpty => Put Void where put :: Void -> Poke RealWorld
put = Void -> Poke RealWorld
forall a. HasCallStack => a
undefined
instance TypeError ENoSum => Put (Either a b) where put :: Either a b -> Poke RealWorld
put = Either a b -> Poke RealWorld
forall a. HasCallStack => a
undefined

instance Put a => Put (Identity a) where put :: Identity a -> Poke RealWorld
put = a -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put (a -> Poke RealWorld)
-> (Identity a -> a) -> Identity a -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity a -> a
forall a. Identity a -> a
runIdentity

instance Put Putter where put :: Poke RealWorld -> Poke RealWorld
put = Poke RealWorld -> Poke RealWorld
forall a. a -> a
id

-- | Unit type serializes to nothing. How zen.
instance Put () where
    {-# INLINE put #-}
    put :: () -> Poke RealWorld
put = () -> Poke RealWorld
forall a. Monoid a => a
mempty

instance (Put l, Put r) => Put (l, r) where
    {-# INLINE put #-}
    put :: (l, r) -> Poke RealWorld
put (l
l, r
r) = l -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put l
l Poke RealWorld -> Poke RealWorld -> Poke RealWorld
forall a. Semigroup a => a -> a -> a
<> r -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put r
r

instance Put a => Put [a] where
    {-# INLINE put #-}
    put :: [a] -> Poke RealWorld
put = [Poke RealWorld] -> Poke RealWorld
forall a. Monoid a => [a] -> a
mconcat ([Poke RealWorld] -> Poke RealWorld)
-> ([a] -> [Poke RealWorld]) -> [a] -> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Poke RealWorld) -> [a] -> [Poke RealWorld]
forall a b. (a -> b) -> [a] -> [b]
map a -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put

instance Put B.ByteString where
    {-# INLINE put #-}
    put :: ByteString -> Poke RealWorld
put = ByteString -> Poke RealWorld
byteString

-- | 8-bit (1-byte) words do not require byte order in order to precisely
--   define their representation.
deriving via ViaPutC Word8 instance Put Word8

-- | 8-bit (1-byte) words do not require byte order in order to precisely
--   define their representation.
deriving via ViaPutC  Int8 instance Put  Int8

-- | Byte order is irrelevant for 8-bit (1-byte) words.
deriving via Word8 instance Put (ByteOrdered end Word8)

-- | Byte order is irrelevant for 8-bit (1-byte) words.
deriving via  Int8 instance Put (ByteOrdered end  Int8)

-- ByteSwap is required on opposite endian platforms, but we're not checking
-- here, so make sure to keep it on both.
-- Stick with ViaPrim here because ByteOrdered is connected to it.
deriving via ViaPrim (ByteOrdered 'LittleEndian a)
    instance (Prim' a, ByteSwap a) => Put (ByteOrdered 'LittleEndian a)
deriving via ViaPrim (ByteOrdered    'BigEndian a)
    instance (Prim' a, ByteSwap a) => Put (ByteOrdered    'BigEndian a)

-- | Put types refined with multiple predicates by wrapping the left
--   predicate with the right. LOL REALLY?
instance Put (Refined pr (Refined pl a)) => Put (Refined (pl `And` pr) a) where
    put :: Refined (And pl pr) a -> Poke RealWorld
put = Refined pr (Refined pl a) -> Poke RealWorld
forall a. Put a => a -> Poke RealWorld
put (Refined pr (Refined pl a) -> Poke RealWorld)
-> (Refined (And pl pr) a -> Refined pr (Refined pl a))
-> Refined (And pl pr) a
-> Poke RealWorld
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x (p :: k1). x -> Refined p x
forall {k} x (p :: k). x -> Refined p x
reallyUnsafeRefine @_ @pr (Refined pl a -> Refined pr (Refined pl a))
-> (Refined (And pl pr) a -> Refined pl a)
-> Refined (And pl pr) a
-> Refined pr (Refined pl a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x (p :: k). x -> Refined p x
forall {k} x (p :: k). x -> Refined p x
reallyUnsafeRefine @_ @pl (a -> Refined pl a)
-> (Refined (And pl pr) a -> a)
-> Refined (And pl pr) a
-> Refined pl a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Refined (And pl pr) a -> a
forall {k} (p :: k) x. Refined p x -> x
unrefine