git.iko.soy/streamly
on d2aaef0 help sign in
streamly/core/src/Streamly/Internal/Data/Array.hs
335 lines 9.9 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE UnboxedTuples #-}
4#include "inline.hs"
6-- |
7-- Module : Streamly.Internal.Data.Array
8-- Copyright : (c) 2019 Composewell Technologies
9--
10-- License : BSD-3-Clause
11-- Maintainer : streamly@composewell.com
12-- Stability : pre-release
13-- Portability : GHC
14--
15module Streamly.Internal.Data.Array
16 ( Array(..)
18 -- * Construction
19 , nil
20 , writeN
21 , write
22 , writeLastN
24 , fromStreamDN
25 , fromStreamD
26 , fromStreamN
27 , fromStream
29 , fromListN
30 , fromList
32 -- * Elimination
33 , length
34 , read
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
53#if !MIN_VERSION_primitive(0,7,1)
54import Control.DeepSeq (NFData(..))
55#endif
56import Control.Monad (when)
57import Control.Monad.IO.Class (liftIO, MonadIO)
58import Control.Monad.Primitive (PrimMonad)
59import Data.Functor.Identity (runIdentity)
60import Data.IORef
61import GHC.Base (Int(..))
62import GHC.IO (unsafePerformIO)
64import Streamly.Internal.Data.Fold.Type (Fold(..))
65import Streamly.Internal.Data.Stream.Type (Stream)
66import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
67import Streamly.Internal.Data.Unfold.Type (Unfold(..))
69import qualified GHC.Exts as Exts
70import qualified Streamly.Internal.Data.Fold.Type as FL
71import qualified Streamly.Internal.Data.Ring as RB
72import qualified Streamly.Internal.Data.Stream.StreamD as D
73import qualified Streamly.Internal.Data.Stream.Type as Stream
75import Data.Primitive.Array hiding (fromList, fromListN)
76import Prelude hiding (foldr, length, read)
78{-# NOINLINE bottomElement #-}
79bottomElement :: a
80bottomElement = undefined
82{-# NOINLINE nil #-}
83nil :: Array a
84nil = unsafePerformIO $ do
85 marr <- liftIO $ newArray 0 bottomElement
86 liftIO $ freezeArray marr 0 0
88-------------------------------------------------------------------------------
89-- Construction - Folds
90-------------------------------------------------------------------------------
92{-# INLINE_NORMAL writeN #-}
93writeN :: MonadIO m => Int -> Fold m a (Array a)
94writeN len = Fold step initial extract
96 where
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
106 initial = do
107 marr <- liftIO $ newArray len bottomElement
108 next marr (-1)
110 step (Tuple' marr i) x = do
111 liftIO $ writeArray marr i x
112 next marr i
114 extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
116{-# INLINE_NORMAL write #-}
117write :: MonadIO m => Fold m a (Array a)
118write = Fold step initial extract
119 where
120 initial = do
121 marr <- liftIO $ newArray 0 bottomElement
122 return $ FL.Partial (Tuple3' marr 0 0)
123 step (Tuple3' marr i capacity) x
124 | i == capacity =
125 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
130 | 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
135-------------------------------------------------------------------------------
136-- Construction - from streams
137-------------------------------------------------------------------------------
139{-# INLINE_NORMAL fromStreamDN #-}
140fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
141fromStreamDN limit str = do
142 marr <- liftIO $ newArray (max limit 0) bottomElement
143 i <-
144 D.foldlM'
145 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
146 (return 0) $
147 D.take limit str
148 liftIO $ freezeArray marr 0 i
150{-# INLINE fromStreamD #-}
151fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
152fromStreamD = D.fold write
154{-# INLINE fromStreamN #-}
155fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)
156fromStreamN n m = do
157 when (n < 0) $ error "fromStreamN: negative write count specified"
158 fromStreamDN n $ Stream.toStreamD m
160{-# INLINE fromStream #-}
161fromStream :: MonadIO m => Stream m a -> m (Array a)
162fromStream = fromStreamD . Stream.toStreamD
164{-# INLINABLE fromListN #-}
165fromListN :: Int -> [a] -> Array a
166fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
168{-# INLINABLE fromList #-}
169fromList :: [a] -> Array a
170fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
172-------------------------------------------------------------------------------
173-- Elimination - Unfolds
174-------------------------------------------------------------------------------
176{-# INLINE length #-}
177length :: Array a -> Int
178length = sizeofArray
180{-# INLINE_NORMAL read #-}
181read :: Monad m => Unfold m (Array a) a
182read = 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)
192-------------------------------------------------------------------------------
193-- Elimination - to streams
194-------------------------------------------------------------------------------
196{-# INLINE_NORMAL toStreamD #-}
197toStreamD :: Monad m => Array a -> D.Stream m a
198toStreamD 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)
208{-# INLINE_NORMAL toStreamDRev #-}
209toStreamDRev :: Monad m => Array a -> D.Stream m a
210toStreamDRev 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)
220{-# INLINE_EARLY toStream #-}
221toStream :: Monad m => Array a -> Stream m a
222toStream = Stream.fromStreamD . toStreamD
224{-# INLINE_EARLY toStreamRev #-}
225toStreamRev :: Monad m => Array a -> Stream m a
226toStreamRev = Stream.fromStreamD . toStreamDRev
228-------------------------------------------------------------------------------
229-- Elimination - using Folds
230-------------------------------------------------------------------------------
232{-# INLINE_NORMAL foldl' #-}
233foldl' :: (b -> a -> b) -> b -> Array a -> b
234foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
236{-# INLINE_NORMAL foldr #-}
237foldr :: (a -> b -> b) -> b -> Array a -> b
238foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
240#if !MIN_VERSION_primitive(0,7,1)
241instance NFData a => NFData (Array a) where
242 {-# INLINE rnf #-}
243 rnf = foldl' (\_ x -> rnf x) ()
244#endif
246{-# INLINE fold #-}
247fold :: Monad m => Fold m a b -> Array a -> m b
248fold f arr = D.fold f (toStreamD arr)
250{-# INLINE streamFold #-}
251streamFold :: Monad m => (Stream m a -> m b) -> Array a -> m b
252streamFold f arr = f (toStream arr)
254-------------------------------------------------------------------------------
255-- Random reads and writes
256-------------------------------------------------------------------------------
258-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
259-- not check the bounds.
260--
261-- @since 0.8.0
262{-# INLINE getIndexUnsafe #-}
263getIndexUnsafe :: Array a -> Int -> a
264getIndexUnsafe = indexArray
266{-# INLINE writeLastN #-}
267writeLastN :: MonadIO m => Int -> Fold m a (Array a)
268writeLastN n
269 | n <= 0 = fmap (const mempty) FL.drain
270 | otherwise = Fold step initial done
272 where
274 initial = do
275 rb <- liftIO $ RB.createRing n
276 return $ FL.Partial $ Tuple' rb (0 :: Int)
278 step (Tuple' rb rh) x = do
279 liftIO $ RB.unsafeInsertRing rb rh x
280 return $ FL.Partial $ Tuple' rb (rh + 1)
282 done (Tuple' rb rh) = do
283 arr' <- liftIO $ newArray (min rh n) (undefined :: a)
284 ref <- liftIO $ readIORef $ RB.ringHead rb
285 if rh < n
286 then
287 liftIO $ copyMutableArray arr' 0 (RB.arr rb) 0 ref
288 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'
293-- XXX This is not efficient as it copies the array. We should support array
294-- slicing so that we can just refer to the underlying array memory instead of
295-- copying.
296--
297-- | Truncate the array at the beginning and end as long as the predicate
298-- holds true.
299strip :: (a -> Bool) -> Array a -> Array a
300strip p arr =
301 let lastIndex = length arr - 1
302 indexR = getIndexR lastIndex -- last predicate failing index
303 in if indexR == -1
304 then nil
305 else
306 let indexL = getIndexL 0 -- first predicate failing index
307 in if indexL == 0 && indexR == lastIndex
308 then arr
309 else cloneArray arr indexL (indexR - indexL + 1)
311 where
313 getIndexR idx
314 | idx < 0 = idx
315 | otherwise =
316 if p (indexArray arr idx) then getIndexR (idx - 1) else idx
318 getIndexL idx = if p (indexArray arr idx) then getIndexL (idx + 1) else idx
320-- | Write an input stream of (index, value) pairs to an array. Throws an
321-- error if any index is out of bounds.
322--
323-- /Pre-release/
324{-# INLINE putIndices #-}
325putIndices :: PrimMonad m
326 => Array a -> Fold m (Int, a) ()
327putIndices arr = FL.rmapM (\ _ -> return ()) (FL.foldlM' step initial)
329 where
331 initial = unsafeThawArray arr
333 step marr (i, x) = do
334 writeArray marr i x
335 return marr