@@ -1,6 +1,5 @@
11{-# LANGUAGE CPP #-}
22{-# LANGUAGE MagicHash #-}
3−{-# LANGUAGE RecordWildCards #-}
43{-# LANGUAGE UnboxedTuples #-}
54
65#include "inline.hs"
@@ -20,18 +19,27 @@ module Streamly.Internal.Data.Array
2019 , foldl'
2120 , foldr
2221
22+ , length
23+
2324 , writeN
25+ , write
2426
2527 , toStreamD
2628 , toStreamDRev
29+
30+ , fromListN
31+ , fromList
2732 , fromStreamDN
33+ , fromStreamD
2834 )
2935where
3036
31−import Prelude hiding (foldr)
37+import Prelude hiding (foldr, length)
38+import Control.DeepSeq (NFData(..))
3239import Control.Monad.ST (stToIO, RealWorld)
3340import Control.Monad.IO.Class (liftIO, MonadIO)
3441import GHC.ST (ST(..))
42+import GHC.IO (unsafePerformIO)
3543import GHC.Base (Int(..))
3644import Data.Functor.Identity (runIdentity)
3745import qualified GHC.Exts as Exts
@@ -42,8 +50,7 @@ import qualified Streamly.Streams.StreamD as D
4250
4351data Array a =
4452 Array
45− { aLen :: {-# UNPACK #-} !Int -- Length of the array
46− , aA :: Exts.Array# a
53+ { array# :: Exts.Array# a
4754 }
4855
4956{-# NOINLINE bottomElement #-}
@@ -75,7 +82,30 @@ freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
7582freezeArray marr (I# start#) (I# len#) =
7683 ST (\s# ->
7784 case Exts.freezeArray# (marray# marr) start# len# s# of
78− (# s'#, arr# #) -> (# s'#, Array (I# len#) arr# #))
85+ (# s'#, arr# #) -> (# s'#, Array arr# #))
86+
87+{-# INLINE copyMutableArray #-}
88+copyMutableArray :: MutableArray s a
89+ -- ^ The source array
90+ -> Int
91+ -- ^ Offset into the source array
92+ -> MutableArray s a
93+ -- ^ Destination array to copy to
94+ -> Int
95+ -- ^ Offset into the destination array
96+ -> Int
97+ -- ^ Number of elements to copy
98+ -> ST s ()
99+copyMutableArray sourceMarr (I# sourceOffset#) destMarr (I# destOffset#) (I# numElements#) =
100+ ST (\s# ->
101+ case Exts.copyMutableArray#
102+ (marray# sourceMarr)
103+ sourceOffset#
104+ (marray# destMarr)
105+ destOffset#
106+ numElements#
107+ s# of
108+ s'# -> (# s'#, () #))
79109
80110{-# INLINE newArrayIO #-}
81111newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
@@ -89,28 +119,44 @@ writeArrayIO marr i x = stToIO $ writeArray marr i x
89119freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
90120freezeArrayIO marr i len = stToIO $ freezeArray marr i len
91121
122+{-# INLINE length #-}
123+length :: Array a -> Int
124+length arr = I# (Exts.sizeofArray# (array# arr))
125+
126+{-# INLINE copyMutableArrayIO #-}
127+copyMutableArrayIO ::
128+ MutableArray RealWorld a
129+ -> Int
130+ -> MutableArray RealWorld a
131+ -> Int
132+ -> Int
133+ -> IO ()
134+copyMutableArrayIO sourceMarr sourceOffset destMarr destOffset numElements =
135+ stToIO $
136+ copyMutableArray sourceMarr sourceOffset destMarr destOffset numElements
137+
92138{-# INLINE_NORMAL toStreamD #-}
93139toStreamD :: Monad m => Array a -> D.Stream m a
94−toStreamD Array {..} = D.Stream step 0
140+toStreamD arr = D.Stream step 0
95141 where
96142 {-# INLINE_LATE step #-}
97143 step _ i
98− | i == aLen = return D.Stop
144+ | i == length arr = return D.Stop
99145 step _ (I# i) =
100146 return $
101− case Exts.indexArray# aA i of
147+ case Exts.indexArray# (array# arr) i of
102148 (# x #) -> D.Yield x ((I# i) + 1)
103149
104150{-# INLINE_NORMAL toStreamDRev #-}
105151toStreamDRev :: Monad m => Array a -> D.Stream m a
106−toStreamDRev Array {..} = D.Stream step (aLen - 1)
152+toStreamDRev arr = D.Stream step (length arr - 1)
107153 where
108154 {-# INLINE_LATE step #-}
109155 step _ i
110156 | i < 0 = return D.Stop
111157 step _ (I# i) =
112158 return $
113− case Exts.indexArray# aA i of
159+ case Exts.indexArray# (array# arr) i of
114160 (# x #) -> D.Yield x ((I# i) - 1)
115161
116162{-# INLINE_NORMAL foldl' #-}
@@ -127,13 +173,33 @@ writeN :: MonadIO m => Int -> Fold m a (Array a)
127173writeN limit = Fold step initial extract
128174 where
129175 initial = do
130− marr <- liftIO $ newArrayIO limit bottomElement
131− return (marr, 0)
132− step (marr,i) x | i == limit = return (marr, i)
133− | otherwise = do
134− liftIO $ writeArrayIO marr i x
135− return (marr, i+1)
136− extract (marr,len) = liftIO $ freezeArrayIO marr 0 len
176+ marr <- liftIO $ newArrayIO limit bottomElement
177+ return (marr, 0)
178+ step (marr, i) x
179+ | i == limit = return (marr, i)
180+ | otherwise = do
181+ liftIO $ writeArrayIO marr i x
182+ return (marr, i + 1)
183+ extract (marr, len) = liftIO $ freezeArrayIO marr 0 len
184+
185+{-# INLINE_NORMAL write #-}
186+write :: MonadIO m => Fold m a (Array a)
187+write = Fold step initial extract
188+ where
189+ initial = do
190+ marr <- liftIO $ newArrayIO 0 bottomElement
191+ return (marr, 0, 0)
192+ step (marr, i, capacity) x
193+ | i == capacity =
194+ let newCapacity = max (capacity * 2) 1
195+ in do newMarr <- liftIO $ newArrayIO newCapacity bottomElement
196+ liftIO $ copyMutableArrayIO marr 0 newMarr 0 i
197+ liftIO $ writeArrayIO newMarr i x
198+ return (newMarr, i + 1, newCapacity)
199+ | otherwise = do
200+ liftIO $ writeArrayIO marr i x
201+ return (marr, i + 1, capacity)
202+ extract (marr, len, _) = liftIO $ freezeArrayIO marr 0 len
137203
138204{-# INLINE_NORMAL fromStreamDN #-}
139205fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
@@ -145,3 +211,19 @@ fromStreamDN limit str = do
145211 0 $
146212 D.take limit str
147213 liftIO $ freezeArrayIO marr 0 i
214+
215+{-# INLINE fromStreamD #-}
216+fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
217+fromStreamD str = D.runFold write str
218+
219+{-# INLINABLE fromListN #-}
220+fromListN :: Int -> [a] -> Array a
221+fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
222+
223+{-# INLINABLE fromList #-}
224+fromList :: [a] -> Array a
225+fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
226+
227+instance NFData a => NFData (Array a) where
228+ {-# INLINE rnf #-}
229+ rnf = foldl' (\_ x -> rnf x) ()