{-# LANGUAGE UnboxedTuples #-}

#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Array.Prim.Pinned.Mut.Type
-- Copyright   : (c) 2020 Composewell Technologies
--
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
module Streamly.Internal.Data.Array.Prim.Pinned.Mut.Type
    (
      Array (..)

    -- * Construction
    , newArray
    , newAlignedArray
    , unsafeWriteIndex

    , spliceTwo
    , unsafeCopy

    , fromListM
    , fromListNM
    , fromStreamDN
    , fromStreamD

    -- * Streams of arrays
    , fromStreamDArraysOf

    , packArraysChunksOf
    , lpackArraysChunksOf

#if !defined(mingw32_HOST_OS)
--    , groupIOVecsOf
#endif

    -- * Elimination
    , unsafeReadIndex
    , length
    , byteLength

    , writeN
    , ArrayUnsafe(..)
    , writeNUnsafe
    , writeNAligned
    , write

    -- * Utilities
    , resizeArray
    , shrinkArray

    , touchArray
    , withArrayAsPtr
    )
where

import GHC.IO (IO(..))

#include "Streamly/Internal/Data/Array/Prim/MutTypesInclude.hs"

-------------------------------------------------------------------------------
-- Allocation (Pinned)
-------------------------------------------------------------------------------

-- XXX we can use a single newArray routine which accepts an allocation
-- function which could be newByteArray#, newPinnedByteArray# or
-- newAlignedPinnedByteArray#. That function can go in the common include file.
--
-- | Allocate an array that is pinned and can hold 'count' items.  The memory of
-- the array is uninitialized.
--
-- Note that this is internal routine, the reference to this array cannot be
-- given out until the array has been written to and frozen.
{-# INLINE newArray #-}
newArray ::
       forall m a. (MonadIO m, Prim a)
    => Int
    -> m (Array a)
newArray :: forall (m :: * -> *) a. (MonadIO m, Prim a) => Int -> m (Array a)
newArray (I# Int#
n#) =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let bytes :: Int#
bytes = Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => a -> Int#
sizeOf# (forall a. HasCallStack => a
undefined :: a)
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState IO)
s# ->
            case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newPinnedByteArray# Int#
bytes State# (PrimState IO)
s# of
                (# State# RealWorld
s1#, MutableByteArray# RealWorld
arr# #) -> (# State# RealWorld
s1#, forall a. MutableByteArray# RealWorld -> Array a
Array MutableByteArray# RealWorld
arr# #)

-- Change order of args?
-- | Allocate a new array aligned to the specified alignment and using pinned
-- memory.
{-# INLINE newAlignedArray #-}
newAlignedArray ::
       forall m a. (MonadIO m, Prim a)
    => Int -- size
    -> Int -- Alignment
    -> m (Array a)
newAlignedArray :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> Int -> m (Array a)
newAlignedArray (I# Int#
n#) (I# Int#
a#) =
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        let bytes :: Int#
bytes = Int#
n# Int# -> Int# -> Int#
*# forall a. Prim a => a -> Int#
sizeOf# (forall a. HasCallStack => a
undefined :: a)
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState IO)
s# ->
            case forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# Int#
bytes Int#
a# State# (PrimState IO)
s# of
                (# State# RealWorld
s1#, MutableByteArray# RealWorld
arr# #) -> (# State# RealWorld
s1#, forall a. MutableByteArray# RealWorld -> Array a
Array MutableByteArray# RealWorld
arr# #)

-- | Resize (pinned) mutable byte array to new specified size (in elem
-- count). The returned array is either the original array resized in-place or,
-- if not possible, a newly allocated (pinned) array (with the original content
-- copied over).
{-# INLINE resizeArray #-}
resizeArray ::
       (MonadIO m, Prim a)
    => Array a
    -> Int -- ^ new size
    -> m (Array a)
resizeArray :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Array a -> Int -> m (Array a)
resizeArray Array a
arr Int
i = do
    Int
len <- forall (m :: * -> *) a. (MonadIO m, Prim a) => Array a -> m Int
length Array a
arr
    if Int
len forall a. Eq a => a -> a -> Bool
== Int
i
    then forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr
    else if Int
i forall a. Ord a => a -> a -> Bool
< Int
len
         then forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Array a -> Int -> m ()
shrinkArray Array a
arr Int
i forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr
         else do
             Array a
nArr <- forall (m :: * -> *) a. (MonadIO m, Prim a) => Int -> m (Array a)
newArray Int
i
             forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Array a -> Int -> Array a -> Int -> Int -> m ()
unsafeCopy Array a
nArr Int
0 Array a
arr Int
0 Int
len
             forall (m :: * -> *) a. Monad m => a -> m a
return Array a
nArr

-------------------------------------------------------------------------------
-- Aligned Construction
-------------------------------------------------------------------------------

-- XXX we can also factor out common code in writeN and writeNAligned in the
-- same way as suggested above.
--
{-# INLINE_NORMAL writeNAligned #-}
writeNAligned ::
       (MonadIO m, Prim a)
    => Int
    -> Int
    -> Fold m a (Array a)
writeNAligned :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> Int -> Fold m a (Array a)
writeNAligned Int
align Int
limit = forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold forall {m :: * -> *} {a}.
(MonadIO m, Prim a) =>
Tuple' (Array a) Int
-> a -> m (Step (Tuple' (Array a) Int) (Array a))
step forall {b}. m (Step (Tuple' (Array a) Int) b)
initial forall {m :: * -> *} {a}.
(MonadIO m, Prim a) =>
Tuple' (Array a) Int -> m (Array a)
extract

    where

    initial :: m (Step (Tuple' (Array a) Int) b)
initial = do
        Array a
marr <- forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Int -> Int -> m (Array a)
newAlignedArray Int
limit Int
align
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> Tuple' a b
Tuple' Array a
marr Int
0

    extract :: Tuple' (Array a) Int -> m (Array a)
extract (Tuple' Array a
marr Int
len) = forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Array a -> Int -> m ()
shrinkArray Array a
marr Int
len forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Array a
marr

    step :: Tuple' (Array a) Int
-> a -> f (Step (Tuple' (Array a) Int) (Array a))
step s :: Tuple' (Array a) Int
s@(Tuple' Array a
marr Int
i) a
x
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
limit = forall s b. b -> Step s b
FL.Done forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {m :: * -> *} {a}.
(MonadIO m, Prim a) =>
Tuple' (Array a) Int -> m (Array a)
extract Tuple' (Array a) Int
s
        | Bool
otherwise = do
            forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Array a -> Int -> a -> m ()
unsafeWriteIndex Array a
marr Int
i a
x
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> Tuple' a b
Tuple' Array a
marr (Int
i forall a. Num a => a -> a -> a
+ Int
1)

-------------------------------------------------------------------------------
-- Mutation with pointers
-------------------------------------------------------------------------------

-- XXX This section can probably go in a common include file for pinned arrays.

-- Change name later.
{-# INLINE toPtr #-}
toPtr :: Array a -> Ptr a
toPtr :: forall a. Array a -> Ptr a
toPtr (Array MutableByteArray# RealWorld
arr#) = forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents# (unsafeCoerce# :: forall a b. a -> b
unsafeCoerce# MutableByteArray# RealWorld
arr#))

{-# INLINE touchArray #-}
touchArray :: Array a -> IO ()
touchArray :: forall a. Array a -> IO ()
touchArray Array a
arr = forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case touch# :: forall a. a -> State# RealWorld -> State# RealWorld
touch# Array a
arr State# RealWorld
s of State# RealWorld
s' -> (# State# RealWorld
s', () #)

{-# INLINE withArrayAsPtr #-}
withArrayAsPtr :: Array a -> (Ptr a -> IO b) -> IO b
withArrayAsPtr :: forall a b. Array a -> (Ptr a -> IO b) -> IO b
withArrayAsPtr Array a
arr Ptr a -> IO b
f = do
    b
r <- Ptr a -> IO b
f (forall a. Array a -> Ptr a
toPtr Array a
arr)
    forall a. Array a -> IO ()
touchArray Array a
arr
    forall (m :: * -> *) a. Monad m => a -> m a
return b
r