{-# LANGUAGE TemplateHaskell #-}

-- This is required as all the instances in this module are orphan instances.
{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- Module      : Streamly.Internal.Data.MutByteArray
-- Copyright   : (c) 2023 Composewell Technologies
-- License     : BSD3-3-Clause
-- Maintainer  : streamly@composewell.com
-- Portability : GHC
--

module Streamly.Internal.Data.MutByteArray
    (
    -- * MutByteArray
      module Streamly.Internal.Data.MutByteArray.Type
    -- * Unbox
    , module Streamly.Internal.Data.Unbox
    , module Streamly.Internal.Data.Unbox.TH
    -- * Serialize
    , module Streamly.Internal.Data.Serialize.Type
    -- * Serialize TH
    , module Streamly.Internal.Data.Serialize.TH
    ) where

--------------------------------------------------------------------------------
-- Imports
--------------------------------------------------------------------------------

import Data.Proxy (Proxy(..))
import Streamly.Internal.Data.Array (Array(..))
import GHC.Exts (Int(..), sizeofByteArray#, unsafeCoerce#)
import GHC.Word (Word8)

#if __GLASGOW_HASKELL__ >= 900
import GHC.Num.Integer (Integer(..))
#else
import GHC.Integer.GMP.Internals (Integer(..), BigNat(..))
#endif

import Streamly.Internal.Data.MutByteArray.Type
import Streamly.Internal.Data.Serialize.TH
import Streamly.Internal.Data.Serialize.Type
import Streamly.Internal.Data.Unbox
import Streamly.Internal.Data.Unbox.TH

--------------------------------------------------------------------------------
-- Common instances
--------------------------------------------------------------------------------

-- Note
-- ====
--
-- Even a non-functional change such as changing the order of constructors will
-- change the instance derivation.
--
-- This will not pose a problem if both, encode, and decode are done by the same
-- version of the application. There *might* be a problem if version that
-- encodes differs from the version that decodes.
--
-- We need to add some compatibility tests using different versions of
-- dependencies.
--
-- Although such chages for the most basic types won't happen we need to detect
-- if it ever happens.
--
-- Should we worry about these kind of changes and this kind of compatibility?
-- This is a problem for all types of derivations that depend on the order of
-- constructors, for example, Enum.

-- Note on Windows build
-- =====================
--
-- On Windows, having template haskell splices here fail the build with the
-- following error:
--
-- @
-- addLibrarySearchPath: C:\...  (Win32 error 3): The system cannot find the path specified.
-- @
--
-- The error might be irrelavant but having these splices triggers it. We should
-- either fix the problem or avoid the use to template haskell splices in this
-- file.
--
-- Similar issue: https://github.com/haskell/cabal/issues/4741

-- $(Serialize.deriveSerialize ''Maybe)
instance Serialize a => Serialize (Maybe a) where

    {-# INLINE addSizeTo #-}
    addSizeTo :: Int -> Maybe a -> Int
addSizeTo Int
acc Maybe a
x =
        case Maybe a
x of
            Maybe a
Nothing -> (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Just a
field0 -> (Int -> a -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) a
field0

    {-# INLINE deserializeAt #-}
    deserializeAt :: Int -> MutByteArray -> Int -> IO (Int, Maybe a)
deserializeAt Int
initialOffset MutByteArray
arr Int
endOffset = do
        (Int
i0, Word8
tag) <- ((Int -> MutByteArray -> Int -> IO (Int, Word8)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
initialOffset) MutByteArray
arr) Int
endOffset
        case Word8
tag :: Word8 of
            Word8
0 -> (Int, Maybe a) -> IO (Int, Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i0, Maybe a
forall a. Maybe a
Nothing)
            Word8
1 -> do (Int
i1, a
a0) <- ((Int -> MutByteArray -> Int -> IO (Int, a)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
i0) MutByteArray
arr) Int
endOffset
                    (Int, Maybe a) -> IO (Int, Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i1, a -> Maybe a
forall a. a -> Maybe a
Just a
a0)
            Word8
_ -> [Char] -> IO (Int, Maybe a)
forall a. HasCallStack => [Char] -> a
error [Char]
"Found invalid tag while peeking (Maybe a)"

    {-# INLINE serializeAt #-}
    serializeAt :: Int -> MutByteArray -> Maybe a -> IO Int
serializeAt Int
initialOffset MutByteArray
arr Maybe a
val =
        case Maybe a
val of
            Maybe a
Nothing -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
0 :: Word8)
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i0
            Just a
field0 -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
1 :: Word8)
                Int
i1 <- ((Int -> MutByteArray -> a -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
i0) MutByteArray
arr) a
field0
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i1

-- $(Serialize.deriveSerialize ''Either)
instance (Serialize a, Serialize b) => Serialize (Either a b) where

    {-# INLINE addSizeTo #-}
    addSizeTo :: Int -> Either a b -> Int
addSizeTo Int
acc Either a b
x =
        case Either a b
x of
            Left a
field0 -> (Int -> a -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) a
field0
            Right b
field0 -> (Int -> b -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) b
field0

    {-# INLINE deserializeAt #-}
    deserializeAt :: Int -> MutByteArray -> Int -> IO (Int, Either a b)
deserializeAt Int
initialOffset MutByteArray
arr Int
endOffset = do
        (Int
i0, Word8
tag) <- ((Int -> MutByteArray -> Int -> IO (Int, Word8)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
initialOffset) MutByteArray
arr) Int
endOffset
        case Word8
tag :: Word8 of
            Word8
0 -> do (Int
i1, a
a0) <- ((Int -> MutByteArray -> Int -> IO (Int, a)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
i0) MutByteArray
arr) Int
endOffset
                    (Int, Either a b) -> IO (Int, Either a b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i1, a -> Either a b
forall a b. a -> Either a b
Left a
a0)
            Word8
1 -> do (Int
i1, b
a0) <- ((Int -> MutByteArray -> Int -> IO (Int, b)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
i0) MutByteArray
arr) Int
endOffset
                    (Int, Either a b) -> IO (Int, Either a b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i1, b -> Either a b
forall a b. b -> Either a b
Right b
a0)
            Word8
_ -> [Char] -> IO (Int, Either a b)
forall a. HasCallStack => [Char] -> a
error [Char]
"Found invalid tag while peeking (Either a b)"

    {-# INLINE serializeAt #-}
    serializeAt :: Int -> MutByteArray -> Either a b -> IO Int
serializeAt Int
initialOffset MutByteArray
arr Either a b
val =
        case Either a b
val of
            Left a
field0 -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
0 :: Word8)
                Int
i1 <- ((Int -> MutByteArray -> a -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
i0) MutByteArray
arr) a
field0
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i1
            Right b
field0 -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
1 :: Word8)
                Int
i1 <- ((Int -> MutByteArray -> b -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
i0) MutByteArray
arr) b
field0
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i1

instance Serialize (Proxy a) where

    {-# INLINE addSizeTo #-}
    addSizeTo :: Int -> Proxy a -> Int
addSizeTo Int
acc Proxy a
_ = (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

    {-# INLINE deserializeAt #-}
    deserializeAt :: Int -> MutByteArray -> Int -> IO (Int, Proxy a)
deserializeAt Int
initialOffset MutByteArray
_ Int
_ = (Int, Proxy a) -> IO (Int, Proxy a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int
initialOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1), Proxy a
forall {k} (t :: k). Proxy t
Proxy)

    {-# INLINE serializeAt #-}
    serializeAt :: Int -> MutByteArray -> Proxy a -> IO Int
serializeAt Int
initialOffset MutByteArray
_ Proxy a
_ = Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
initialOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

--------------------------------------------------------------------------------
-- Integer
--------------------------------------------------------------------------------

data LiftedInteger
    = LIS Int
    | LIP (Array Word)
    | LIN (Array Word)

-- $(Serialize.deriveSerialize ''LiftedInteger)
instance Serialize LiftedInteger where

    {-# INLINE addSizeTo #-}
    addSizeTo :: Int -> LiftedInteger -> Int
addSizeTo Int
acc LiftedInteger
x =
        case LiftedInteger
x of
            LIS Int
field0 -> (Int -> Int -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Int
field0
            LIP Array Word
field0 -> (Int -> Array Word -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Array Word
field0
            LIN Array Word
field0 -> (Int -> Array Word -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Array Word
field0

    {-# INLINE deserializeAt #-}
    deserializeAt :: Int -> MutByteArray -> Int -> IO (Int, LiftedInteger)
deserializeAt Int
initialOffset MutByteArray
arr Int
endOffset = do
        (Int
i0, Word8
tag) <- ((Int -> MutByteArray -> Int -> IO (Int, Word8)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
initialOffset) MutByteArray
arr) Int
endOffset
        case Word8
tag :: Word8 of
            Word8
0 -> do (Int
i1, Int
a0) <- ((Int -> MutByteArray -> Int -> IO (Int, Int)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
i0) MutByteArray
arr) Int
endOffset
                    (Int, LiftedInteger) -> IO (Int, LiftedInteger)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i1, Int -> LiftedInteger
LIS Int
a0)
            Word8
1 -> do (Int
i1, Array Word
a0) <- ((Int -> MutByteArray -> Int -> IO (Int, Array Word)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
i0) MutByteArray
arr) Int
endOffset
                    (Int, LiftedInteger) -> IO (Int, LiftedInteger)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i1, Array Word -> LiftedInteger
LIP Array Word
a0)
            Word8
2 -> do (Int
i1, Array Word
a0) <- ((Int -> MutByteArray -> Int -> IO (Int, Array Word)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
i0) MutByteArray
arr) Int
endOffset
                    (Int, LiftedInteger) -> IO (Int, LiftedInteger)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
i1, Array Word -> LiftedInteger
LIN Array Word
a0)
            Word8
_ -> [Char] -> IO (Int, LiftedInteger)
forall a. HasCallStack => [Char] -> a
error [Char]
"Found invalid tag while peeking (LiftedInteger)"

    {-# INLINE serializeAt #-}
    serializeAt :: Int -> MutByteArray -> LiftedInteger -> IO Int
serializeAt Int
initialOffset MutByteArray
arr LiftedInteger
val =
        case LiftedInteger
val of
            LIS Int
field0 -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
0 :: Word8)
                Int
i1 <- ((Int -> MutByteArray -> Int -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
i0) MutByteArray
arr) Int
field0
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i1
            LIP Array Word
field0 -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
1 :: Word8)
                Int
i1 <- ((Int -> MutByteArray -> Array Word -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
i0) MutByteArray
arr) Array Word
field0
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i1
            LIN Array Word
field0 -> do
                Int
i0 <- ((Int -> MutByteArray -> Word8 -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
initialOffset) MutByteArray
arr) (Word8
2 :: Word8)
                Int
i1 <- ((Int -> MutByteArray -> Array Word -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
i0) MutByteArray
arr) Array Word
field0
                Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
i1

#if __GLASGOW_HASKELL__ >= 900

{-# INLINE liftInteger #-}
liftInteger :: Integer -> LiftedInteger
liftInteger :: Integer -> LiftedInteger
liftInteger (IS Int#
x) = Int -> LiftedInteger
LIS (Int# -> Int
I# Int#
x)
liftInteger (IP ByteArray#
x) =
    Array Word -> LiftedInteger
LIP (MutByteArray -> Int -> Int -> Array Word
forall a. MutByteArray -> Int -> Int -> Array a
Array (MutableByteArray# RealWorld -> MutByteArray
MutByteArray (ByteArray# -> MutableByteArray# RealWorld
forall a b. a -> b
unsafeCoerce# ByteArray#
x)) Int
0 (Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
x)))
liftInteger (IN ByteArray#
x) =
    Array Word -> LiftedInteger
LIN (MutByteArray -> Int -> Int -> Array Word
forall a. MutByteArray -> Int -> Int -> Array a
Array (MutableByteArray# RealWorld -> MutByteArray
MutByteArray (ByteArray# -> MutableByteArray# RealWorld
forall a b. a -> b
unsafeCoerce# ByteArray#
x)) Int
0 (Int# -> Int
I# (ByteArray# -> Int#
sizeofByteArray# ByteArray#
x)))

{-# INLINE unliftInteger #-}
unliftInteger :: LiftedInteger -> Integer
unliftInteger :: LiftedInteger -> Integer
unliftInteger (LIS (I# Int#
x)) = Int# -> Integer
IS Int#
x
unliftInteger (LIP (Array (MutByteArray MutableByteArray# RealWorld
x) Int
_ Int
_)) = ByteArray# -> Integer
IP (MutableByteArray# RealWorld -> ByteArray#
forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
x)
unliftInteger (LIN (Array (MutByteArray MutableByteArray# RealWorld
x) Int
_ Int
_)) = ByteArray# -> Integer
IN (MutableByteArray# RealWorld -> ByteArray#
forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
x)

#else

{-# INLINE liftInteger #-}
liftInteger :: Integer -> LiftedInteger
liftInteger (S# x) = LIS (I# x)
liftInteger (Jp# (BN# x)) =
    LIP (Array (MutByteArray (unsafeCoerce# x)) 0 (I# (sizeofByteArray# x)))
liftInteger (Jn# (BN# x)) =
    LIN (Array (MutByteArray (unsafeCoerce# x)) 0 (I# (sizeofByteArray# x)))

{-# INLINE unliftInteger #-}
unliftInteger :: LiftedInteger -> Integer
unliftInteger (LIS (I# x)) = S# x
unliftInteger (LIP (Array (MutByteArray x) _ _)) =
    Jp# (BN# (unsafeCoerce# x))
unliftInteger (LIN (Array (MutByteArray x) _ _)) =
    Jn# (BN# (unsafeCoerce# x))

#endif

instance Serialize Integer where
    {-# INLINE addSizeTo #-}
    addSizeTo :: Int -> Integer -> Int
addSizeTo Int
i Integer
a = Int -> LiftedInteger -> Int
forall a. Serialize a => Int -> a -> Int
addSizeTo Int
i (Integer -> LiftedInteger
liftInteger Integer
a)

    {-# INLINE deserializeAt #-}
    deserializeAt :: Int -> MutByteArray -> Int -> IO (Int, Integer)
deserializeAt Int
off MutByteArray
arr Int
end =
        (LiftedInteger -> Integer)
-> (Int, LiftedInteger) -> (Int, Integer)
forall a b. (a -> b) -> (Int, a) -> (Int, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LiftedInteger -> Integer
unliftInteger ((Int, LiftedInteger) -> (Int, Integer))
-> IO (Int, LiftedInteger) -> IO (Int, Integer)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> MutByteArray -> Int -> IO (Int, LiftedInteger)
forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
off MutByteArray
arr Int
end

    {-# INLINE serializeAt #-}
    serializeAt :: Int -> MutByteArray -> Integer -> IO Int
serializeAt Int
off MutByteArray
arr Integer
val = Int -> MutByteArray -> LiftedInteger -> IO Int
forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
off MutByteArray
arr (Integer -> LiftedInteger
liftInteger Integer
val)