git.iko.soy/streamly
on 56569bc help sign in
streamly/core/src/Streamly/Internal/Data/Array.hs
377 lines 11 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Array
7-- Copyright : (c) 2019 Composewell Technologies
8--
9-- License : BSD-3-Clause
10-- Maintainer : streamly@composewell.com
11-- Stability : pre-release
12-- Portability : GHC
13--
14module Streamly.Internal.Data.Array
15 ( Array(..)
17 -- * Construction
18 , nil
19 , writeN
20 , write
21 , writeLastN
23 , fromStreamDN
24 , fromStreamD
25 , fromStreamN
26 , fromStream
28 , fromListN
29 , fromList
31 -- * Elimination
32 , length
33 , read
35 , toList
36 , toStreamD
37 , toStreamDRev
38 , toStream
39 , toStreamRev
41 , foldl'
42 , foldr
43 , streamFold
44 , fold
46 -- * Random Access
47 , getIndexUnsafe
48 , strip
49 , putIndices
50 )
51where
53import Control.DeepSeq (NFData(..))
54import Control.Monad (when, replicateM)
55import Control.Monad.IO.Class (liftIO, MonadIO)
56import GHC.Base (MutableArray#, RealWorld)
57import GHC.IO (unsafePerformIO)
58import Text.Read (readPrec)
60import Streamly.Internal.Data.Fold.Type (Fold(..))
61import Streamly.Internal.Data.Stream.Type (Stream)
62import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
63import Streamly.Internal.Data.Unfold.Type (Unfold(..))
65import qualified Streamly.Internal.Data.Array.Mut.Type as MArray
66import qualified Streamly.Internal.Data.Fold.Type as FL
67import qualified Streamly.Internal.Data.Ring as RB
68import qualified Streamly.Internal.Data.Stream.StreamD as D
69import qualified Streamly.Internal.Data.Stream.Type as Stream
70import qualified Streamly.Internal.Data.Unfold as Unfold
71import qualified Text.ParserCombinators.ReadPrec as ReadPrec
73import Data.IORef
74import Prelude hiding (foldr, length, read)
76-------------------------------------------------------------------------------
77-- Array Data Type
78-------------------------------------------------------------------------------
80data Array a =
81 Array
82 { arrContents# :: MutableArray# RealWorld a
83 -- ^ The internal contents of the array representing the entire array.
85 , arrStart :: {-# UNPACK #-}!Int
86 -- ^ The starting index of this slice.
88 , arrLen :: {-# UNPACK #-}!Int
89 -- ^ The length of this slice.
90 }
92unsafeFreeze :: MArray.Array a -> Array a
93unsafeFreeze (MArray.Array cont# arrS arrL _) = Array cont# arrS arrL
95unsafeThaw :: Array a -> MArray.Array a
96unsafeThaw (Array cont# arrS arrL) = MArray.Array cont# arrS arrL arrL
98{-# NOINLINE nil #-}
99nil :: Array a
100nil = unsafePerformIO $ do
101 marr <- MArray.newArray 0
102 return $ unsafeFreeze marr
104-------------------------------------------------------------------------------
105-- Construction - Folds
106-------------------------------------------------------------------------------
108{-# INLINE_NORMAL writeN #-}
109writeN :: MonadIO m => Int -> Fold m a (Array a)
110writeN = fmap unsafeFreeze <$> MArray.writeN
112{-# INLINE_NORMAL write #-}
113write :: MonadIO m => Fold m a (Array a)
114write = Fold step initial extract
115 where
116 initial = do
117 marr <- liftIO $ MArray.newArray 0
118 return $ FL.Partial (Tuple3' marr 0 0)
119 step (Tuple3' marr i capacity) x
120 | i == capacity =
121 let newCapacity = max (capacity * 2) 1
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
125 | otherwise = do
126 marr1 <- liftIO $ MArray.snocUnsafe marr x
127 return $ FL.Partial $ Tuple3' marr1 (i + 1) capacity
128 extract (Tuple3' marr _ _) =
129 return $ unsafeFreeze marr
131-------------------------------------------------------------------------------
132-- Construction - from streams
133-------------------------------------------------------------------------------
135{-# INLINE_NORMAL fromStreamDN #-}
136fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
137fromStreamDN limit str = do
138 marr <- liftIO $ MArray.newArray (max limit 0)
139 i <-
140 D.foldlM'
141 (\i x -> i `seq` liftIO $ MArray.putIndexUnsafe marr i x >> return (i + 1))
142 (return 0) $
143 D.take limit str
144 return $ unsafeFreeze $ marr { MArray.arrLen = i }
146{-# INLINE fromStreamD #-}
147fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
148fromStreamD = D.fold write
150{-# INLINE fromStreamN #-}
151fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)
152fromStreamN n m = do
153 when (n < 0) $ error "fromStreamN: negative write count specified"
154 fromStreamDN n $ Stream.toStreamD m
156{-# INLINE fromStream #-}
157fromStream :: MonadIO m => Stream m a -> m (Array a)
158fromStream = fromStreamD . Stream.toStreamD
160{-# INLINABLE fromListN #-}
161fromListN :: Int -> [a] -> Array a
162fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
164{-# INLINABLE fromList #-}
165fromList :: [a] -> Array a
166fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
168-------------------------------------------------------------------------------
169-- Elimination - Unfolds
170-------------------------------------------------------------------------------
172{-# INLINE length #-}
173length :: Array a -> Int
174length = arrLen
176{-# INLINE_NORMAL read #-}
177read :: MonadIO m => Unfold m (Array a) a
178read = Unfold.lmap unsafeThaw MArray.read
180-------------------------------------------------------------------------------
181-- Elimination - to streams
182-------------------------------------------------------------------------------
184{-# INLINE_NORMAL toList #-}
185toList :: Array a -> [a]
186toList arr = loop 0
188 where
190 len = length arr
191 loop i | i == len = []
192 loop i = getIndexUnsafe arr i : loop (i + 1)
194{-# INLINE_NORMAL toStreamD #-}
195toStreamD :: MonadIO m => Array a -> D.Stream m a
196toStreamD = MArray.toStreamD . unsafeThaw
198{-# INLINE_NORMAL toStreamDRev #-}
199toStreamDRev :: Monad m => Array a -> D.Stream m a
200toStreamDRev arr@Array{..} =
201 D.map (getIndexUnsafe arr)
202 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
204{-# INLINE_EARLY toStream #-}
205toStream :: MonadIO m => Array a -> Stream m a
206toStream = Stream.fromStreamD . toStreamD
208{-# INLINE_EARLY toStreamRev #-}
209toStreamRev :: Monad m => Array a -> Stream m a
210toStreamRev = Stream.fromStreamD . toStreamDRev
212-------------------------------------------------------------------------------
213-- Elimination - using Folds
214-------------------------------------------------------------------------------
216{-# INLINE_NORMAL foldl' #-}
217foldl' :: (b -> a -> b) -> b -> Array a -> b
218foldl' f z arr = unsafePerformIO $ D.foldl' f z $ toStreamD arr
220{-# INLINE_NORMAL foldr #-}
221foldr :: (a -> b -> b) -> b -> Array a -> b
222foldr f z arr = unsafePerformIO $ D.foldr f z $ toStreamD arr
224instance NFData a => NFData (Array a) where
225 {-# INLINE rnf #-}
226 rnf = foldl' (\_ x -> rnf x) ()
228{-# INLINE fold #-}
229fold :: MonadIO m => Fold m a b -> Array a -> m b
230fold f arr = D.fold f (toStreamD arr)
232{-# INLINE streamFold #-}
233streamFold :: MonadIO m => (Stream m a -> m b) -> Array a -> m b
234streamFold f arr = f (toStream arr)
236-------------------------------------------------------------------------------
237-- Random reads and writes
238-------------------------------------------------------------------------------
240-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
241-- not check the bounds.
242--
243-- @since 0.8.0
244{-# INLINE getIndexUnsafe #-}
245getIndexUnsafe :: Array a -> Int -> a
246getIndexUnsafe arr i =
247 unsafePerformIO $ MArray.getIndexUnsafe (unsafeThaw arr) i
249{-# INLINE writeLastN #-}
250writeLastN :: MonadIO m => Int -> Fold m a (Array a)
251writeLastN n
252 | n <= 0 = fmap (const nil) FL.drain
253 | otherwise = Fold step initial done
255 where
257 initial = do
258 rb <- liftIO $ RB.createRing n
259 return $ FL.Partial $ Tuple' rb (0 :: Int)
261 step (Tuple' rb rh) x = do
262 liftIO $ RB.unsafeInsertRing rb rh x
263 return $ FL.Partial $ Tuple' rb (rh + 1)
265 done (Tuple' rb rh) = do
266 arr' <- liftIO $ MArray.newArray (min rh n)
267 ref <- liftIO $ readIORef $ RB.ringHead rb
268 if rh < n
269 then
270 MArray.putSliceUnsafe (RB.arr rb) 0 arr' 0 ref
271 else do
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'
276sliceUnsafe :: Int -> Int -> Array a -> Array a
277sliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len
279-- XXX This is not efficient as it copies the array. We should support array
280-- slicing so that we can just refer to the underlying array memory instead of
281-- copying.
282--
283-- | Truncate the array at the beginning and end as long as the predicate
284-- holds true.
285strip :: (a -> Bool) -> Array a -> Array a
286strip p arr =
287 let lastIndex = length arr - 1
288 indexR = getIndexR lastIndex -- last predicate failing index
289 in if indexR == -1
290 then nil
291 else
292 let indexL = getIndexL 0 -- first predicate failing index
293 in if indexL == 0 && indexR == lastIndex
294 then arr
295 else unsafeFreeze $ unsafePerformIO $ do
296 let newLen = indexR - indexL + 1
297 arrThawed = unsafeThaw (sliceUnsafe indexL newLen arr)
298 MArray.clone arrThawed
300 where
302 getIndexR idx
303 | idx < 0 = idx
304 | otherwise =
305 if p (getIndexUnsafe arr idx) then getIndexR (idx - 1) else idx
307 getIndexL idx = if p (getIndexUnsafe arr idx)
308 then getIndexL (idx + 1)
309 else idx
311-- | Write an input stream of (index, value) pairs to an array. Throws an
312-- error if any index is out of bounds.
313--
314-- /Pre-release/
315{-# INLINE putIndices #-}
316putIndices :: MonadIO m
317 => Array a -> Fold m (Int, a) ()
318putIndices arr = FL.rmapM (\ _ -> return ()) (FL.foldlM' step initial)
320 where
322 initial = return $ unsafeThaw arr
324 step marr (i, x) = do
325 MArray.putIndexUnsafe marr i x
326 return marr
328-------------------------------------------------------------------------------
329-- Instances
330-------------------------------------------------------------------------------
332instance Eq a => Eq (Array a) where
333 a1 == a2 = lenA1 == lenA2 && loop (lenA1 - 1)
335 where
337 lenA1 = length a1
338 lenA2 = length a2
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)
347instance 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
354 where
356 lenA1 = length a1
357 lenA2 = length a2
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)
369instance Show a => Show (Array a) where
370 show arr = "fromList " ++ show (toList arr)
372instance 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