{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Packed.Needs (
    -- * Type
    Needs (..),

    -- * Builders
    NeedsBuilder,
    runBuilder,
    (>>=),
    (>=>),
    NeedsWriter,
    NeedsWriter',

    -- * Mixing Needs together
    concatNeeds,
    applyNeeds,

    -- * Internal
    getCursor,
    getOffset,
    getOrigin,
    writeStorable,

    -- * Dangerous
    unsafeWriteTag,
    unsafeCastNeeds,
    guardRealloc,
    unsafeShiftNeedsPtr,
) where

import qualified Control.Functor.Linear as L
import Data.ByteString.Internal
import Data.Int
import Data.Kind
import qualified Data.Num.Linear as L
import qualified Data.Ord.Linear as L
import Data.Packed.Internal
import Data.Packed.Packed
import Data.Packed.Utils ((:++:))
import Data.Unrestricted.Linear
import Foreign (Storable (..))
import GHC.Exts
import GHC.ForeignPtr
import GHC.IO (IO (..))
import qualified System.IO.Linear as L
import Unsafe.Linear
import Prelude hiding ((>>=))

-- | A buffer where packed values can be written
-- The order to write these values is defined by the 'l' type list
--
-- If 'p' is an empty list, then a value of type 't' can be extracted from that buffer.
-- (See the signature of 'runReader')
data Needs (p :: [Type]) (t :: [Type])
    = Needs
        {-# UNPACK #-} !(MutableByteArray# RealWorld)
        {-# UNPACK #-} !Addr#
        {-# UNPACK #-} !Int

{-# INLINE getCursor #-}
getCursor :: Needs a b %1 -> (# Addr#, Needs a b #)
getCursor :: forall (a :: [*]) (b :: [*]).
Needs a b %1 -> (# Addr#, Needs a b #)
getCursor = (Needs a b -> (# Addr#, Needs a b #))
%1 -> Needs a b %1 -> (# Addr#, Needs a b #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
toLinear (\n :: Needs a b
n@(Needs MutableByteArray# RealWorld
_ Addr#
cursor Int
_) -> (# Addr#
cursor, Needs a b
n #))

{-# INLINE getOffset #-}
getOffset :: Needs a b %1 -> (# Int#, Needs a b #)
getOffset :: forall (a :: [*]) (b :: [*]). Needs a b %1 -> (# Int#, Needs a b #)
getOffset = (Needs a b -> (# Int#, Needs a b #))
%1 -> Needs a b %1 -> (# Int#, Needs a b #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
toLinear (\n :: Needs a b
n@(Needs MutableByteArray# RealWorld
origin Addr#
cursor Int
_) -> (# Addr#
cursor Addr# -> Addr# -> Int#
`minusAddr#` MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
origin, Needs a b
n #))

{-# INLINE getOrigin #-}
getOrigin :: Needs a b %1 -> (# Addr#, Needs a b #)
getOrigin :: forall (a :: [*]) (b :: [*]).
Needs a b %1 -> (# Addr#, Needs a b #)
getOrigin = (Needs a b -> (# Addr#, Needs a b #))
%1 -> Needs a b %1 -> (# Addr#, Needs a b #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
toLinear (\n :: Needs a b
n@(Needs MutableByteArray# RealWorld
origin Addr#
_ Int
_) -> (# MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
origin, Needs a b
n #))

{-# INLINE getSpaceLeft #-}
getSpaceLeft :: Needs a b %1 -> (# Int, Needs a b #)
getSpaceLeft :: forall (a :: [*]) (b :: [*]). Needs a b %1 -> (# Int, Needs a b #)
getSpaceLeft = (Needs a b -> (# Int, Needs a b #))
%1 -> Needs a b %1 -> (# Int, Needs a b #)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
toLinear (\n :: Needs a b
n@(Needs MutableByteArray# RealWorld
_ Addr#
_ Int
spaceLeft) -> (# Int
spaceLeft, Needs a b
n #))

-- | Casts a 'Data.Packed.Needs'
{-# INLINE unsafeCastNeeds #-}
unsafeCastNeeds :: Needs a b %1 -> Needs c d
unsafeCastNeeds :: forall (a :: [*]) (b :: [*]) (c :: [*]) (d :: [*]).
Needs a b %1 -> Needs c d
unsafeCastNeeds (Needs MutableByteArray# RealWorld
a Addr#
b Int
c) = MutableByteArray# RealWorld -> Addr# -> Int -> Needs c d
forall (p :: [*]) (t :: [*]).
MutableByteArray# RealWorld -> Addr# -> Int -> Needs p t
Needs MutableByteArray# RealWorld
a Addr#
b Int
c

--- Needs Builder

-- | A wrapper around a function that builds a 'Needs'
--
-- 'ps': The type of the expected input of the source 'Needs'
--
-- 'ts': The type of the final packed data of the source 'Needs'
--
-- 'pd': The type of the expected input of the resuling 'Needs'
--
-- 'td': The type of the final packed data of the resulting 'Needs'
--
-- __Note:__ It is an indexed monad.
type NeedsBuilder p1 t1 p2 t2 = Needs p1 t1 %1 -> L.IO (Needs p2 t2)

-- | Similar to 'Prelude.>>='
{-# INLINE (>>=) #-}
(>>=) :: L.IO (Needs p1 t1) %1 -> (Needs p1 t1 %1 -> L.IO (Needs p2 t2)) -> L.IO (Needs p2 t2)
>>= :: forall (p1 :: [*]) (t1 :: [*]) (p2 :: [*]) (t2 :: [*]).
IO (Needs p1 t1)
%1 -> (Needs p1 t1 %1 -> IO (Needs p2 t2)) -> IO (Needs p2 t2)
(>>=) IO (Needs p1 t1)
a Needs p1 t1 %1 -> IO (Needs p2 t2)
b = L.do
    !x <- IO (Needs p1 t1)
a
    !x1 <- b x
    L.return x1

-- | Similar to 'Control.Monad.>=>'
{-# INLINE (>=>) #-}
(>=>) :: (a %1 -> L.IO t) %1 -> (t %1 -> L.IO b) -> a %1 -> L.IO b
>=> :: forall a t b. (a %1 -> IO t) %1 -> (t %1 -> IO b) -> a %1 -> IO b
(>=>) a %1 -> IO t
a t %1 -> IO b
b !a
c = L.do
    !x <- a %1 -> IO t
a a
c
    !x1 <- b x
    L.return x1

-- | Shortcut type for 'NeedsBuilder'\'s that simply write a value to a 'Needs' without changing the final packed type
type NeedsWriter a r t = NeedsBuilder (a ': r) t r t

-- | Shortcut type for 'NeedsBuilder'\'s that simply write multiple values to a 'Needs' without changing the final packed type
type NeedsWriter' a r t = NeedsBuilder (a :++: r) t r t

baseBufferSize :: Int
baseBufferSize :: Int
baseBufferSize = Int
1000

-- | Runs a 'Data.Packed.Needs.NeedsBuilder' computation. Produces a 'Data.Packed.Packed'
{-# INLINE runBuilder #-}
runBuilder :: NeedsBuilder p1 r '[] r -> Packed r
runBuilder :: forall (p1 :: [*]) (r :: [*]). NeedsBuilder p1 r '[] r -> Packed r
runBuilder NeedsBuilder p1 r '[] r
builder = IO (Packed r) -> Packed r
forall a. IO a -> a
unsafeDupablePerformIO (IO (Packed r) -> Packed r) -> IO (Packed r) -> Packed r
forall a b. (a -> b) -> a -> b
$ do
    srcNeeds <- (State# RealWorld -> (# State# RealWorld, Needs p1 r #))
-> IO (Needs p1 r)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Needs p1 r #))
 -> IO (Needs p1 r))
-> (State# RealWorld -> (# State# RealWorld, Needs p1 r #))
-> IO (Needs p1 r)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case Int#
-> Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# (Int %1 -> Int#
unInt Int
baseBufferSize) Int#
0# State# RealWorld
s of
        (# State# RealWorld
s', MutableByteArray# RealWorld
byteArray #) -> (# State# RealWorld
s', MutableByteArray# RealWorld -> Addr# -> Int -> Needs p1 r
forall (p :: [*]) (t :: [*]).
MutableByteArray# RealWorld -> Addr# -> Int -> Needs p t
Needs MutableByteArray# RealWorld
byteArray (MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
byteArray) Int
baseBufferSize #)
    !finalNeeds <- L.withLinearIO $ L.fmap (toLinear Ur) (builder srcNeeds)
    finish finalNeeds

{-# INLINE finish #-}
finish :: Needs '[] a -> IO (Packed a)
finish :: forall (a :: [*]). Needs '[] a -> IO (Packed a)
finish (Needs MutableByteArray# RealWorld
og Addr#
cursor Int
_) = do
    let !contentLen :: Int#
contentLen = Addr#
cursor Addr# -> Addr# -> Int#
`minusAddr#` MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
og
    () <- (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case MutableByteArray# RealWorld
-> Int# -> State# RealWorld -> State# RealWorld
forall d. MutableByteArray# d -> Int# -> State# d -> State# d
shrinkMutableByteArray# MutableByteArray# RealWorld
og Int#
contentLen State# RealWorld
s of
        State# RealWorld
s' -> (# State# RealWorld
s', () #)
    -- SRC: https://hackage.haskell.org/package/ghc-internal-9.1201.0/docs/src/GHC.Internal.ForeignPtr.html#mallocPlainForeignPtrBytes
    let !fptr = Addr# -> ForeignPtrContents -> ForeignPtr Word8
forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr (MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
og) (MutableByteArray# RealWorld -> ForeignPtrContents
PlainPtr MutableByteArray# RealWorld
og)
    -- BS calls this function when alloacting an ForeignPtr, let's do that too
    -- https://hackage-content.haskell.org/package/bytestring-0.12.2.0/docs/src/Data.ByteString.Internal.Type.html#mkDeferredByteString
    !bs <- mkDeferredByteString fptr (I# contentLen)
    return $! unsafeToPacked bs

-- | Appends a 'Data.Needs.Needs'
{-# INLINE concatNeeds #-}
concatNeeds :: Needs p t %1 -> NeedsBuilder '[] t1 p (t1 :++: t)
concatNeeds :: forall (p :: [*]) (t :: [*]) (t1 :: [*]).
Needs p t %1 -> NeedsBuilder '[] t1 p (t1 :++: t)
concatNeeds = (Needs p t -> Needs '[] t1 %1 -> IO (Needs p (t1 :++: t)))
%1 -> Needs p t %1 -> Needs '[] t1 %1 -> IO (Needs p (t1 :++: t))
forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
toLinear2 Needs p t -> Needs '[] t1 %1 -> IO (Needs p (t1 :++: t))
forall (a :: [*]) (b :: [*]) (a' :: [*]) (b' :: [*]) (c :: [*])
       (d :: [*]).
Needs a b -> NeedsBuilder a' b' c d
appendNeeds

-- | Uses a 'Data.Packed.Needs.Needs' to fill in the required values for the other 'Data.Packed.Needs.Needs'
{-# INLINE applyNeeds #-}
applyNeeds :: Needs '[] t1 %1 -> NeedsBuilder (t1 :++: r) t r t
applyNeeds :: forall (t1 :: [*]) (r :: [*]) (t :: [*]).
Needs '[] t1 %1 -> NeedsBuilder (t1 :++: r) t r t
applyNeeds = (Needs '[] t1 -> Needs (t1 :++: r) t %1 -> IO (Needs r t))
%1 -> Needs '[] t1 %1 -> Needs (t1 :++: r) t %1 -> IO (Needs r t)
forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
toLinear2 Needs '[] t1 -> Needs (t1 :++: r) t %1 -> IO (Needs r t)
forall (a :: [*]) (b :: [*]) (a' :: [*]) (b' :: [*]) (c :: [*])
       (d :: [*]).
Needs a b -> NeedsBuilder a' b' c d
appendNeeds

-- | Writes a 'Storable'
{-# INLINE writeStorable #-}
writeStorable :: (Storable a) => a -> NeedsBuilder (a ': r) t r t
writeStorable :: forall a (r :: [*]) (t :: [*]).
Storable a =>
a -> NeedsBuilder (a : r) t r t
writeStorable !a
a !Needs (a : r) t
needs = L.do
    !newNeeds <- Int -> Needs (a : r) t %1 -> IO (Needs (a : r) t)
forall (p1 :: [*]) (t1 :: [*]).
Int -> Needs p1 t1 %1 -> IO (Needs p1 t1)
guardRealloc (a -> Int
forall a. Storable a => a -> Int
sizeOf a
a) Needs (a : r) t
needs
    let !(# cursor, newNeeds1 #) = getCursor newNeeds
    () <- toLinear (\Addr#
cursor' -> IO () %1 -> IO ()
forall a. IO a %1 -> IO a
L.fromSystemIO (Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
cursor') a
a)) cursor
    L.return (unsafeShiftNeedsPtr (sizeOf a) newNeeds1)

-- Internal

{-# INLINE guardRealloc #-}
guardRealloc :: Int -> Needs p1 t1 %1 -> L.IO (Needs p1 t1)
guardRealloc :: forall (p1 :: [*]) (t1 :: [*]).
Int -> Needs p1 t1 %1 -> IO (Needs p1 t1)
guardRealloc Int
neededSpace Needs p1 t1
needs =
    if Int
neededSpace Int %1 -> Int %1 -> Bool
forall a. Ord a => a %1 -> a %1 -> Bool
L.> Int
spaceLeft
        then Needs p1 t1 %1 -> Int -> IO (Needs p1 t1)
forall (p :: [*]) (t :: [*]). Needs p t %1 -> Int -> IO (Needs p t)
reallocNeeds Needs p1 t1
needs1 Int
neededSpace
        else Needs p1 t1 %1 -> IO (Needs p1 t1)
forall (m :: * -> *) a. Monad m => a %1 -> m a
L.return Needs p1 t1
needs1
  where
    !(# Int
spaceLeft, Needs p1 t1
needs1 #) = Needs p1 t1 %1 -> (# Int, Needs p1 t1 #)
forall (a :: [*]) (b :: [*]). Needs a b %1 -> (# Int, Needs a b #)
getSpaceLeft Needs p1 t1
needs

{-# INLINE reallocNeeds #-}
reallocNeeds :: Needs p t %1 -> Int -> L.IO (Needs p t)
reallocNeeds :: forall (p :: [*]) (t :: [*]). Needs p t %1 -> Int -> IO (Needs p t)
reallocNeeds = (Needs p t -> Int -> IO (Needs p t))
%1 -> Needs p t %1 -> Int -> IO (Needs p t)
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
toLinear Needs p t -> Int -> IO (Needs p t)
forall (p :: [*]) (t :: [*]). Needs p t -> Int -> IO (Needs p t)
reallocNeeds'

reallocNeeds' :: Needs p t -> Int -> L.IO (Needs p t)
reallocNeeds' :: forall (p :: [*]) (t :: [*]). Needs p t -> Int -> IO (Needs p t)
reallocNeeds' (Needs MutableByteArray# RealWorld
origin Addr#
cursor Int
spaceLeft) Int
additionalNeededSpace =
    let !cursorOffset :: Int
cursorOffset = Int# -> Int
I# (Int# -> Int) -> Int# -> Int
forall a b. (a -> b) -> a -> b
$ Addr#
cursor Addr# -> Addr# -> Int#
`minusAddr#` MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
origin
        !oldBufferSize :: Int
oldBufferSize = Int
spaceLeft Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
cursorOffset
        !newBufferSize :: Int
newBufferSize = Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
oldBufferSize Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Int
additionalNeededSpace)
     in IO (Needs p t) %1 -> IO (Needs p t)
forall a. IO a %1 -> IO a
L.fromSystemIO
            ( (State# RealWorld -> (# State# RealWorld, Needs p t #))
-> IO (Needs p t)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Needs p t #))
 -> IO (Needs p t))
-> (State# RealWorld -> (# State# RealWorld, Needs p t #))
-> IO (Needs p t)
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case MutableByteArray# RealWorld
-> Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d.
MutableByteArray# d
-> Int# -> State# d -> (# State# d, MutableByteArray# d #)
resizeMutableByteArray# MutableByteArray# RealWorld
origin (Int %1 -> Int#
unInt Int
newBufferSize) State# RealWorld
s of
                (# State# RealWorld
s', MutableByteArray# RealWorld
reallocedBA #) ->
                    (#
                        State# RealWorld
s'
                        , MutableByteArray# RealWorld -> Addr# -> Int -> Needs p t
forall (p :: [*]) (t :: [*]).
MutableByteArray# RealWorld -> Addr# -> Int -> Needs p t
Needs
                            MutableByteArray# RealWorld
reallocedBA
                            (MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
reallocedBA Addr# -> Int# -> Addr#
`plusAddr#` Int %1 -> Int#
unInt Int
cursorOffset)
                            (Int
newBufferSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cursorOffset)
                    #)
            )

{-# NOINLINE appendNeeds #-}
appendNeeds :: Needs a b -> NeedsBuilder a' b' c d
appendNeeds :: forall (a :: [*]) (b :: [*]) (a' :: [*]) (b' :: [*]) (c :: [*])
       (d :: [*]).
Needs a b -> NeedsBuilder a' b' c d
appendNeeds (Needs MutableByteArray# RealWorld
srcMa Addr#
cursor Int
_) Needs a' b'
dest = L.do
    let !srcLen :: Int#
srcLen = Addr#
cursor Addr# -> Addr# -> Int#
`minusAddr#` MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
srcMa
    !reallocedDest <- Int -> Needs a' b' %1 -> IO (Needs a' b')
forall (p1 :: [*]) (t1 :: [*]).
Int -> Needs p1 t1 %1 -> IO (Needs p1 t1)
guardRealloc (Int# -> Int
I# Int#
srcLen) Needs a' b'
dest
    let %1 !(# destOg, reallocedDest1 #) = getOrigin reallocedDest
        %1 !() = toLinear (\Addr#
destAddr -> case (State# RealWorld -> State# RealWorld) -> State# RealWorld
forall o. (State# RealWorld -> o) -> o
runRW# ((State# RealWorld -> State# RealWorld) -> State# RealWorld)
-> (State# RealWorld -> State# RealWorld) -> State# RealWorld
forall a b. (a -> b) -> a -> b
$ Addr# -> Addr# -> Int# -> State# RealWorld -> State# RealWorld
copyAddrToAddr# (MutableByteArray# RealWorld -> Addr#
forall d. MutableByteArray# d -> Addr#
mutableByteArrayContents# MutableByteArray# RealWorld
srcMa) Addr#
destAddr Int#
srcLen of !State# RealWorld
_ -> ()) destOg
    L.return (unsafeCastNeeds reallocedDest1)

-- Utils

{-# INLINE unsafeWriteTag #-}
unsafeWriteTag :: Int8 -> NeedsBuilder a b a' b
unsafeWriteTag :: forall (a :: [*]) (b :: [*]) (a' :: [*]).
Int8 -> NeedsBuilder a b a' b
unsafeWriteTag Int8
tag Needs a b
needs = Int8 -> NeedsBuilder (Int8 : a') b a' b
forall a (r :: [*]) (t :: [*]).
Storable a =>
a -> NeedsBuilder (a : r) t r t
writeStorable Int8
tag (Needs a b %1 -> Needs (Int8 : a') b
forall (a :: [*]) (b :: [*]) (c :: [*]) (d :: [*]).
Needs a b %1 -> Needs c d
unsafeCastNeeds Needs a b
needs)

{-# INLINE unsafeShiftNeedsPtr #-}
unsafeShiftNeedsPtr :: Int -> Needs a b %1 -> Needs a' b
unsafeShiftNeedsPtr :: forall (a :: [*]) (b :: [*]) (a' :: [*]).
Int -> Needs a b %1 -> Needs a' b
unsafeShiftNeedsPtr Int
n (Needs MutableByteArray# RealWorld
og Addr#
cursor Int
spaceLeft) =
    MutableByteArray# RealWorld -> Addr# -> Int -> Needs a' b
forall (p :: [*]) (t :: [*]).
MutableByteArray# RealWorld -> Addr# -> Int -> Needs p t
Needs MutableByteArray# RealWorld
og ((Addr# -> Int# -> Addr#) %1 -> Addr# %1 -> Int# -> Addr#
forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
toLinear2 Addr# -> Int# -> Addr#
plusAddr# Addr#
cursor (Int %1 -> Int#
unInt Int
n)) (Int
spaceLeft Int %1 -> Int %1 -> Int
forall a. AdditiveGroup a => a %1 -> a %1 -> a
L.- Int
n)