@@ -0,0 +1,147 @@
2+{-# LANGUAGE MagicHash #-}
3+{-# LANGUAGE RecordWildCards #-}
4+{-# LANGUAGE UnboxedTuples #-}
9+-- Module : Streamly.Internal.Data.Array
10+-- Copyright : (c) 2019 Composewell Technologies
12+-- License : BSD-3-Clause
13+-- Maintainer : streamly@composewell.com
14+-- Stability : experimental
17+module Streamly.Internal.Data.Array
31+import Prelude hiding (foldr)
32+import Control.Monad.ST (stToIO, RealWorld)
33+import Control.Monad.IO.Class (liftIO, MonadIO)
34+import GHC.ST (ST(..))
35+import GHC.Base (Int(..))
36+import Data.Functor.Identity (runIdentity)
37+import qualified GHC.Exts as Exts
39+import Streamly.Internal.Data.Fold.Types (Fold(..))
41+import qualified Streamly.Streams.StreamD as D
45+ { aLen :: {-# UNPACK #-} !Int -- Length of the array
46+ , aA :: Exts.Array# a
49+{-# NOINLINE bottomElement #-}
51+bottomElement = undefined
53+data MutableArray s a =
55+ { marray# :: Exts.MutableArray# s a
58+{-# INLINE newArray #-}
59+newArray :: Int -> a -> ST s (MutableArray s a)
60+newArray (I# n#) x = ST
61+ (\s# -> case Exts.newArray# n# x s# of
63+ let ma = MutableArray marr#
66+{-# INLINE writeArray #-}
67+writeArray :: MutableArray s a -> Int -> a -> ST s ()
68+writeArray marr (I# i#) x =
70+ case Exts.writeArray# (marray# marr) i# x s# of
71+ s'# -> (# s'#, () #))
73+{-# INLINE freezeArray #-}
74+freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
75+freezeArray marr (I# start#) (I# len#) =
77+ case Exts.freezeArray# (marray# marr) start# len# s# of
78+ (# s'#, arr# #) -> (# s'#, Array (I# len#) arr# #))
80+{-# INLINE newArrayIO #-}
81+newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
82+newArrayIO i x = stToIO $ newArray i x
84+{-# INLINE writeArrayIO #-}
85+writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
86+writeArrayIO marr i x = stToIO $ writeArray marr i x
88+{-# INLINE freezeArrayIO #-}
89+freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
90+freezeArrayIO marr i len = stToIO $ freezeArray marr i len
92+{-# INLINE_NORMAL toStreamD #-}
93+toStreamD :: Monad m => Array a -> D.Stream m a
94+toStreamD Array {..} = D.Stream step 0
96+ {-# INLINE_LATE step #-}
98+ | i == aLen = return D.Stop
101+ case Exts.indexArray# aA i of
102+ (# x #) -> D.Yield x ((I# i) + 1)
104+{-# INLINE_NORMAL toStreamDRev #-}
105+toStreamDRev :: Monad m => Array a -> D.Stream m a
106+toStreamDRev Array {..} = D.Stream step (aLen - 1)
108+ {-# INLINE_LATE step #-}
110+ | i < 0 = return D.Stop
113+ case Exts.indexArray# aA i of
114+ (# x #) -> D.Yield x ((I# i) - 1)
116+{-# INLINE_NORMAL foldl' #-}
117+foldl' :: (b -> a -> b) -> b -> Array a -> b
118+foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
120+{-# INLINE_NORMAL foldr #-}
121+foldr :: (a -> b -> b) -> b -> Array a -> b
122+foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
124+-- writeN n = S.evertM (fromStreamDN n)
125+{-# INLINE_NORMAL writeN #-}
126+writeN :: MonadIO m => Int -> Fold m a (Array a)
127+writeN limit = Fold step initial extract
130+ marr <- liftIO $ newArrayIO limit bottomElement
132+ step (marr,i) x | i == limit = return (marr, i)
134+ liftIO $ writeArrayIO marr i x
136+ extract (marr,len) = liftIO $ freezeArrayIO marr 0 len
138+{-# INLINE_NORMAL fromStreamDN #-}
139+fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
140+fromStreamDN limit str = do
141+ marr <- liftIO $ newArrayIO (max limit 0) bottomElement
144+ (\i x -> i `seq` (liftIO $ writeArrayIO marr i x) >> return (i + 1))
147+ liftIO $ freezeArrayIO marr 0 i