-- |
-- Module      : Streamly.Internal.Data.SmallArray
-- Copyright   : (c) 2019 Composewell Technologies
--
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC

{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE UnboxedTuples #-}

#include "inline.hs"

module Streamly.Internal.Data.SmallArray
  (
    -- XXX should it be just Array instead? We should be able to replace one
    -- array type with another easily.
    SmallArray(..)

  , foldl'
  , foldr

  , length

  , writeN

  , toStreamD
  , toStreamDRev

  , toStream
  , toStreamRev
  , read

  , fromListN
  , fromStreamDN
  , fromStreamN

  , streamFold
  , fold
  )
where

import Prelude hiding (foldr, length, read)
import Control.DeepSeq (NFData(..))
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO, liftIO)
import GHC.IO (unsafePerformIO)
import Data.Functor.Identity (runIdentity)

import Streamly.Internal.Data.SmallArray.Type

import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Stream.Serial (SerialT(..))

import qualified Streamly.Internal.Data.Stream.StreamD as D
import qualified Streamly.Internal.Data.Fold.Type as FL

{-# NOINLINE bottomElement #-}
bottomElement :: a
bottomElement :: forall a. a
bottomElement = forall a. HasCallStack => a
undefined

{-# INLINE length #-}
length :: SmallArray a -> Int
length :: forall a. SmallArray a -> Int
length = forall a. SmallArray a -> Int
sizeofSmallArray

{-# INLINE_NORMAL toStreamD #-}
toStreamD :: Monad m => SmallArray a -> D.Stream m a
toStreamD :: forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamD SmallArray a
arr = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {p}. Monad m => p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. SmallArray a -> Int
length SmallArray a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
        | Bool
otherwise =
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            case forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i of
                (# a
x #) -> forall s a. a -> s -> Step s a
D.Yield a
x (Int
i forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE_NORMAL toStreamDRev #-}
toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
toStreamDRev :: forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamDRev SmallArray a
arr = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream forall {m :: * -> *} {p}. Monad m => p -> Int -> m (Step Int a)
step (forall a. SmallArray a -> Int
length SmallArray a
arr forall a. Num a => a -> a -> a
- Int
1)
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
        | Bool
otherwise =
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            case forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i of
                (# a
x #) -> forall s a. a -> s -> Step s a
D.Yield a
x (Int
i forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL foldl' #-}
foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> SmallArray a -> b
foldl' b -> a -> b
f b
z SmallArray a
arr = forall a. Identity a -> a
runIdentity 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 => SmallArray a -> Stream m a
toStreamD SmallArray a
arr

{-# INLINE_NORMAL foldr #-}
foldr :: (a -> b -> b) -> b -> SmallArray a -> b
foldr :: forall a b. (a -> b -> b) -> b -> SmallArray a -> b
foldr a -> b -> b
f b
z SmallArray a
arr = forall a. Identity a -> a
runIdentity 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 => SmallArray a -> Stream m a
toStreamD SmallArray a
arr

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'SmallArray'.
--
-- Since we are folding to a 'SmallArray' @n@ should be <= 128, for larger number
-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Data.Array.Foreign".
{-# INLINE_NORMAL writeN #-}
writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
writeN :: forall (m :: * -> *) a. MonadIO m => Int -> Fold m a (SmallArray a)
writeN Int
len = 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 =>
Tuple' (SmallMutableArray RealWorld a) Int
-> a
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
step forall {a}.
m (Step
     (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
initial forall {m :: * -> *} {a}.
MonadIO m =>
Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
extract

    where

    {-# INLINE next #-}
    next :: SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
next SmallMutableArray RealWorld a
marr Int
i = do
        let i1 :: Int
i1 = Int
i forall a. Num a => a -> a -> a
+ Int
1
            st :: Tuple' (SmallMutableArray RealWorld a) Int
st = forall a b. a -> b -> Tuple' a b
Tuple' SmallMutableArray RealWorld a
marr Int
i1
        if Int
len forall a. Ord a => a -> a -> Bool
> Int
i1
        then 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 Tuple' (SmallMutableArray RealWorld a) Int
st
        else forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s b. b -> Step s b
FL.Done forall a b. (a -> b) -> a -> b
$ forall {m :: * -> *} {a}.
MonadIO m =>
Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
extract Tuple' (SmallMutableArray RealWorld a) Int
st

    initial :: m (Step
     (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
initial = do
        SmallMutableArray RealWorld a
marr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len forall a. a
bottomElement
        forall {m :: * -> *} {a}.
MonadIO m =>
SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
next SmallMutableArray RealWorld a
marr (-Int
1)

    step :: Tuple' (SmallMutableArray RealWorld a) Int
-> a
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
step (Tuple' SmallMutableArray RealWorld a
marr Int
i) a
x = do
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld a
marr Int
i a
x
        forall {m :: * -> *} {a}.
MonadIO m =>
SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
next SmallMutableArray RealWorld a
marr Int
i

    extract :: Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
extract (Tuple' SmallMutableArray RealWorld a
marr Int
l) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray RealWorld a
marr Int
0 Int
l


{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
fromStreamDN :: forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (SmallArray a)
fromStreamDN Int
limit Stream m a
str = do
    SmallMutableArray RealWorld a
marr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (forall a. Ord a => a -> a -> a
max Int
limit Int
0) forall a. a
bottomElement
    Int
i <-
        forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
D.foldlM'
            (\Int
i a
x -> Int
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld a
marr Int
i a
x) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i forall a. Num a => a -> a -> a
+ Int
1))
            (forall (m :: * -> *) a. Monad m => a -> m a
return Int
0) forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) a.
Applicative m =>
Int -> Stream m a -> Stream m a
D.take Int
limit Stream m a
str
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray RealWorld a
marr Int
0 Int
i

-- | Create a 'SmallArray' from the first @n@ elements of a list. The
-- array may hold less than @n@ elements if the length of the list <=
-- @n@.
--
-- It is recommended to use a value of @n@ <= 128. For larger sized
-- arrays, use an 'Array' from "Streamly.Data.Array" or
-- "Streamly.Data.Array.Foreign"
{-# INLINABLE fromListN #-}
fromListN :: Int -> [a] -> SmallArray a
fromListN :: forall a. Int -> [a] -> SmallArray 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 (SmallArray a)
fromStreamDN Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

instance NFData a => NFData (SmallArray a) where
    {-# INLINE rnf #-}
    rnf :: SmallArray a -> ()
rnf = forall b a. (b -> a -> b) -> b -> SmallArray a -> b
foldl' (\()
_ a
x -> forall a. NFData a => a -> ()
rnf a
x) ()

-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
-- array is allocated to size @n@, if the stream terminates before @n@
-- elements then the array may hold less than @n@ elements.
--
-- For optimal performance use this with @n@ <= 128.
{-# INLINE fromStreamN #-}
fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
fromStreamN :: forall (m :: * -> *) a.
MonadIO m =>
Int -> SerialT m a -> m (SmallArray a)
fromStreamN Int
n (SerialT Stream m a
m) = do
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => [Char] -> a
error [Char]
"fromStreamN: negative write count specified"
    forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (SmallArray a)
fromStreamDN Int
n forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Applicative m => Stream m a -> Stream m a
D.fromStreamK Stream m a
m

{-# INLINE_EARLY toStream #-}
toStream :: Monad m => SmallArray a -> SerialT m a
toStream :: forall (m :: * -> *) a. Monad m => SmallArray a -> SerialT m a
toStream = forall (m :: * -> *) a. Stream m a -> SerialT m a
SerialT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
D.toStreamK forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamD

{-# INLINE_EARLY toStreamRev #-}
toStreamRev :: Monad m => SmallArray a -> SerialT m a
toStreamRev :: forall (m :: * -> *) a. Monad m => SmallArray a -> SerialT m a
toStreamRev = forall (m :: * -> *) a. Stream m a -> SerialT m a
SerialT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
D.toStreamK forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamDRev

{-# INLINE fold #-}
fold :: Monad m => Fold m a b -> SmallArray a -> m b
fold :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> SmallArray a -> m b
fold Fold m a b
f SmallArray 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 => SmallArray a -> Stream m a
toStreamD SmallArray a
arr)

{-# INLINE streamFold #-}
streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
streamFold :: forall (m :: * -> *) a b.
Monad m =>
(SerialT m a -> m b) -> SmallArray a -> m b
streamFold SerialT m a -> m b
f SmallArray a
arr = SerialT m a -> m b
f (forall (m :: * -> *) a. Monad m => SmallArray a -> SerialT m a
toStream SmallArray a
arr)

{-# INLINE_NORMAL read #-}
read :: Monad m => Unfold m (SmallArray a) a
read :: forall (m :: * -> *) a. Monad m => Unfold m (SmallArray a) a
read = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {m :: * -> *} {a}.
Monad m =>
(SmallArray a, Int) -> m (Step (SmallArray a, Int) a)
step forall {m :: * -> *} {b} {a}. (Monad m, Num b) => a -> m (a, b)
inject
  where
    inject :: a -> m (a, b)
inject a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return (a
arr, b
0)
    step :: (SmallArray a, Int) -> m (Step (SmallArray a, Int) a)
step (SmallArray a
arr, Int
i)
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. SmallArray a -> Int
length SmallArray a
arr = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
D.Stop
        | Bool
otherwise =
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            case forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i of
                (# a
x #) -> forall s a. a -> s -> Step s a
D.Yield a
x (SmallArray a
arr, Int
i forall a. Num a => a -> a -> a
+ Int
1)