@@ -1,6 +1,5 @@
11{-# OPTIONS_GHC -fno-warn-orphans #-}
22
3−{-# LANGUAGE UnboxedTuples #-}
43#include "inline.hs"
54
65-- |
@@ -33,6 +32,7 @@ module Streamly.Internal.Data.Array
3332 , length
3433 , read
3534
35+ , toList
3636 , toStreamD
3737 , toStreamDRev
3838 , toStream
@@ -50,40 +50,56 @@ module Streamly.Internal.Data.Array
5050 )
5151where
5252
53−#if !MIN_VERSION_primitive(0,7,1)
5453import Control.DeepSeq (NFData(..))
55−#endif
56−import Control.Monad (when)
54+import Control.Monad (when, replicateM)
5755import Control.Monad.IO.Class (liftIO, MonadIO)
58−import Control.Monad.Primitive (PrimMonad)
59−import Data.Functor.Identity (runIdentity)
60−import Data.IORef
61−import GHC.Base (Int(..))
56+import GHC.Base (MutableArray#, RealWorld)
6257import GHC.IO (unsafePerformIO)
58+import Text.Read (readPrec)
6359
6460import Streamly.Internal.Data.Fold.Type (Fold(..))
6561import Streamly.Internal.Data.Stream.Type (Stream)
6662import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
6763import Streamly.Internal.Data.Unfold.Type (Unfold(..))
6864
69−import qualified GHC.Exts as Exts
65+import qualified Streamly.Internal.Data.Array.Mut.Type as MArray
7066import qualified Streamly.Internal.Data.Fold.Type as FL
7167import qualified Streamly.Internal.Data.Ring as RB
7268import qualified Streamly.Internal.Data.Stream.StreamD as D
7369import qualified Streamly.Internal.Data.Stream.Type as Stream
70+import qualified Streamly.Internal.Data.Unfold as Unfold
71+import qualified Text.ParserCombinators.ReadPrec as ReadPrec
7472
75−import Data.Primitive.Array hiding (fromList, fromListN)
73+import Data.IORef
7674import Prelude hiding (foldr, length, read)
7775
78−{-# NOINLINE bottomElement #-}
79−bottomElement :: a
80−bottomElement = undefined
76+-------------------------------------------------------------------------------
77+-- Array Data Type
78+-------------------------------------------------------------------------------
79+
80+data Array a =
81+ Array
82+ { arrContents# :: MutableArray# RealWorld a
83+ -- ^ The internal contents of the array representing the entire array.
84+
85+ , arrStart :: {-# UNPACK #-}!Int
86+ -- ^ The starting index of this slice.
87+
88+ , arrLen :: {-# UNPACK #-}!Int
89+ -- ^ The length of this slice.
90+ }
91+
92+unsafeFreeze :: MArray.Array a -> Array a
93+unsafeFreeze (MArray.Array cont# arrS arrL _) = Array cont# arrS arrL
94+
95+unsafeThaw :: Array a -> MArray.Array a
96+unsafeThaw (Array cont# arrS arrL) = MArray.Array cont# arrS arrL arrL
8197
8298{-# NOINLINE nil #-}
8399nil :: Array a
84100nil = unsafePerformIO $ do
85− marr <- liftIO $ newArray 0 bottomElement
86− liftIO $ freezeArray marr 0 0
101+ marr <- MArray.newArray 0
102+ return $ unsafeFreeze marr
87103
88104-------------------------------------------------------------------------------
89105-- Construction - Folds
@@ -91,46 +107,26 @@ nil = unsafePerformIO $ do
91107
92108{-# INLINE_NORMAL writeN #-}
93109writeN :: MonadIO m => Int -> Fold m a (Array a)
94−writeN len = Fold step initial extract
95−
96− where
97−
98− {-# INLINE next #-}
99− next marr i = do
100− let i1 = i + 1
101− st = Tuple' marr i1
102− if len > i1
103− then return $ FL.Partial st
104− else fmap FL.Done $ extract st
105−
106− initial = do
107− marr <- liftIO $ newArray len bottomElement
108− next marr (-1)
109−
110− step (Tuple' marr i) x = do
111− liftIO $ writeArray marr i x
112− next marr i
113−
114− extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
110+writeN = fmap unsafeFreeze <$> MArray.writeN
115111
116112{-# INLINE_NORMAL write #-}
117113write :: MonadIO m => Fold m a (Array a)
118114write = Fold step initial extract
119115 where
120116 initial = do
121− marr <- liftIO $ newArray 0 bottomElement
117+ marr <- liftIO $ MArray.newArray 0
122118 return $ FL.Partial (Tuple3' marr 0 0)
123119 step (Tuple3' marr i capacity) x
124120 | i == capacity =
125121 let newCapacity = max (capacity * 2) 1
126− in do newMarr <- liftIO $ newArray newCapacity bottomElement
127− liftIO $ copyMutableArray newMarr 0 marr 0 i
128− liftIO $ writeArray newMarr i x
129− return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
122+ in do newMarr <- liftIO $ MArray.realloc newCapacity marr
123+ marr1 <- liftIO $ MArray.snocUnsafe newMarr x
124+ return $ FL.Partial $ Tuple3' marr1 (i + 1) newCapacity
130125 | otherwise = do
131− liftIO $ writeArray marr i x
132− return $ FL.Partial $ Tuple3' marr (i + 1) capacity
133− extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
126+ marr1 <- liftIO $ MArray.snocUnsafe marr x
127+ return $ FL.Partial $ Tuple3' marr1 (i + 1) capacity
128+ extract (Tuple3' marr _ _) =
129+ return $ unsafeFreeze marr
134130
135131-------------------------------------------------------------------------------
136132-- Construction - from streams
@@ -139,13 +135,13 @@ write = Fold step initial extract
139135{-# INLINE_NORMAL fromStreamDN #-}
140136fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
141137fromStreamDN limit str = do
142− marr <- liftIO $ newArray (max limit 0) bottomElement
138+ marr <- liftIO $ MArray.newArray (max limit 0)
143139 i <-
144140 D.foldlM'
145− (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
141+ (\i x -> i `seq` liftIO $ MArray.putIndexUnsafe marr i x >> return (i + 1))
146142 (return 0) $
147143 D.take limit str
148− liftIO $ freezeArray marr 0 i
144+ return $ unsafeFreeze $ marr { MArray.arrLen = i }
149145
150146{-# INLINE fromStreamD #-}
151147fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
@@ -175,50 +171,38 @@ fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
175171
176172{-# INLINE length #-}
177173length :: Array a -> Int
178−length = sizeofArray
174+length = arrLen
179175
180176{-# INLINE_NORMAL read #-}
181−read :: Monad m => Unfold m (Array a) a
182−read = Unfold step inject
183− where
184− inject arr = return (arr, 0)
185− step (arr, i)
186− | i == length arr = return D.Stop
187− step (arr, I# i) =
188− return $
189− case Exts.indexArray# (array# arr) i of
190− (# x #) -> D.Yield x (arr, I# i + 1)
177+read :: MonadIO m => Unfold m (Array a) a
178+read = Unfold.lmap unsafeThaw MArray.read
191179
192180-------------------------------------------------------------------------------
193181-- Elimination - to streams
194182-------------------------------------------------------------------------------
195183
184+{-# INLINE_NORMAL toList #-}
185+toList :: Array a -> [a]
186+toList arr = loop 0
187+
188+ where
189+
190+ len = length arr
191+ loop i | i == len = []
192+ loop i = getIndexUnsafe arr i : loop (i + 1)
193+
196194{-# INLINE_NORMAL toStreamD #-}
197−toStreamD :: Monad m => Array a -> D.Stream m a
198−toStreamD arr = D.Stream step 0
199− where
200− {-# INLINE_LATE step #-}
201− step _ i
202− | i == length arr = return D.Stop
203− step _ (I# i) =
204− return $
205− case Exts.indexArray# (array# arr) i of
206− (# x #) -> D.Yield x (I# i + 1)
195+toStreamD :: MonadIO m => Array a -> D.Stream m a
196+toStreamD = MArray.toStreamD . unsafeThaw
207197
208198{-# INLINE_NORMAL toStreamDRev #-}
209199toStreamDRev :: Monad m => Array a -> D.Stream m a
210−toStreamDRev arr = D.Stream step (length arr - 1)
211− where
212− {-# INLINE_LATE step #-}
213− step _ i
214− | i < 0 = return D.Stop
215− step _ (I# i) =
216− return $
217− case Exts.indexArray# (array# arr) i of
218− (# x #) -> D.Yield x (I# i - 1)
200+toStreamDRev arr@Array{..} =
201+ D.map (getIndexUnsafe arr)
202+ $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
219203
220204{-# INLINE_EARLY toStream #-}
221−toStream :: Monad m => Array a -> Stream m a
205+toStream :: MonadIO m => Array a -> Stream m a
222206toStream = Stream.fromStreamD . toStreamD
223207
224208{-# INLINE_EARLY toStreamRev #-}
@@ -231,24 +215,22 @@ toStreamRev = Stream.fromStreamD . toStreamDRev
231215
232216{-# INLINE_NORMAL foldl' #-}
233217foldl' :: (b -> a -> b) -> b -> Array a -> b
234−foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
218+foldl' f z arr = unsafePerformIO $ D.foldl' f z $ toStreamD arr
235219
236220{-# INLINE_NORMAL foldr #-}
237221foldr :: (a -> b -> b) -> b -> Array a -> b
238−foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
222+foldr f z arr = unsafePerformIO $ D.foldr f z $ toStreamD arr
239223
240−#if !MIN_VERSION_primitive(0,7,1)
241224instance NFData a => NFData (Array a) where
242225 {-# INLINE rnf #-}
243226 rnf = foldl' (\_ x -> rnf x) ()
244−#endif
245227
246228{-# INLINE fold #-}
247−fold :: Monad m => Fold m a b -> Array a -> m b
229+fold :: MonadIO m => Fold m a b -> Array a -> m b
248230fold f arr = D.fold f (toStreamD arr)
249231
250232{-# INLINE streamFold #-}
251−streamFold :: Monad m => (Stream m a -> m b) -> Array a -> m b
233+streamFold :: MonadIO m => (Stream m a -> m b) -> Array a -> m b
252234streamFold f arr = f (toStream arr)
253235
254236-------------------------------------------------------------------------------
@@ -261,12 +243,13 @@ streamFold f arr = f (toStream arr)
261243-- @since 0.8.0
262244{-# INLINE getIndexUnsafe #-}
263245getIndexUnsafe :: Array a -> Int -> a
264−getIndexUnsafe = indexArray
246+getIndexUnsafe arr i =
247+ unsafePerformIO $ MArray.getIndexUnsafe (unsafeThaw arr) i
265248
266249{-# INLINE writeLastN #-}
267250writeLastN :: MonadIO m => Int -> Fold m a (Array a)
268251writeLastN n
269− | n <= 0 = fmap (const mempty) FL.drain
252+ | n <= 0 = fmap (const nil) FL.drain
270253 | otherwise = Fold step initial done
271254
272255 where
@@ -280,15 +263,18 @@ writeLastN n
280263 return $ FL.Partial $ Tuple' rb (rh + 1)
281264
282265 done (Tuple' rb rh) = do
283− arr' <- liftIO $ newArray (min rh n) (undefined :: a)
266+ arr' <- liftIO $ MArray.newArray (min rh n)
284267 ref <- liftIO $ readIORef $ RB.ringHead rb
285268 if rh < n
286269 then
287− liftIO $ copyMutableArray arr' 0 (RB.arr rb) 0 ref
270+ MArray.putSliceUnsafe (RB.arr rb) 0 arr' 0 ref
288271 else do
289− liftIO $ copyMutableArray arr' 0 (RB.arr rb) ref (n - ref)
290− liftIO $ copyMutableArray arr' (n - ref) (RB.arr rb) 0 ref
291− liftIO $ unsafeFreezeArray arr'
272+ MArray.putSliceUnsafe (RB.arr rb) ref arr' 0 (n - ref)
273+ MArray.putSliceUnsafe (RB.arr rb) 0 arr' (n - ref) ref
274+ return $ unsafeFreeze arr'
275+
276+sliceUnsafe :: Int -> Int -> Array a -> Array a
277+sliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len
292278
293279-- XXX This is not efficient as it copies the array. We should support array
294280-- slicing so that we can just refer to the underlying array memory instead of
@@ -306,30 +292,86 @@ strip p arr =
306292 let indexL = getIndexL 0 -- first predicate failing index
307293 in if indexL == 0 && indexR == lastIndex
308294 then arr
309− else cloneArray arr indexL (indexR - indexL + 1)
295+ else unsafeFreeze $ unsafePerformIO $ do
296+ let newLen = indexR - indexL + 1
297+ arrThawed = unsafeThaw (sliceUnsafe indexL newLen arr)
298+ MArray.clone arrThawed
310299
311300 where
312301
313302 getIndexR idx
314303 | idx < 0 = idx
315304 | otherwise =
316− if p (indexArray arr idx) then getIndexR (idx - 1) else idx
305+ if p (getIndexUnsafe arr idx) then getIndexR (idx - 1) else idx
317306
318− getIndexL idx = if p (indexArray arr idx) then getIndexL (idx + 1) else idx
307+ getIndexL idx = if p (getIndexUnsafe arr idx)
308+ then getIndexL (idx + 1)
309+ else idx
319310
320311-- | Write an input stream of (index, value) pairs to an array. Throws an
321312-- error if any index is out of bounds.
322313--
323314-- /Pre-release/
324315{-# INLINE putIndices #-}
325−putIndices :: PrimMonad m
316+putIndices :: MonadIO m
326317 => Array a -> Fold m (Int, a) ()
327318putIndices arr = FL.rmapM (\ _ -> return ()) (FL.foldlM' step initial)
328319
329320 where
330321
331− initial = unsafeThawArray arr
322+ initial = return $ unsafeThaw arr
332323
333324 step marr (i, x) = do
334− writeArray marr i x
325+ MArray.putIndexUnsafe marr i x
335326 return marr
327+
328+-------------------------------------------------------------------------------
329+-- Instances
330+-------------------------------------------------------------------------------
331+
332+instance Eq a => Eq (Array a) where
333+ a1 == a2 = lenA1 == lenA2 && loop (lenA1 - 1)
334+
335+ where
336+
337+ lenA1 = length a1
338+ lenA2 = length a2
339+
340+ loop i
341+ | i < 0 = True
342+ | otherwise =
343+ let v1 = getIndexUnsafe a1 i
344+ v2 = getIndexUnsafe a2 i
345+ in v1 == v2 && loop (i - 1)
346+
347+instance Ord a => Ord (Array a) where
348+ compare a1 a2 =
349+ case compare lenA1 lenA2 of
350+ LT -> LT
351+ GT -> GT
352+ EQ -> loop 0
353+
354+ where
355+
356+ lenA1 = length a1
357+ lenA2 = length a2
358+
359+ loop i
360+ | i >= lenA1 = EQ
361+ | otherwise =
362+ let v1 = getIndexUnsafe a1 i
363+ v2 = getIndexUnsafe a2 i
364+ in case compare v1 v2 of
365+ LT -> LT
366+ GT -> GT
367+ EQ -> loop (i + 1)
368+
369+instance Show a => Show (Array a) where
370+ show arr = "fromList " ++ show (toList arr)
371+
372+instance Read a => Read (Array a) where
373+ readPrec = do
374+ fromListWord <- replicateM 9 ReadPrec.get
375+ if fromListWord == "fromList "
376+ then fromList <$> readPrec
377+ else ReadPrec.pfail