git.iko.soy/streamly
on d134ea3 help sign in
streamly/core/src/Streamly/Internal/Data/Array/Generic.hs
375 lines 11 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Generic
3-- Copyright : (c) 2019 Composewell Technologies
4--
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : pre-release
8-- Portability : GHC
9--
10module Streamly.Internal.Data.Array.Generic
11 ( Array(..)
13 -- * Construction
14 , nil
15 , writeN
16 , write
17 , writeLastN
19 , fromStreamDN
20 , fromStreamD
21 , fromStreamN
22 , fromStream
24 , fromListN
25 , fromList
27 -- * Elimination
28 , length
29 , reader
31 , toList
32 , readStreamD
33 , readRevStreamD
34 , read
35 , readRev
37 , foldl'
38 , foldr
39 , streamFold
40 , fold
42 -- * Random Access
43 , getIndexUnsafe
44 , strip
45 , putIndices
46 )
47where
49#include "inline.hs"
51import Control.Monad (when, replicateM)
52import Control.Monad.IO.Class (liftIO, MonadIO)
53import GHC.Base (MutableArray#, RealWorld)
54import GHC.IO (unsafePerformIO)
55import Text.Read (readPrec)
57import Streamly.Internal.Data.Fold.Type (Fold(..))
58import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
59import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
60import Streamly.Internal.Data.Unfold.Type (Unfold(..))
61import Streamly.Internal.System.IO (unsafeInlineIO)
63import qualified Streamly.Internal.Data.Array.Generic.Mut.Type as MArray
64import qualified Streamly.Internal.Data.Fold.Type as FL
65import qualified Streamly.Internal.Data.Producer.Type as Producer
66import qualified Streamly.Internal.Data.Producer as Producer
67import qualified Streamly.Internal.Data.Ring as RB
68import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
69import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D
70import qualified Text.ParserCombinators.ReadPrec as ReadPrec
72import Data.IORef
73import Prelude hiding (foldr, length, read)
75-------------------------------------------------------------------------------
76-- Array Data Type
77-------------------------------------------------------------------------------
79data Array a =
80 Array
81 { arrContents# :: MutableArray# RealWorld a
82 -- ^ The internal contents of the array representing the entire array.
84 , arrStart :: {-# UNPACK #-}!Int
85 -- ^ The starting index of this slice.
87 , arrLen :: {-# UNPACK #-}!Int
88 -- ^ The length of this slice.
89 }
91unsafeFreeze :: MArray.Array a -> Array a
92unsafeFreeze (MArray.Array cont# arrS arrL _) = Array cont# arrS arrL
94unsafeThaw :: Array a -> MArray.Array a
95unsafeThaw (Array cont# arrS arrL) = MArray.Array cont# arrS arrL arrL
97{-# NOINLINE nil #-}
98nil :: Array a
99nil = unsafePerformIO $ do
100 marr <- MArray.new 0
101 return $ unsafeFreeze marr
103-------------------------------------------------------------------------------
104-- Construction - Folds
105-------------------------------------------------------------------------------
107{-# INLINE_NORMAL writeN #-}
108writeN :: MonadIO m => Int -> Fold m a (Array a)
109writeN = fmap unsafeFreeze <$> MArray.writeN
111{-# INLINE_NORMAL write #-}
112write :: MonadIO m => Fold m a (Array a)
113write = Fold step initial extract
114 where
115 initial = do
116 marr <- liftIO $ MArray.new 0
117 return $ FL.Partial (Tuple3' marr 0 0)
118 step (Tuple3' marr i capacity) x
119 | i == capacity =
120 let newCapacity = max (capacity * 2) 1
121 in do newMarr <- liftIO $ MArray.realloc newCapacity marr
122 marr1 <- liftIO $ MArray.snocUnsafe newMarr x
123 return $ FL.Partial $ Tuple3' marr1 (i + 1) newCapacity
124 | otherwise = do
125 marr1 <- liftIO $ MArray.snocUnsafe marr x
126 return $ FL.Partial $ Tuple3' marr1 (i + 1) capacity
127 extract (Tuple3' marr _ _) =
128 return $ unsafeFreeze marr
130-------------------------------------------------------------------------------
131-- Construction - from streams
132-------------------------------------------------------------------------------
134{-# INLINE_NORMAL fromStreamDN #-}
135fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
136fromStreamDN limit str = do
137 marr <- liftIO $ MArray.new (max limit 0)
138 i <-
139 D.foldlM'
140 (\i x -> i `seq` liftIO $ MArray.putIndexUnsafe i marr x >> return (i + 1))
141 (return 0) $
142 D.take limit str
143 return $ unsafeFreeze $ marr { MArray.arrLen = i }
145{-# INLINE fromStreamD #-}
146fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
147fromStreamD = D.fold write
149{-# INLINE fromStreamN #-}
150fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)
151fromStreamN n m = do
152 when (n < 0) $ error "fromStreamN: negative write count specified"
153 fromStreamDN n m
155{-# INLINE fromStream #-}
156fromStream :: MonadIO m => Stream m a -> m (Array a)
157fromStream = fromStreamD
159{-# INLINABLE fromListN #-}
160fromListN :: Int -> [a] -> Array a
161fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
163{-# INLINABLE fromList #-}
164fromList :: [a] -> Array a
165fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
167-------------------------------------------------------------------------------
168-- Elimination - Unfolds
169-------------------------------------------------------------------------------
171{-# INLINE length #-}
172length :: Array a -> Int
173length = arrLen
175{-# INLINE_NORMAL reader #-}
176reader :: Monad m => Unfold m (Array a) a
177reader =
178 Producer.simplify
179 $ Producer.translate unsafeThaw unsafeFreeze
180 $ MArray.producerWith (return . unsafeInlineIO)
182-------------------------------------------------------------------------------
183-- Elimination - to streams
184-------------------------------------------------------------------------------
186{-# INLINE_NORMAL toList #-}
187toList :: Array a -> [a]
188toList arr = loop 0
190 where
192 len = length arr
193 loop i | i == len = []
194 loop i = getIndexUnsafe arr i : loop (i + 1)
196{-# INLINE_NORMAL readStreamD #-}
197readStreamD :: MonadIO m => Array a -> D.Stream m a
198readStreamD = MArray.toStreamD . unsafeThaw
200{-# INLINE_NORMAL readRevStreamD #-}
201readRevStreamD :: Monad m => Array a -> D.Stream m a
202readRevStreamD arr@Array{..} =
203 D.map (getIndexUnsafe arr)
204 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
206{-# INLINE_EARLY read #-}
207read :: MonadIO m => Array a -> Stream m a
208read = readStreamD
210{-# INLINE_EARLY readRev #-}
211readRev :: Monad m => Array a -> Stream m a
212readRev = readRevStreamD
214-------------------------------------------------------------------------------
215-- Elimination - using Folds
216-------------------------------------------------------------------------------
218{-# INLINE_NORMAL foldl' #-}
219foldl' :: (b -> a -> b) -> b -> Array a -> b
220foldl' f z arr = unsafePerformIO $ D.foldl' f z $ readStreamD arr
222{-# INLINE_NORMAL foldr #-}
223foldr :: (a -> b -> b) -> b -> Array a -> b
224foldr f z arr = unsafePerformIO $ D.foldr f z $ readStreamD arr
226{-# INLINE fold #-}
227fold :: MonadIO m => Fold m a b -> Array a -> m b
228fold f arr = D.fold f (readStreamD arr)
230{-# INLINE streamFold #-}
231streamFold :: MonadIO m => (Stream m a -> m b) -> Array a -> m b
232streamFold f arr = f (read arr)
234-------------------------------------------------------------------------------
235-- Random reads and writes
236-------------------------------------------------------------------------------
238-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
239-- not check the bounds.
240--
241-- @since 0.8.0
242{-# INLINE getIndexUnsafe #-}
243getIndexUnsafe :: Array a -> Int -> a
244getIndexUnsafe arr i =
245 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)
247{-# INLINE writeLastN #-}
248writeLastN :: MonadIO m => Int -> Fold m a (Array a)
249writeLastN n
250 | n <= 0 = fmap (const nil) FL.drain
251 | otherwise = Fold step initial done
253 where
255 initial = do
256 rb <- liftIO $ RB.createRing n
257 return $ FL.Partial $ Tuple' rb (0 :: Int)
259 step (Tuple' rb rh) x = do
260 liftIO $ RB.unsafeInsertRing rb rh x
261 return $ FL.Partial $ Tuple' rb (rh + 1)
263 done (Tuple' rb rh) = do
264 arr' <- liftIO $ MArray.new (min rh n)
265 ref <- liftIO $ readIORef $ RB.ringHead rb
266 if rh < n
267 then
268 MArray.putSliceUnsafe (RB.arr rb) 0 arr' 0 ref
269 else do
270 MArray.putSliceUnsafe (RB.arr rb) ref arr' 0 (n - ref)
271 MArray.putSliceUnsafe (RB.arr rb) 0 arr' (n - ref) ref
272 return $ unsafeFreeze arr'
274sliceUnsafe :: Int -> Int -> Array a -> Array a
275sliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len
277-- XXX This is not efficient as it copies the array. We should support array
278-- slicing so that we can just refer to the underlying array memory instead of
279-- copying.
280--
281-- | Truncate the array at the beginning and end as long as the predicate
282-- holds true.
283strip :: (a -> Bool) -> Array a -> Array a
284strip p arr =
285 let lastIndex = length arr - 1
286 indexR = getIndexR lastIndex -- last predicate failing index
287 in if indexR == -1
288 then nil
289 else
290 let indexL = getIndexL 0 -- first predicate failing index
291 in if indexL == 0 && indexR == lastIndex
292 then arr
293 else unsafeFreeze $ unsafePerformIO $ do
294 let newLen = indexR - indexL + 1
295 arrThawed = unsafeThaw (sliceUnsafe indexL newLen arr)
296 MArray.clone arrThawed
298 where
300 getIndexR idx
301 | idx < 0 = idx
302 | otherwise =
303 if p (getIndexUnsafe arr idx) then getIndexR (idx - 1) else idx
305 getIndexL idx = if p (getIndexUnsafe arr idx)
306 then getIndexL (idx + 1)
307 else idx
309-- | Write an input stream of (index, value) pairs to an array. Throws an
310-- error if any index is out of bounds.
311--
312-- /Pre-release/
313{-# INLINE putIndices #-}
314putIndices :: MonadIO m
315 => Array a -> Fold m (Int, a) ()
316putIndices arr = FL.rmapM (\ _ -> return ()) (FL.foldlM' step initial)
318 where
320 initial = return $ unsafeThaw arr
322 step marr (i, x) = do
323 MArray.putIndexUnsafe i marr x
324 return marr
326-------------------------------------------------------------------------------
327-- Instances
328-------------------------------------------------------------------------------
330instance Eq a => Eq (Array a) where
331 a1 == a2 = lenA1 == lenA2 && loop (lenA1 - 1)
333 where
335 lenA1 = length a1
336 lenA2 = length a2
338 loop i
339 | i < 0 = True
340 | otherwise =
341 let v1 = getIndexUnsafe a1 i
342 v2 = getIndexUnsafe a2 i
343 in v1 == v2 && loop (i - 1)
345instance Ord a => Ord (Array a) where
346 compare a1 a2 =
347 case compare lenA1 lenA2 of
348 LT -> LT
349 GT -> GT
350 EQ -> loop 0
352 where
354 lenA1 = length a1
355 lenA2 = length a2
357 loop i
358 | i >= lenA1 = EQ
359 | otherwise =
360 let v1 = getIndexUnsafe a1 i
361 v2 = getIndexUnsafe a2 i
362 in case compare v1 v2 of
363 LT -> LT
364 GT -> GT
365 EQ -> loop (i + 1)
367instance Show a => Show (Array a) where
368 show arr = "fromList " ++ show (toList arr)
370instance Read a => Read (Array a) where
371 readPrec = do
372 fromListWord <- replicateM 9 ReadPrec.get
373 if fromListWord == "fromList "
374 then fromList <$> readPrec
375 else ReadPrec.pfail