{-# LANGUAGE ScopedTypeVariables #-}

module Data.Packed.FieldSize (
    FieldSize (..),
    skipWithFieldSize,
    isolate,
    getFieldSizeFromPacked,
    writeWithFieldSize,
    readerWithFieldSize,
    applyNeedsWithFieldSize,
) where

import ByteString.StrictBuilder (builderLength)
import Control.Monad
import qualified Data.ByteString as BS
import Data.Int (Int32)
import Data.Packed.Needs
import qualified Data.Packed.Needs as N
import Data.Packed.Packable
import Data.Packed.Packed
import Data.Packed.Reader hiding (return)
import qualified Data.Packed.Reader as R
import Data.Packed.Skippable (Skippable (..), unsafeSkipN)
import Data.Packed.Unpackable
import Foreign.Ptr
import Foreign.Storable
import Prelude hiding (read)

-- | Type representation for the size of a packed data.
-- The size is in bytes.
--
-- __Note__: Take a look at the 'Data.Packed.TH.PackingFlag's to understand how to use it
newtype FieldSize = FieldSize Int32

instance {-# OVERLAPPING #-} Packable FieldSize where
    write :: forall (r :: [*]) (t :: [*]).
FieldSize -> NeedsWriter FieldSize r t
write (FieldSize Int32
value) = (Needs (FieldSize : r) t -> Needs (Int32 : r) t)
-> NeedsBuilder (FieldSize : r) t (Int32 : r) t
forall (ps :: [*]) (ts :: [*]) (pd :: [*]) (td :: [*]).
(Needs ps ts -> Needs pd td) -> NeedsBuilder ps ts pd td
mkNeedsBuilder Needs (FieldSize : r) t -> Needs (Int32 : r) t
forall (a :: [*]) (b :: [*]) (c :: [*]) (d :: [*]).
Needs a b -> Needs c d
unsafeCastNeeds NeedsBuilder (FieldSize : r) t (Int32 : r) t
-> NeedsBuilder (Int32 : r) t r t
-> NeedsBuilder (FieldSize : r) t r t
forall (p1 :: [*]) (t1 :: [*]) (p2 :: [*]) (t2 :: [*]) (p3 :: [*])
       (t3 :: [*]).
NeedsBuilder p1 t1 p2 t2
-> NeedsBuilder p2 t2 p3 t3 -> NeedsBuilder p1 t1 p3 t3
N.>> Int32 -> NeedsBuilder (Int32 : r) t r t
forall (r :: [*]) (t :: [*]). Int32 -> NeedsWriter Int32 r t
forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter a r t
write Int32
value

instance {-# OVERLAPPING #-} Unpackable FieldSize where
    reader :: forall (r :: [*]). PackedReader '[FieldSize] r FieldSize
reader = (ReaderPtr ('[FieldSize] :++: r)
 -> Int -> IO (FieldSize, ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader '[FieldSize] r FieldSize
forall (p :: [*]) (r :: [*]) v.
(ReaderPtr ('[FieldSize] :++: r)
 -> Int -> IO (v, ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader p r v
mkPackedReader ((ReaderPtr ('[FieldSize] :++: r)
  -> Int -> IO (FieldSize, ReaderPtr ('[FieldSize] :++: r), Int))
 -> PackedReader '[FieldSize] r FieldSize)
-> (ReaderPtr ('[FieldSize] :++: r)
    -> Int -> IO (FieldSize, ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader '[FieldSize] r FieldSize
forall a b. (a -> b) -> a -> b
$ \ReaderPtr ('[FieldSize] :++: r)
packed Int
l -> do
        (fieldSize, rest, l1) <- PackedReader '[Int32] Any Int32
-> ReaderPtr ('[FieldSize] :++: r)
-> Int
-> IO (Int32, ReaderPtr ('[FieldSize] :++: r), Int)
forall (p :: [*]) (r :: [*]) v.
PackedReader p r v
-> ReaderPtr ('[FieldSize] :++: r)
-> Int
-> IO (v, ReaderPtr ('[FieldSize] :++: r), Int)
runPackedReader PackedReader '[Int32] Any Int32
forall (r :: [*]). PackedReader '[Int32] r Int32
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader ReaderPtr ('[FieldSize] :++: r)
packed Int
l
        return (FieldSize fieldSize, rest, l1)

instance {-# OVERLAPPING #-} Skippable FieldSize where
    skip :: forall (r :: [*]). PackedReader '[FieldSize] r ()
skip = Int -> PackedReader '[FieldSize] r ()
forall (a :: [*]) (r :: [*]). Int -> PackedReader a r ()
unsafeSkipN (Int32 -> Int
forall a. Storable a => a -> Int
sizeOf (Int32
1 :: Int32))

{-# INLINE getFieldSizeFromPacked #-}

-- | Returns the size of the packed value.
--
-- __Warning:__ For this to be accurate, there should only be one value packed in the binary strea.
getFieldSizeFromPacked :: Packed '[a] -> FieldSize
getFieldSizeFromPacked :: forall a. Packed '[a] -> FieldSize
getFieldSizeFromPacked Packed '[a]
packed = Int32 -> FieldSize
FieldSize (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length (Packed '[a] -> ByteString
forall (a :: [*]). Packed a -> ByteString
fromPacked Packed '[a]
packed))

{-# INLINE skipWithFieldSize #-}

-- | Allows skipping over a field without having to unpack it
skipWithFieldSize :: PackedReader '[FieldSize, a] r ()
skipWithFieldSize :: forall a (r :: [*]). PackedReader '[FieldSize, a] r ()
skipWithFieldSize = (ReaderPtr ('[FieldSize] :++: r)
 -> Int -> IO ((), ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader '[FieldSize, a] r ()
forall (p :: [*]) (r :: [*]) v.
(ReaderPtr ('[FieldSize] :++: r)
 -> Int -> IO (v, ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader p r v
mkPackedReader ((ReaderPtr ('[FieldSize] :++: r)
  -> Int -> IO ((), ReaderPtr ('[FieldSize] :++: r), Int))
 -> PackedReader '[FieldSize, a] r ())
-> (ReaderPtr ('[FieldSize] :++: r)
    -> Int -> IO ((), ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader '[FieldSize, a] r ()
forall a b. (a -> b) -> a -> b
$ \ReaderPtr ('[FieldSize] :++: r)
packed Int
l -> do
    (FieldSize s, packed1, l1) <- PackedReader '[FieldSize] Any FieldSize
-> ReaderPtr ('[FieldSize] :++: r)
-> Int
-> IO (FieldSize, ReaderPtr ('[FieldSize] :++: r), Int)
forall (p :: [*]) (r :: [*]) v.
PackedReader p r v
-> ReaderPtr ('[FieldSize] :++: r)
-> Int
-> IO (v, ReaderPtr ('[FieldSize] :++: r), Int)
runPackedReader PackedReader '[FieldSize] Any FieldSize
forall (r :: [*]). PackedReader '[FieldSize] r FieldSize
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader ReaderPtr ('[FieldSize] :++: r)
packed Int
l
    let size64 = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
s
    return ((), packed1 `plusPtr` size64, l1 - size64)

{-# INLINE writeWithFieldSize #-}

-- | Write a value into a 'Data.Packed.Needs.Needs', along with its 'FieldSize'
--
-- Note: Universal quantifier is nedded for GHC < 9.10, because of ScopedTypeVariables
writeWithFieldSize :: forall a r t. (Packable a) => a -> NeedsWriter' '[FieldSize, a] r t
writeWithFieldSize :: forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter' '[FieldSize, a] r t
writeWithFieldSize a
a = FieldSize -> NeedsWriter FieldSize (a : r) t
forall (r :: [*]) (t :: [*]).
FieldSize -> NeedsWriter FieldSize r t
forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter a r t
write (Int32 -> FieldSize
FieldSize Int32
size) NeedsWriter FieldSize (a : r) t
-> NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
forall (p1 :: [*]) (t1 :: [*]) (p2 :: [*]) (t2 :: [*]) (p3 :: [*])
       (t3 :: [*]).
NeedsBuilder p1 t1 p2 t2
-> NeedsBuilder p2 t2 p3 t3 -> NeedsBuilder p1 t1 p3 t3
N.>> Needs '[] '[a] -> NeedsBuilder ('[a] :++: r) t r t
forall (t1 :: [*]) (r :: [*]) (t :: [*]).
Needs '[] t1 -> NeedsBuilder (t1 :++: r) t r t
applyNeeds Needs '[] '[a]
aNeeds
  where
    size :: Int32
size = Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Builder -> Int
builderLength Builder
aBuilder)
    aNeeds :: Needs '[] '[a]
    aNeeds :: Needs '[] '[a]
aNeeds@(Needs Builder
aBuilder) = NeedsBuilder '[a] '[a] '[] '[a] -> Needs '[] '[a]
forall (a :: [*]) (b :: [*]) (x :: [*]) (y :: [*]).
NeedsBuilder a b x y -> Needs x y
withEmptyNeeds (a -> NeedsBuilder '[a] '[a] '[] '[a]
forall (r :: [*]) (t :: [*]). a -> NeedsWriter a r t
forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter a r t
write a
a)

{-# INLINE applyNeedsWithFieldSize #-}
applyNeedsWithFieldSize :: Needs '[] '[a] -> NeedsWriter' (FieldSize ': a ': '[]) r t
applyNeedsWithFieldSize :: forall a (r :: [*]) (t :: [*]).
Needs '[] '[a] -> NeedsWriter' '[FieldSize, a] r t
applyNeedsWithFieldSize n :: Needs '[] '[a]
n@(Needs Builder
builder) = FieldSize -> NeedsWriter FieldSize (a : r) t
forall (r :: [*]) (t :: [*]).
FieldSize -> NeedsWriter FieldSize r t
forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter a r t
write (Int32 -> FieldSize
FieldSize (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Builder -> Int
builderLength Builder
builder))) NeedsWriter FieldSize (a : r) t
-> NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
forall (p1 :: [*]) (t1 :: [*]) (p2 :: [*]) (t2 :: [*]) (p3 :: [*])
       (t3 :: [*]).
NeedsBuilder p1 t1 p2 t2
-> NeedsBuilder p2 t2 p3 t3 -> NeedsBuilder p1 t1 p3 t3
N.>> Needs '[] '[a] -> NeedsBuilder ('[a] :++: r) t r t
forall (t1 :: [*]) (r :: [*]) (t :: [*]).
Needs '[] t1 -> NeedsBuilder (t1 :++: r) t r t
applyNeeds Needs '[] '[a]
n

{-# INLINE readerWithFieldSize #-}

-- | Produces a reader for a value preceded by its 'FieldSize'
readerWithFieldSize :: (Unpackable a) => PackedReader '[FieldSize, a] r a
readerWithFieldSize :: forall a (r :: [*]).
Unpackable a =>
PackedReader '[FieldSize, a] r a
readerWithFieldSize = PackedReader '[FieldSize] (a : r) ()
PackedReader '[FieldSize] ('[a] :++: r) ()
forall (r :: [*]). PackedReader '[FieldSize] r ()
forall a (r :: [*]). Skippable a => PackedReader '[a] r ()
skip PackedReader '[FieldSize] ('[a] :++: r) ()
-> PackedReader '[a] r a
-> PackedReader ('[FieldSize] :++: '[a]) r a
forall (p :: [*]) (r1 :: [*]) (r2 :: [*]) v v'.
PackedReader p (r1 :++: r2) v
-> PackedReader r1 r2 v' -> PackedReader (p :++: r1) r2 v'
R.>> PackedReader '[a] r a
forall (r :: [*]). PackedReader '[a] r a
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader

{-# INLINE isolate #-}

-- | Splits the 'Packed' value, and isolate the first encoded value.
isolate :: PackedReader '[FieldSize, a] r (Packed '[a])
isolate :: forall a (r :: [*]). PackedReader '[FieldSize, a] r (Packed '[a])
isolate = (ReaderPtr ('[FieldSize] :++: r)
 -> Int -> IO (Packed '[a], ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader '[FieldSize, a] r (Packed '[a])
forall (p :: [*]) (r :: [*]) v.
(ReaderPtr ('[FieldSize] :++: r)
 -> Int -> IO (v, ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader p r v
mkPackedReader ((ReaderPtr ('[FieldSize] :++: r)
  -> Int -> IO (Packed '[a], ReaderPtr ('[FieldSize] :++: r), Int))
 -> PackedReader '[FieldSize, a] r (Packed '[a]))
-> (ReaderPtr ('[FieldSize] :++: r)
    -> Int -> IO (Packed '[a], ReaderPtr ('[FieldSize] :++: r), Int))
-> PackedReader '[FieldSize, a] r (Packed '[a])
forall a b. (a -> b) -> a -> b
$ \ReaderPtr ('[FieldSize] :++: r)
packed Int
l -> do
    (FieldSize s, packed1, l1) <- PackedReader '[FieldSize] Any FieldSize
-> ReaderPtr ('[FieldSize] :++: r)
-> Int
-> IO (FieldSize, ReaderPtr ('[FieldSize] :++: r), Int)
forall (p :: [*]) (r :: [*]) v.
PackedReader p r v
-> ReaderPtr ('[FieldSize] :++: r)
-> Int
-> IO (v, ReaderPtr ('[FieldSize] :++: r), Int)
runPackedReader PackedReader '[FieldSize] Any FieldSize
forall (r :: [*]). PackedReader '[FieldSize] r FieldSize
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader ReaderPtr ('[FieldSize] :++: r)
packed Int
l
    let sizeInt = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
s
    return (unsafeToPacked' packed1 sizeInt, packed1 `plusPtr` sizeInt, l1 - sizeInt)