git.iko.soy/streamly
on 89b95bd help sign in
streamly/core/src/Streamly/Internal/Data/Array/Generic.hs
377 lines 12 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Array.Generic
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.Generic
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 , reader
35 , toList
36 , readStreamD
37 , readRevStreamD
38 , read
39 , readRev
41 , foldl'
42 , foldr
43 , streamFold
44 , fold
46 -- * Random Access
47 , getIndexUnsafe
48 , strip
49 , putIndices
50 )
51where
53import Control.Monad (when, replicateM)
54import Control.Monad.IO.Class (liftIO, MonadIO)
55import GHC.Base (MutableArray#, RealWorld)
56import GHC.IO (unsafePerformIO)
57import Text.Read (readPrec)
59import Streamly.Internal.Data.Fold.Type (Fold(..))
60import Streamly.Internal.Data.Stream.Type (Stream)
61import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
62import Streamly.Internal.Data.Unfold.Type (Unfold(..))
63import Streamly.Internal.System.IO (unsafeInlineIO)
65import qualified Streamly.Internal.Data.Array.Generic.Mut.Type as MArray
66import qualified Streamly.Internal.Data.Fold.Type as FL
67import qualified Streamly.Internal.Data.Producer.Type as Producer
68import qualified Streamly.Internal.Data.Producer as Producer
69import qualified Streamly.Internal.Data.Ring as RB
70import qualified Streamly.Internal.Data.Stream.StreamD as D
71import qualified Streamly.Internal.Data.Stream.Type as Stream
72import qualified Text.ParserCombinators.ReadPrec as ReadPrec
74import Data.IORef
75import Prelude hiding (foldr, length, read)
77-------------------------------------------------------------------------------
78-- Array Data Type
79-------------------------------------------------------------------------------
81data Array a =
82 Array
83 { arrContents# :: MutableArray# RealWorld a
84 -- ^ The internal contents of the array representing the entire array.
86 , arrStart :: {-# UNPACK #-}!Int
87 -- ^ The starting index of this slice.
89 , arrLen :: {-# UNPACK #-}!Int
90 -- ^ The length of this slice.
91 }
93unsafeFreeze :: MArray.Array a -> Array a
94unsafeFreeze (MArray.Array cont# arrS arrL _) = Array cont# arrS arrL
96unsafeThaw :: Array a -> MArray.Array a
97unsafeThaw (Array cont# arrS arrL) = MArray.Array cont# arrS arrL arrL
99{-# NOINLINE nil #-}
100nil :: Array a
101nil = unsafePerformIO $ do
102 marr <- MArray.newArray 0
103 return $ unsafeFreeze marr
105-------------------------------------------------------------------------------
106-- Construction - Folds
107-------------------------------------------------------------------------------
109{-# INLINE_NORMAL writeN #-}
110writeN :: MonadIO m => Int -> Fold m a (Array a)
111writeN = fmap unsafeFreeze <$> MArray.writeN
113{-# INLINE_NORMAL write #-}
114write :: MonadIO m => Fold m a (Array a)
115write = Fold step initial extract
116 where
117 initial = do
118 marr <- liftIO $ MArray.newArray 0
119 return $ FL.Partial (Tuple3' marr 0 0)
120 step (Tuple3' marr i capacity) x
121 | i == capacity =
122 let newCapacity = max (capacity * 2) 1
123 in do newMarr <- liftIO $ MArray.realloc newCapacity marr
124 marr1 <- liftIO $ MArray.snocUnsafe newMarr x
125 return $ FL.Partial $ Tuple3' marr1 (i + 1) newCapacity
126 | otherwise = do
127 marr1 <- liftIO $ MArray.snocUnsafe marr x
128 return $ FL.Partial $ Tuple3' marr1 (i + 1) capacity
129 extract (Tuple3' marr _ _) =
130 return $ unsafeFreeze marr
132-------------------------------------------------------------------------------
133-- Construction - from streams
134-------------------------------------------------------------------------------
136{-# INLINE_NORMAL fromStreamDN #-}
137fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
138fromStreamDN limit str = do
139 marr <- liftIO $ MArray.newArray (max limit 0)
140 i <-
141 D.foldlM'
142 (\i x -> i `seq` liftIO $ MArray.putIndexUnsafe marr i x >> return (i + 1))
143 (return 0) $
144 D.take limit str
145 return $ unsafeFreeze $ marr { MArray.arrLen = i }
147{-# INLINE fromStreamD #-}
148fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
149fromStreamD = D.fold write
151{-# INLINE fromStreamN #-}
152fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)
153fromStreamN n m = do
154 when (n < 0) $ error "fromStreamN: negative write count specified"
155 fromStreamDN n $ Stream.toStreamD m
157{-# INLINE fromStream #-}
158fromStream :: MonadIO m => Stream m a -> m (Array a)
159fromStream = fromStreamD . Stream.toStreamD
161{-# INLINABLE fromListN #-}
162fromListN :: Int -> [a] -> Array a
163fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
165{-# INLINABLE fromList #-}
166fromList :: [a] -> Array a
167fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
169-------------------------------------------------------------------------------
170-- Elimination - Unfolds
171-------------------------------------------------------------------------------
173{-# INLINE length #-}
174length :: Array a -> Int
175length = arrLen
177{-# INLINE_NORMAL reader #-}
178reader :: Monad m => Unfold m (Array a) a
179reader =
180 Producer.simplify
181 $ Producer.translate unsafeThaw unsafeFreeze
182 $ MArray.producerWith (return . unsafeInlineIO)
184-------------------------------------------------------------------------------
185-- Elimination - to streams
186-------------------------------------------------------------------------------
188{-# INLINE_NORMAL toList #-}
189toList :: Array a -> [a]
190toList arr = loop 0
192 where
194 len = length arr
195 loop i | i == len = []
196 loop i = getIndexUnsafe arr i : loop (i + 1)
198{-# INLINE_NORMAL readStreamD #-}
199readStreamD :: MonadIO m => Array a -> D.Stream m a
200readStreamD = MArray.toStreamD . unsafeThaw
202{-# INLINE_NORMAL readRevStreamD #-}
203readRevStreamD :: Monad m => Array a -> D.Stream m a
204readRevStreamD arr@Array{..} =
205 D.map (getIndexUnsafe arr)
206 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
208{-# INLINE_EARLY read #-}
209read :: MonadIO m => Array a -> Stream m a
210read = Stream.fromStreamD . readStreamD
212{-# INLINE_EARLY readRev #-}
213readRev :: Monad m => Array a -> Stream m a
214readRev = Stream.fromStreamD . readRevStreamD
216-------------------------------------------------------------------------------
217-- Elimination - using Folds
218-------------------------------------------------------------------------------
220{-# INLINE_NORMAL foldl' #-}
221foldl' :: (b -> a -> b) -> b -> Array a -> b
222foldl' f z arr = unsafePerformIO $ D.foldl' f z $ readStreamD arr
224{-# INLINE_NORMAL foldr #-}
225foldr :: (a -> b -> b) -> b -> Array a -> b
226foldr f z arr = unsafePerformIO $ D.foldr f z $ readStreamD arr
228{-# INLINE fold #-}
229fold :: MonadIO m => Fold m a b -> Array a -> m b
230fold f arr = D.fold f (readStreamD arr)
232{-# INLINE streamFold #-}
233streamFold :: MonadIO m => (Stream m a -> m b) -> Array a -> m b
234streamFold f arr = f (read 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