-- |
-- Module      : Streamly.Internal.Data.Array.Generic
-- Copyright   : (c) 2019 Composewell Technologies
--
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : pre-release
-- Portability : GHC
--
module Streamly.Internal.Data.Array.Generic
    ( Array(..)

    -- * Construction
    , nil
    , writeN
    , write
    , writeWith
    , writeLastN

    , fromStreamN
    , fromStream

    , fromListN
    , fromList

    -- * Elimination
    , length
    , reader

    , toList
    , read
    , readRev

    , foldl'
    , foldr
    , streamFold
    , fold

    -- * Random Access
    , getIndexUnsafe
    , getSliceUnsafe
    , strip
    )
where

#include "inline.hs"

import Control.Monad (replicateM)
import Control.Monad.IO.Class (MonadIO)
import GHC.Base (MutableArray#, RealWorld)
import GHC.IO (unsafePerformIO)
import Text.Read (readPrec)

import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.System.IO (unsafeInlineIO)

import qualified Streamly.Internal.Data.Array.Generic.Mut.Type as MArray
import qualified Streamly.Internal.Data.Fold.Type as FL
import qualified Streamly.Internal.Data.Producer.Type as Producer
import qualified Streamly.Internal.Data.Producer as Producer
import qualified Streamly.Internal.Data.Ring as RB
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D
import qualified Text.ParserCombinators.ReadPrec as ReadPrec

import Prelude hiding (foldr, length, read)

-------------------------------------------------------------------------------
-- Array Data Type
-------------------------------------------------------------------------------

data Array a =
    Array
        { forall a. Array a -> MutableArray# RealWorld a
arrContents# :: MutableArray# RealWorld a
          -- ^ The internal contents of the array representing the entire array.

        , forall a. Array a -> Int
arrStart :: {-# UNPACK #-}!Int
          -- ^ The starting index of this slice.

        , forall a. Array a -> Int
arrLen :: {-# UNPACK #-}!Int
          -- ^ The length of this slice.
        }

unsafeFreeze :: MArray.MutArray a -> Array a
unsafeFreeze :: forall a. MutArray a -> Array a
unsafeFreeze (MArray.MutArray MutableArray# RealWorld a
cont# Int
arrS Int
arrL Int
_) = forall a. MutableArray# RealWorld a -> Int -> Int -> Array a
Array MutableArray# RealWorld a
cont# Int
arrS Int
arrL

unsafeThaw :: Array a -> MArray.MutArray a
unsafeThaw :: forall a. Array a -> MutArray a
unsafeThaw (Array MutableArray# RealWorld a
cont# Int
arrS Int
arrL) = forall a.
MutableArray# RealWorld a -> Int -> Int -> Int -> MutArray a
MArray.MutArray MutableArray# RealWorld a
cont# Int
arrS Int
arrL Int
arrL

{-# NOINLINE nil #-}
nil :: Array a
nil :: forall a. Array a
nil = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall a. MutArray a -> Array a
unsafeFreeze forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadIO m => m (MutArray a)
MArray.nil

-------------------------------------------------------------------------------
-- Construction - Folds
-------------------------------------------------------------------------------

{-# INLINE_NORMAL writeN #-}
writeN :: MonadIO m => Int -> Fold m a (Array a)
writeN :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Array a)
writeN = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. MutArray a -> Array a
unsafeFreeze forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
MArray.writeN

{-# INLINE_NORMAL writeWith #-}
writeWith :: MonadIO m => Int -> Fold m a (Array a)
writeWith :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Array a)
writeWith Int
elemCount = forall a. MutArray a -> Array a
unsafeFreeze forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (MutArray a)
MArray.writeWith Int
elemCount

-- | Fold the whole input to a single array.
--
-- /Caution! Do not use this on infinite streams./
--
{-# INLINE write #-}
write :: MonadIO m => Fold m a (Array a)
write :: forall (m :: * -> *) a. MonadIO m => Fold m a (Array a)
write = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. MutArray a -> Array a
unsafeFreeze forall (m :: * -> *) a. MonadIO m => Fold m a (MutArray a)
MArray.write

-------------------------------------------------------------------------------
-- Construction - from streams
-------------------------------------------------------------------------------

{-# INLINE fromStreamN #-}
fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)
fromStreamN :: forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (Array a)
fromStreamN Int
n = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold (forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Array a)
writeN Int
n)

{-# INLINE fromStream #-}
fromStream :: MonadIO m => Stream m a -> m (Array a)
fromStream :: forall (m :: * -> *) a. MonadIO m => Stream m a -> m (Array a)
fromStream = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold forall (m :: * -> *) a. MonadIO m => Fold m a (Array a)
write

-- XXX Consider foldr/build fusion in toList/fromList

{-# INLINABLE fromListN #-}
fromListN :: Int -> [a] -> Array a
fromListN :: forall a. Int -> [a] -> Array a
fromListN Int
n [a]
xs = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (Array a)
fromStreamN Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

{-# INLINABLE fromList #-}
fromList :: [a] -> Array a
fromList :: forall a. [a] -> Array a
fromList [a]
xs = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Stream m a -> m (Array a)
fromStream forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

-------------------------------------------------------------------------------
-- Elimination - Unfolds
-------------------------------------------------------------------------------

{-# INLINE length #-}
length :: Array a -> Int
length :: forall a. Array a -> Int
length = forall a. Array a -> Int
arrLen

{-# INLINE_NORMAL reader #-}
reader :: Monad m => Unfold m (Array a) a
reader :: forall (m :: * -> *) a. Monad m => Unfold m (Array a) a
reader =
    forall (m :: * -> *) a b. Producer m a b -> Unfold m a b
Producer.simplify
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a c b.
Functor m =>
(a -> c) -> (c -> a) -> Producer m c b -> Producer m a b
Producer.translate forall a. Array a -> MutArray a
unsafeThaw forall a. MutArray a -> Array a
unsafeFreeze
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
Monad m =>
(forall b. IO b -> m b) -> Producer m (MutArray a) a
MArray.producerWith (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IO a -> a
unsafeInlineIO)

-------------------------------------------------------------------------------
-- Elimination - to streams
-------------------------------------------------------------------------------

{-# INLINE_NORMAL toList #-}
toList :: Array a -> [a]
toList :: forall a. Array a -> [a]
toList Array a
arr = Int -> [a]
loop Int
0

    where

    len :: Int
len = forall a. Array a -> Int
length Array a
arr
    loop :: Int -> [a]
loop Int
i | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = []
    loop Int
i = forall a. Int -> Array a -> a
getIndexUnsafe Int
i Array a
arr forall a. a -> [a] -> [a]
: Int -> [a]
loop (Int
i forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE_NORMAL read #-}
read :: Monad m => Array a -> Stream m a
read :: forall (m :: * -> *) a. Monad m => Array a -> Stream m a
read arr :: Array a
arr@Array{Int
MutableArray# RealWorld a
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} =
    forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
D.map (forall a. Int -> Array a -> a
`getIndexUnsafe` Array a
arr) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> Stream m a
D.enumerateFromToIntegral Int
0 (Int
arrLen forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL readRev #-}
readRev :: Monad m => Array a -> Stream m a
readRev :: forall (m :: * -> *) a. Monad m => Array a -> Stream m a
readRev arr :: Array a
arr@Array{Int
MutableArray# RealWorld a
arrLen :: Int
arrStart :: Int
arrContents# :: MutableArray# RealWorld a
arrLen :: forall a. Array a -> Int
arrStart :: forall a. Array a -> Int
arrContents# :: forall a. Array a -> MutableArray# RealWorld a
..} =
    forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
D.map (forall a. Int -> Array a -> a
`getIndexUnsafe` Array a
arr)
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
D.enumerateFromThenToIntegral (Int
arrLen forall a. Num a => a -> a -> a
- Int
1) (Int
arrLen forall a. Num a => a -> a -> a
- Int
2) Int
0

-------------------------------------------------------------------------------
-- Elimination - using Folds
-------------------------------------------------------------------------------

{-# INLINE_NORMAL foldl' #-}
foldl' :: (b -> a -> b) -> b -> Array a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Array a -> b
foldl' b -> a -> b
f b
z Array a
arr = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
D.foldl' b -> a -> b
f b
z forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Array a -> Stream m a
read Array a
arr

{-# INLINE_NORMAL foldr #-}
foldr :: (a -> b -> b) -> b -> Array a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
foldr a -> b -> b
f b
z Array a
arr = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr a -> b -> b
f b
z forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Array a -> Stream m a
read Array a
arr

{-# INLINE fold #-}
fold :: Monad m => Fold m a b -> Array a -> m b
fold :: forall (m :: * -> *) a b. Monad m => Fold m a b -> Array a -> m b
fold Fold m a b
f Array a
arr = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold Fold m a b
f (forall (m :: * -> *) a. Monad m => Array a -> Stream m a
read Array a
arr)

{-# INLINE streamFold #-}
streamFold :: Monad m => (Stream m a -> m b) -> Array a -> m b
streamFold :: forall (m :: * -> *) a b.
Monad m =>
(Stream m a -> m b) -> Array a -> m b
streamFold Stream m a -> m b
f Array a
arr = Stream m a -> m b
f (forall (m :: * -> *) a. Monad m => Array a -> Stream m a
read Array a
arr)

-------------------------------------------------------------------------------
-- Random reads and writes
-------------------------------------------------------------------------------

-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
-- not check the bounds.
--
-- @since 0.8.0
{-# INLINE getIndexUnsafe #-}
getIndexUnsafe :: Int -> Array a -> a
getIndexUnsafe :: forall a. Int -> Array a -> a
getIndexUnsafe Int
i Array a
arr =
    forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => Int -> MutArray a -> m a
MArray.getIndexUnsafe Int
i (forall a. Array a -> MutArray a
unsafeThaw Array a
arr)

{-# INLINE writeLastN #-}
writeLastN :: MonadIO m => Int -> Fold m a (Array a)
writeLastN :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Array a)
writeLastN Int
n = forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> Fold m a b -> Fold m a c
FL.rmapM forall {m :: * -> *} {a}. MonadIO m => Ring a -> m (Array a)
f (forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (Ring a)
RB.writeLastN Int
n)

    where

    f :: Ring a -> m (Array a)
f Ring a
rb = do
        MutArray a
arr <- forall (m :: * -> *) a.
MonadIO m =>
Int -> Int -> Ring a -> m (MutArray a)
RB.toMutArray Int
0 Int
n Ring a
rb
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. MutArray a -> Array a
unsafeFreeze MutArray a
arr

{-# INLINE getSliceUnsafe #-}
getSliceUnsafe :: Int -> Int -> Array a -> Array a
getSliceUnsafe :: forall a. Int -> Int -> Array a -> Array a
getSliceUnsafe Int
offset Int
len (Array MutableArray# RealWorld a
cont Int
off1 Int
_) = forall a. MutableArray# RealWorld a -> Int -> Int -> Array a
Array MutableArray# RealWorld a
cont (Int
off1 forall a. Num a => a -> a -> a
+ Int
offset) Int
len

-- XXX This is not efficient as it copies the array. We should support array
-- slicing so that we can just refer to the underlying array memory instead of
-- copying.

-- | Truncate the array at the beginning and end as long as the predicate
-- holds true. Returns a slice of the original array.
{-# INLINE strip #-}
strip :: (a -> Bool) -> Array a -> Array a
strip :: forall a. (a -> Bool) -> Array a -> Array a
strip a -> Bool
p Array a
arr = forall a. MutArray a -> Array a
unsafeFreeze forall a b. (a -> b) -> a -> b
$ forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
MonadIO m =>
(a -> Bool) -> MutArray a -> m (MutArray a)
MArray.strip a -> Bool
p (forall a. Array a -> MutArray a
unsafeThaw Array a
arr)

-------------------------------------------------------------------------------
-- Instances
-------------------------------------------------------------------------------

instance Eq a => Eq (Array a) where
    {-# INLINE (==) #-}
    Array a
arr1 == :: Array a -> Array a -> Bool
== Array a
arr2 =
        forall a. IO a -> a
unsafeInlineIO forall a b. (a -> b) -> a -> b
$! forall a. Array a -> MutArray a
unsafeThaw Array a
arr1 forall (m :: * -> *) a.
(MonadIO m, Eq a) =>
MutArray a -> MutArray a -> m Bool
`MArray.eq` forall a. Array a -> MutArray a
unsafeThaw Array a
arr2

instance Ord a => Ord (Array a) where
    {-# INLINE compare #-}
    compare :: Array a -> Array a -> Ordering
compare Array a
arr1 Array a
arr2 =
        forall a. IO a -> a
unsafeInlineIO forall a b. (a -> b) -> a -> b
$! forall a. Array a -> MutArray a
unsafeThaw Array a
arr1 forall (m :: * -> *) a.
(MonadIO m, Ord a) =>
MutArray a -> MutArray a -> m Ordering
`MArray.cmp` forall a. Array a -> MutArray a
unsafeThaw Array a
arr2

    -- Default definitions defined in base do not have an INLINE on them, so we
    -- replicate them here with an INLINE.
    {-# INLINE (<) #-}
    Array a
x < :: Array a -> Array a -> Bool
<  Array a
y = case forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
LT -> Bool
True;  Ordering
_ -> Bool
False }

    {-# INLINE (<=) #-}
    Array a
x <= :: Array a -> Array a -> Bool
<= Array a
y = case forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
GT -> Bool
False; Ordering
_ -> Bool
True }

    {-# INLINE (>) #-}
    Array a
x > :: Array a -> Array a -> Bool
>  Array a
y = case forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
GT -> Bool
True;  Ordering
_ -> Bool
False }

    {-# INLINE (>=) #-}
    Array a
x >= :: Array a -> Array a -> Bool
>= Array a
y = case forall a. Ord a => a -> a -> Ordering
compare Array a
x Array a
y of { Ordering
LT -> Bool
False; Ordering
_ -> Bool
True }

    -- These two default methods use '<=' rather than 'compare'
    -- because the latter is often more expensive
    {-# INLINE max #-}
    max :: Array a -> Array a -> Array a
max Array a
x Array a
y = if Array a
x forall a. Ord a => a -> a -> Bool
<= Array a
y then Array a
y else Array a
x

    {-# INLINE min #-}
    min :: Array a -> Array a -> Array a
min Array a
x Array a
y = if Array a
x forall a. Ord a => a -> a -> Bool
<= Array a
y then Array a
x else Array a
y

instance Show a => Show (Array a) where
    {-# INLINE show #-}
    show :: Array a -> String
show Array a
arr = String
"fromList " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall a. Array a -> [a]
toList Array a
arr)

instance Read a => Read (Array a) where
    {-# INLINE readPrec #-}
    readPrec :: ReadPrec (Array a)
readPrec = do
        String
fromListWord <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
9 ReadPrec Char
ReadPrec.get
        if String
fromListWord forall a. Eq a => a -> a -> Bool
== String
"fromList "
        then forall a. [a] -> Array a
fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Read a => ReadPrec a
readPrec
        else forall a. ReadPrec a
ReadPrec.pfail