#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Array.Foreign.Mut
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Unboxed pinned mutable array type for 'Storable' types with an option to use
-- foreign (non-GHC) memory allocators. Fulfils the following goals:
--
-- * Random access (array)
-- * Efficient storage (unboxed)
-- * Performance (unboxed access)
-- * Performance - in-place operations (mutable)
-- * Performance - GC (pinned, mutable)
-- * interfacing with OS (pinned)
-- * Fragmentation control (foreign allocators)
--
-- Stream and Fold APIs allow easy, efficient and convenient operations on
-- arrays.

module Streamly.Internal.Data.Array.Foreign.Mut
    (
      module Streamly.Internal.Data.Array.Foreign.Mut.Type
    , splitOn
    , genSlicesFromLen
    , getSlicesFromLen
    )
where

import Prelude hiding (foldr, length, read, splitAt)

import Streamly.Internal.Data.Array.Foreign.Mut.Type
import qualified Streamly.Internal.Data.Stream.StreamD as D
import Foreign.Storable (Storable)
import Streamly.Internal.Data.Stream.IsStream.Type (SerialT)
import qualified Streamly.Internal.Data.Stream.IsStream.Type as IsStream
import qualified Streamly.Internal.Data.Unfold as Unfold
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Control.Monad.IO.Class (MonadIO(..))

-- | Split the array into a stream of slices using a predicate. The element
-- matching the predicate is dropped.
--
-- /Pre-release/
{-# INLINE splitOn #-}
splitOn :: (MonadIO m, Storable a) =>
    (a -> Bool) -> Array a -> SerialT m (Array a)
splitOn :: (a -> Bool) -> Array a -> SerialT m (Array a)
splitOn a -> Bool
predicate Array a
arr =
    Stream m (Array a) -> SerialT m (Array a)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
IsStream.fromStreamD
        (Stream m (Array a) -> SerialT m (Array a))
-> Stream m (Array a) -> SerialT m (Array a)
forall a b. (a -> b) -> a -> b
$ ((Int, Int) -> Array a)
-> Stream m (Int, Int) -> Stream m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
i, Int
len) -> Int -> Int -> Array a -> Array a
forall a. Storable a => Int -> Int -> Array a -> Array a
getSliceUnsafe Int
i Int
len Array a
arr)
        (Stream m (Int, Int) -> Stream m (Array a))
-> Stream m (Int, Int) -> Stream m (Array a)
forall a b. (a -> b) -> a -> b
$ (a -> Bool) -> Stream m a -> Stream m (Int, Int)
forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m (Int, Int)
D.sliceOnSuffix a -> Bool
predicate (Array a -> Stream m a
forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Array a -> Stream m a
toStreamD Array a
arr)

-- | Generate a stream of array slice descriptors ((index, len)) of specified
-- length from an array, starting from the supplied array index. The last slice
-- may be shorter than the requested length depending on the array length.
--
-- /Pre-release/
{-# INLINE genSlicesFromLen #-}
genSlicesFromLen :: forall m a. (Monad m, Storable a)
    => Int -- ^ from index
    -> Int -- ^ length of the slice
    -> Unfold m (Array a) (Int, Int)
genSlicesFromLen :: Int -> Int -> Unfold m (Array a) (Int, Int)
genSlicesFromLen Int
from Int
len =
    let fromThenTo :: c -> (Int, Int, c)
fromThenTo c
n = (Int
from, Int
from Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len, c
n c -> c -> c
forall a. Num a => a -> a -> a
- c
1)
        mkSlice :: Int -> Int -> m (Int, Int)
mkSlice Int
n Int
i = (Int, Int) -> m (Int, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i, Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
len (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i))
     in (Array a -> Int)
-> Unfold m Int (Int, Int) -> Unfold m (Array a) (Int, Int)
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
Unfold.lmap Array a -> Int
forall a. Storable a => Array a -> Int
length
        (Unfold m Int (Int, Int) -> Unfold m (Array a) (Int, Int))
-> Unfold m Int (Int, Int) -> Unfold m (Array a) (Int, Int)
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> m (Int, Int))
-> Unfold m Int Int -> Unfold m Int (Int, Int)
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Unfold m a b -> Unfold m a c
Unfold.mapMWithInput Int -> Int -> m (Int, Int)
forall (m :: * -> *). Monad m => Int -> Int -> m (Int, Int)
mkSlice
        (Unfold m Int Int -> Unfold m Int (Int, Int))
-> Unfold m Int Int -> Unfold m Int (Int, Int)
forall a b. (a -> b) -> a -> b
$ (Int -> (Int, Int, Int))
-> Unfold m (Int, Int, Int) Int -> Unfold m Int Int
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
Unfold.lmap Int -> (Int, Int, Int)
forall c. Num c => c -> (Int, Int, c)
fromThenTo Unfold m (Int, Int, Int) Int
forall a (m :: * -> *).
(Enumerable a, Monad m) =>
Unfold m (a, a, a) a
Unfold.enumerateFromThenTo

-- | Generate a stream of slices of specified length from an array, starting
-- from the supplied array index. The last slice may be shorter than the
-- requested length depending on the array length.
--
-- /Pre-release/
{-# INLINE getSlicesFromLen #-}
getSlicesFromLen :: forall m a. (Monad m, Storable a)
    => Int -- ^ from index
    -> Int -- ^ length of the slice
    -> Unfold m (Array a) (Array a)
getSlicesFromLen :: Int -> Int -> Unfold m (Array a) (Array a)
getSlicesFromLen Int
from Int
len =
    let mkSlice :: Array a -> (Int, Int) -> m (Array a)
mkSlice Array a
arr (Int
i, Int
n) = Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array a -> m (Array a)) -> Array a -> m (Array a)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Array a -> Array a
forall a. Storable a => Int -> Int -> Array a -> Array a
getSliceUnsafe Int
i Int
n Array a
arr
     in (Array a -> (Int, Int) -> m (Array a))
-> Unfold m (Array a) (Int, Int) -> Unfold m (Array a) (Array a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Unfold m a b -> Unfold m a c
Unfold.mapMWithInput Array a -> (Int, Int) -> m (Array a)
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Array a -> (Int, Int) -> m (Array a)
mkSlice (Int -> Int -> Unfold m (Array a) (Int, Int)
forall (m :: * -> *) a.
(Monad m, Storable a) =>
Int -> Int -> Unfold m (Array a) (Int, Int)
genSlicesFromLen Int
from Int
len)