{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UnboxedTuples #-}

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

import qualified Control.Functor.Linear as L
import Control.Monad.Identity (Identity (Identity))
import qualified Data.ByteString as BS
import Data.Packed.Internal
import Data.Packed.Needs
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 GHC.Exts
import GHC.Int
import qualified System.IO.Linear as L
import Unsafe.Linear
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 deriving (Integer -> FieldSize
FieldSize -> FieldSize
FieldSize -> FieldSize -> FieldSize
(FieldSize -> FieldSize -> FieldSize)
-> (FieldSize -> FieldSize -> FieldSize)
-> (FieldSize -> FieldSize -> FieldSize)
-> (FieldSize -> FieldSize)
-> (FieldSize -> FieldSize)
-> (FieldSize -> FieldSize)
-> (Integer -> FieldSize)
-> Num FieldSize
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
$c+ :: FieldSize -> FieldSize -> FieldSize
+ :: FieldSize -> FieldSize -> FieldSize
$c- :: FieldSize -> FieldSize -> FieldSize
- :: FieldSize -> FieldSize -> FieldSize
$c* :: FieldSize -> FieldSize -> FieldSize
* :: FieldSize -> FieldSize -> FieldSize
$cnegate :: FieldSize -> FieldSize
negate :: FieldSize -> FieldSize
$cabs :: FieldSize -> FieldSize
abs :: FieldSize -> FieldSize
$csignum :: FieldSize -> FieldSize
signum :: FieldSize -> FieldSize
$cfromInteger :: Integer -> FieldSize
fromInteger :: Integer -> FieldSize
Num, Int -> FieldSize
FieldSize -> Int
FieldSize -> [FieldSize]
FieldSize -> FieldSize
FieldSize -> FieldSize -> [FieldSize]
FieldSize -> FieldSize -> FieldSize -> [FieldSize]
(FieldSize -> FieldSize)
-> (FieldSize -> FieldSize)
-> (Int -> FieldSize)
-> (FieldSize -> Int)
-> (FieldSize -> [FieldSize])
-> (FieldSize -> FieldSize -> [FieldSize])
-> (FieldSize -> FieldSize -> [FieldSize])
-> (FieldSize -> FieldSize -> FieldSize -> [FieldSize])
-> Enum FieldSize
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: FieldSize -> FieldSize
succ :: FieldSize -> FieldSize
$cpred :: FieldSize -> FieldSize
pred :: FieldSize -> FieldSize
$ctoEnum :: Int -> FieldSize
toEnum :: Int -> FieldSize
$cfromEnum :: FieldSize -> Int
fromEnum :: FieldSize -> Int
$cenumFrom :: FieldSize -> [FieldSize]
enumFrom :: FieldSize -> [FieldSize]
$cenumFromThen :: FieldSize -> FieldSize -> [FieldSize]
enumFromThen :: FieldSize -> FieldSize -> [FieldSize]
$cenumFromTo :: FieldSize -> FieldSize -> [FieldSize]
enumFromTo :: FieldSize -> FieldSize -> [FieldSize]
$cenumFromThenTo :: FieldSize -> FieldSize -> FieldSize -> [FieldSize]
enumFromThenTo :: FieldSize -> FieldSize -> FieldSize -> [FieldSize]
Enum, Num FieldSize
Ord FieldSize
(Num FieldSize, Ord FieldSize) =>
(FieldSize -> Rational) -> Real FieldSize
FieldSize -> Rational
forall a. (Num a, Ord a) => (a -> Rational) -> Real a
$ctoRational :: FieldSize -> Rational
toRational :: FieldSize -> Rational
Real, Eq FieldSize
Eq FieldSize =>
(FieldSize -> FieldSize -> Ordering)
-> (FieldSize -> FieldSize -> Bool)
-> (FieldSize -> FieldSize -> Bool)
-> (FieldSize -> FieldSize -> Bool)
-> (FieldSize -> FieldSize -> Bool)
-> (FieldSize -> FieldSize -> FieldSize)
-> (FieldSize -> FieldSize -> FieldSize)
-> Ord FieldSize
FieldSize -> FieldSize -> Bool
FieldSize -> FieldSize -> Ordering
FieldSize -> FieldSize -> FieldSize
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: FieldSize -> FieldSize -> Ordering
compare :: FieldSize -> FieldSize -> Ordering
$c< :: FieldSize -> FieldSize -> Bool
< :: FieldSize -> FieldSize -> Bool
$c<= :: FieldSize -> FieldSize -> Bool
<= :: FieldSize -> FieldSize -> Bool
$c> :: FieldSize -> FieldSize -> Bool
> :: FieldSize -> FieldSize -> Bool
$c>= :: FieldSize -> FieldSize -> Bool
>= :: FieldSize -> FieldSize -> Bool
$cmax :: FieldSize -> FieldSize -> FieldSize
max :: FieldSize -> FieldSize -> FieldSize
$cmin :: FieldSize -> FieldSize -> FieldSize
min :: FieldSize -> FieldSize -> FieldSize
Ord, FieldSize -> FieldSize -> Bool
(FieldSize -> FieldSize -> Bool)
-> (FieldSize -> FieldSize -> Bool) -> Eq FieldSize
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: FieldSize -> FieldSize -> Bool
== :: FieldSize -> FieldSize -> Bool
$c/= :: FieldSize -> FieldSize -> Bool
/= :: FieldSize -> FieldSize -> Bool
Eq)

deriving instance Integral FieldSize

instance {-# OVERLAPPING #-} Packable FieldSize where
    write :: forall (r :: [*]) (t :: [*]).
FieldSize -> NeedsWriter FieldSize r t
write (FieldSize Int32
value) Needs (FieldSize : r) t
needs = Int32 -> NeedsWriter Int32 r t
forall (r :: [*]) (t :: [*]). Int32 -> NeedsWriter Int32 r t
forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter a r t
write Int32
value (Needs (FieldSize : r) t %1 -> Needs (Int32 : r) t
forall (a :: [*]) (b :: [*]) (c :: [*]) (d :: [*]).
Needs a b %1 -> Needs c d
unsafeCastNeeds Needs (FieldSize : r) t
needs)

instance {-# OVERLAPPING #-} Unpackable FieldSize where
    reader :: forall (r :: [*]). PackedReader '[FieldSize] r FieldSize
reader = (PackedFragment ('[FieldSize] :++: r)
 -> Identity (FieldSize, PackedFragment r))
-> PackedReader '[FieldSize] r FieldSize
forall (p :: [*]) (r :: [*]) v.
(PackedFragment (p :++: r) -> Identity (v, PackedFragment r))
-> PackedReader p r v
mkPackedReader ((PackedFragment ('[FieldSize] :++: r)
  -> Identity (FieldSize, PackedFragment r))
 -> PackedReader '[FieldSize] r FieldSize)
-> (PackedFragment ('[FieldSize] :++: r)
    -> Identity (FieldSize, PackedFragment r))
-> PackedReader '[FieldSize] r FieldSize
forall a b. (a -> b) -> a -> b
$ \PackedFragment ('[FieldSize] :++: r)
pf ->
        let
            Identity !(!Int32
fieldSize, !PackedFragment r
pf1) = PackedReader '[Int32] r Int32
-> PackedFragment ('[Int32] :++: r)
-> Identity (Int32, PackedFragment r)
forall (p :: [*]) (r :: [*]) v.
PackedReader p r v
-> PackedFragment (p :++: r) -> Identity (v, PackedFragment r)
runReaderStep PackedReader '[Int32] r Int32
forall (r :: [*]). PackedReader '[Int32] r Int32
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader (PackedFragment (FieldSize : r) -> PackedFragment (Int32 : r)
forall (p :: [*]) (t :: [*]). PackedFragment p -> PackedFragment t
castPackedFragment PackedFragment (FieldSize : r)
PackedFragment ('[FieldSize] :++: r)
pf)
         in
            (FieldSize, PackedFragment r)
-> Identity (FieldSize, PackedFragment r)
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int32 -> FieldSize
FieldSize Int32
fieldSize, PackedFragment r
pf1)

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 = (PackedFragment ('[FieldSize, a] :++: r)
 -> Identity ((), PackedFragment r))
-> PackedReader '[FieldSize, a] r ()
forall (p :: [*]) (r :: [*]) v.
(PackedFragment (p :++: r) -> Identity (v, PackedFragment r))
-> PackedReader p r v
mkPackedReader ((PackedFragment ('[FieldSize, a] :++: r)
  -> Identity ((), PackedFragment r))
 -> PackedReader '[FieldSize, a] r ())
-> (PackedFragment ('[FieldSize, a] :++: r)
    -> Identity ((), PackedFragment r))
-> PackedReader '[FieldSize, a] r ()
forall a b. (a -> b) -> a -> b
$ \PackedFragment ('[FieldSize, a] :++: r)
pf ->
    let
        Identity !(FieldSize Int32
s, PF Ptr Word8
packed1 Int
l1) = PackedReader '[FieldSize] (a : r) FieldSize
-> PackedFragment ('[FieldSize] :++: (a : r))
-> Identity (FieldSize, PackedFragment (a : r))
forall (p :: [*]) (r :: [*]) v.
PackedReader p r v
-> PackedFragment (p :++: r) -> Identity (v, PackedFragment r)
runReaderStep PackedReader '[FieldSize] (a : r) FieldSize
forall (r :: [*]). PackedReader '[FieldSize] r FieldSize
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader PackedFragment ('[FieldSize, a] :++: r)
PackedFragment ('[FieldSize] :++: (a : r))
pf
        !size64 :: Int
size64 = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
s
     in
        ((), PackedFragment r) -> Identity ((), PackedFragment r)
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ((), Ptr Word8 -> Int -> PackedFragment r
forall (p :: [*]). Ptr Word8 -> Int -> PackedFragment p
PF (Ptr Word8
packed1 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size64) (Int
l1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
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 = NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
forall a (r :: [*]) (t :: [*]).
NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
withFieldSize (a -> NeedsBuilder (a : r) t r t
forall (r :: [*]) (t :: [*]). a -> NeedsWriter a r t
forall a (r :: [*]) (t :: [*]).
Packable a =>
a -> NeedsWriter a r t
write a
a)

{-# INLINE withFieldSize #-}
withFieldSize :: NeedsBuilder (a ': r) t r t -> NeedsBuilder (FieldSize ': a ': r) t r t
withFieldSize :: forall a (r :: [*]) (t :: [*]).
NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
withFieldSize NeedsBuilder (a : r) t r t
cont Needs (FieldSize : a : r) t
needs = L.do
    let !indirectionSize :: Int
indirectionSize = Int32 -> Int
forall a. Storable a => a -> Int
sizeOf (Int32
0 :: Int32)
    -- Reallocating the buffer so that the fieldsize can fit
    !newNeeds <- Int
-> Needs (FieldSize : a : r) t
%1 -> IO (Needs (FieldSize : a : r) t)
forall (p1 :: [*]) (t1 :: [*]).
Int -> Needs p1 t1 %1 -> IO (Needs p1 t1)
guardRealloc Int
indirectionSize Needs (FieldSize : a : r) t
needs
    -- Get the position of the buffer where the FS will be
    let !(# fieldSizeOffset, newNeeds1 #) = getOffset newNeeds
    -- Shift the cursor
    !writtenNeeds <- cont (unsafeShiftNeedsPtr indirectionSize newNeeds1)
    -- Get the final position of the cursor
    let !(# finalCursor, writtenNeeds1 #) = getOffset writtenNeeds
        !(# og, writtenNeeds2 #) = getOrigin writtenNeeds1
    () <-
        toLinear3
            ( \Int#
finalCursor' Int#
fsPosition Addr#
og' Int
fsSize ->
                let
                    -- Count the number of bytes that were written
                    !writtenBytes :: Int32#
writtenBytes = Int# -> Int32#
intToInt32# (Int#
finalCursor' Int# -> Int# -> Int#
-# (Int#
fsPosition Int# -> Int# -> Int#
+# Int %1 -> Int#
unInt Int
fsSize))
                 in
                    -- And write it

                    IO () %1 -> IO ()
forall a. IO a %1 -> IO a
L.fromSystemIO (Ptr Int32 -> Int32 -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Addr# -> Ptr Int32
forall a. Addr# -> Ptr a
Ptr (Addr# -> Ptr Int32) -> Addr# -> Ptr Int32
forall a b. (a -> b) -> a -> b
$ Addr#
og' Addr# -> Int# -> Addr#
`plusAddr#` Int#
fsPosition) (Int32# -> Int32
I32# Int32#
writtenBytes))
            )
            finalCursor
            fieldSizeOffset
            og
            indirectionSize
    L.return writtenNeeds2

{-# INLINE applyNeedsWithFieldSize #-}
applyNeedsWithFieldSize :: Needs '[] '[a] -> NeedsWriter' (FieldSize ': a ': '[]) r t
applyNeedsWithFieldSize :: forall a (r :: [*]) (t :: [*]).
Needs '[] '[a] -> NeedsWriter' '[FieldSize, a] r t
applyNeedsWithFieldSize Needs '[] '[a]
n = NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
forall a (r :: [*]) (t :: [*]).
NeedsBuilder (a : r) t r t
-> NeedsBuilder (FieldSize : a : r) t r t
withFieldSize (Needs '[] '[a] %1 -> NeedsBuilder ('[a] :++: r) t r t
forall (t1 :: [*]) (r :: [*]) (t :: [*]).
Needs '[] t1 %1 -> 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 (PackedFragment '[a])
isolate :: forall a (r :: [*]).
PackedReader '[FieldSize, a] r (PackedFragment '[a])
isolate = (PackedFragment ('[FieldSize, a] :++: r)
 -> Identity (PackedFragment '[a], PackedFragment r))
-> PackedReader '[FieldSize, a] r (PackedFragment '[a])
forall (p :: [*]) (r :: [*]) v.
(PackedFragment (p :++: r) -> Identity (v, PackedFragment r))
-> PackedReader p r v
mkPackedReader ((PackedFragment ('[FieldSize, a] :++: r)
  -> Identity (PackedFragment '[a], PackedFragment r))
 -> PackedReader '[FieldSize, a] r (PackedFragment '[a]))
-> (PackedFragment ('[FieldSize, a] :++: r)
    -> Identity (PackedFragment '[a], PackedFragment r))
-> PackedReader '[FieldSize, a] r (PackedFragment '[a])
forall a b. (a -> b) -> a -> b
$ \PackedFragment ('[FieldSize, a] :++: r)
pf ->
    let
        Identity !(FieldSize Int32
s, PF Ptr Word8
packed1 Int
l1) = PackedReader '[FieldSize] (a : r) FieldSize
-> PackedFragment ('[FieldSize] :++: (a : r))
-> Identity (FieldSize, PackedFragment (a : r))
forall (p :: [*]) (r :: [*]) v.
PackedReader p r v
-> PackedFragment (p :++: r) -> Identity (v, PackedFragment r)
runReaderStep PackedReader '[FieldSize] (a : r) FieldSize
forall (r :: [*]). PackedReader '[FieldSize] r FieldSize
forall a (r :: [*]). Unpackable a => PackedReader '[a] r a
reader PackedFragment ('[FieldSize, a] :++: r)
PackedFragment ('[FieldSize] :++: (a : r))
pf
        !sizeInt :: Int
sizeInt = Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
s
     in
        (PackedFragment '[a], PackedFragment r)
-> Identity (PackedFragment '[a], PackedFragment r)
forall a. a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr Word8 -> Int -> PackedFragment '[a]
forall (p :: [*]). Ptr Word8 -> Int -> PackedFragment p
PF Ptr Word8
packed1 Int
sizeInt, Ptr Word8 -> Int -> PackedFragment r
forall (p :: [*]). Ptr Word8 -> Int -> PackedFragment p
PF (Ptr Word8
packed1 Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
sizeInt) (Int
l1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sizeInt))