git.iko.soy/streamly
on 1946eb8 help sign in
streamly/core/src/Streamly/Internal/Data/Array/Generic.hs
354 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 , writeWith
18 , writeLastN
20 , fromStreamDN
21 , fromStreamD
22 , fromStreamN
23 , fromStream
25 , fromListN
26 , fromList
28 -- * Elimination
29 , length
30 , reader
32 , toList
33 , readStreamD
34 , readRevStreamD
35 , read
36 , readRev
38 , foldl'
39 , foldr
40 , streamFold
41 , fold
43 -- * Random Access
44 , getIndexUnsafe
45 , strip
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'(..))
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.MutArray a -> Array a
92unsafeFreeze (MArray.MutArray cont# arrS arrL _) = Array cont# arrS arrL
94unsafeThaw :: Array a -> MArray.MutArray a
95unsafeThaw (Array cont# arrS arrL) = MArray.MutArray 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 writeWith #-}
112writeWith :: MonadIO m => Int -> Fold m a (Array a)
113writeWith elemCount = unsafeFreeze <$> MArray.writeWith elemCount
115-- | Fold the whole input to a single array.
116--
117-- /Caution! Do not use this on infinite streams./
118--
119{-# INLINE write #-}
120write :: MonadIO m => Fold m a (Array a)
121write = fmap unsafeFreeze MArray.write
123-------------------------------------------------------------------------------
124-- Construction - from streams
125-------------------------------------------------------------------------------
127{-# INLINE_NORMAL fromStreamDN #-}
128fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
129fromStreamDN limit str = do
130 marr <- liftIO $ MArray.new (max limit 0)
131 i <-
132 D.foldlM'
133 (\i x -> i `seq` liftIO $ MArray.putIndexUnsafe i marr x >> return (i + 1))
134 (return 0) $
135 D.take limit str
136 return $ unsafeFreeze $ marr { MArray.arrLen = i }
138{-# INLINE fromStreamD #-}
139fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
140fromStreamD = D.fold write
142{-# INLINE fromStreamN #-}
143fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)
144fromStreamN n m = do
145 when (n < 0) $ error "fromStreamN: negative write count specified"
146 fromStreamDN n m
148{-# INLINE fromStream #-}
149fromStream :: MonadIO m => Stream m a -> m (Array a)
150fromStream = fromStreamD
152{-# INLINABLE fromListN #-}
153fromListN :: Int -> [a] -> Array a
154fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
156{-# INLINABLE fromList #-}
157fromList :: [a] -> Array a
158fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
160-------------------------------------------------------------------------------
161-- Elimination - Unfolds
162-------------------------------------------------------------------------------
164{-# INLINE length #-}
165length :: Array a -> Int
166length = arrLen
168{-# INLINE_NORMAL reader #-}
169reader :: Monad m => Unfold m (Array a) a
170reader =
171 Producer.simplify
172 $ Producer.translate unsafeThaw unsafeFreeze
173 $ MArray.producerWith (return . unsafeInlineIO)
175-------------------------------------------------------------------------------
176-- Elimination - to streams
177-------------------------------------------------------------------------------
179{-# INLINE_NORMAL toList #-}
180toList :: Array a -> [a]
181toList arr = loop 0
183 where
185 len = length arr
186 loop i | i == len = []
187 loop i = getIndexUnsafe arr i : loop (i + 1)
189{-# INLINE_NORMAL readStreamD #-}
190readStreamD :: MonadIO m => Array a -> D.Stream m a
191readStreamD = MArray.toStreamD . unsafeThaw
193{-# INLINE_NORMAL readRevStreamD #-}
194readRevStreamD :: Monad m => Array a -> D.Stream m a
195readRevStreamD arr@Array{..} =
196 D.map (getIndexUnsafe arr)
197 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
199{-# INLINE_EARLY read #-}
200read :: MonadIO m => Array a -> Stream m a
201read = readStreamD
203{-# INLINE_EARLY readRev #-}
204readRev :: Monad m => Array a -> Stream m a
205readRev = readRevStreamD
207-------------------------------------------------------------------------------
208-- Elimination - using Folds
209-------------------------------------------------------------------------------
211{-# INLINE_NORMAL foldl' #-}
212foldl' :: (b -> a -> b) -> b -> Array a -> b
213foldl' f z arr = unsafePerformIO $ D.foldl' f z $ readStreamD arr
215{-# INLINE_NORMAL foldr #-}
216foldr :: (a -> b -> b) -> b -> Array a -> b
217foldr f z arr = unsafePerformIO $ D.foldr f z $ readStreamD arr
219{-# INLINE fold #-}
220fold :: MonadIO m => Fold m a b -> Array a -> m b
221fold f arr = D.fold f (readStreamD arr)
223{-# INLINE streamFold #-}
224streamFold :: MonadIO m => (Stream m a -> m b) -> Array a -> m b
225streamFold f arr = f (read arr)
227-------------------------------------------------------------------------------
228-- Random reads and writes
229-------------------------------------------------------------------------------
231-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
232-- not check the bounds.
233--
234-- @since 0.8.0
235{-# INLINE getIndexUnsafe #-}
236getIndexUnsafe :: Array a -> Int -> a
237getIndexUnsafe arr i =
238 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)
240{-# INLINE writeLastN #-}
241writeLastN :: MonadIO m => Int -> Fold m a (Array a)
242writeLastN n
243 | n <= 0 = fmap (const nil) FL.drain
244 | otherwise = Fold step initial done
246 where
248 initial = do
249 rb <- liftIO $ RB.createRing n
250 return $ FL.Partial $ Tuple' rb (0 :: Int)
252 step (Tuple' rb cnt) x = do
253 liftIO $ RB.unsafeInsertRing rb x
254 return $ FL.Partial $ Tuple' rb (cnt + 1)
256 done (Tuple' rb cnt) = do
257 let len = min cnt n
258 arr <- liftIO $ MArray.new len
259 arr1 <- MArray.uninit arr len
260 -- XXX Should use an API for retrieving the ring head
261 rhead <- liftIO $ readIORef $ RB.ringHead rb
262 if cnt <= n
263 then do
264 MArray.putSliceUnsafe (RB.ringArr rb) 0 arr1 0 cnt
265 else do
266 MArray.putSliceUnsafe (RB.ringArr rb) rhead arr1 0 (n - rhead)
267 MArray.putSliceUnsafe (RB.ringArr rb) 0 arr1 (n - rhead) rhead
268 return $ unsafeFreeze arr1
270sliceUnsafe :: Int -> Int -> Array a -> Array a
271sliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len
273-- XXX This is not efficient as it copies the array. We should support array
274-- slicing so that we can just refer to the underlying array memory instead of
275-- copying.
276--
277-- | Truncate the array at the beginning and end as long as the predicate
278-- holds true.
279strip :: (a -> Bool) -> Array a -> Array a
280strip p arr =
281 let lastIndex = length arr - 1
282 indexR = getIndexR lastIndex -- last predicate failing index
283 in if indexR == -1
284 then nil
285 else
286 let indexL = getIndexL 0 -- first predicate failing index
287 in if indexL == 0 && indexR == lastIndex
288 then arr
289 else unsafeFreeze $ unsafePerformIO $ do
290 let newLen = indexR - indexL + 1
291 arrThawed = unsafeThaw (sliceUnsafe indexL newLen arr)
292 MArray.clone arrThawed
294 where
296 getIndexR idx
297 | idx < 0 = idx
298 | otherwise =
299 if p (getIndexUnsafe arr idx) then getIndexR (idx - 1) else idx
301 getIndexL idx = if p (getIndexUnsafe arr idx)
302 then getIndexL (idx + 1)
303 else idx
305-------------------------------------------------------------------------------
306-- Instances
307-------------------------------------------------------------------------------
309instance Eq a => Eq (Array a) where
310 a1 == a2 = lenA1 == lenA2 && loop (lenA1 - 1)
312 where
314 lenA1 = length a1
315 lenA2 = length a2
317 loop i
318 | i < 0 = True
319 | otherwise =
320 let v1 = getIndexUnsafe a1 i
321 v2 = getIndexUnsafe a2 i
322 in v1 == v2 && loop (i - 1)
324instance Ord a => Ord (Array a) where
325 compare a1 a2 =
326 case compare lenA1 lenA2 of
327 LT -> LT
328 GT -> GT
329 EQ -> loop 0
331 where
333 lenA1 = length a1
334 lenA2 = length a2
336 loop i
337 | i >= lenA1 = EQ
338 | otherwise =
339 let v1 = getIndexUnsafe a1 i
340 v2 = getIndexUnsafe a2 i
341 in case compare v1 v2 of
342 LT -> LT
343 GT -> GT
344 EQ -> loop (i + 1)
346instance Show a => Show (Array a) where
347 show arr = "fromList " ++ show (toList arr)
349instance Read a => Read (Array a) where
350 readPrec = do
351 fromListWord <- replicateM 9 ReadPrec.get
352 if fromListWord == "fromList "
353 then fromList <$> readPrec
354 else ReadPrec.pfail