git.iko.soy/streamly
on 2aae2da help sign in
streamly/core/src/Streamly/Internal/Data/Array.hs
334 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.Serial (SerialT(..))
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
74import Data.Primitive.Array hiding (fromList, fromListN)
75import Prelude hiding (foldr, length, read)
77{-# NOINLINE bottomElement #-}
78bottomElement :: a
79bottomElement = undefined
81{-# NOINLINE nil #-}
82nil :: Array a
83nil = unsafePerformIO $ do
84 marr <- liftIO $ newArray 0 bottomElement
85 liftIO $ freezeArray marr 0 0
87-------------------------------------------------------------------------------
88-- Construction - Folds
89-------------------------------------------------------------------------------
91{-# INLINE_NORMAL writeN #-}
92writeN :: MonadIO m => Int -> Fold m a (Array a)
93writeN len = Fold step initial extract
95 where
97 {-# INLINE next #-}
98 next marr i = do
99 let i1 = i + 1
100 st = Tuple' marr i1
101 if len > i1
102 then return $ FL.Partial st
103 else fmap FL.Done $ extract st
105 initial = do
106 marr <- liftIO $ newArray len bottomElement
107 next marr (-1)
109 step (Tuple' marr i) x = do
110 liftIO $ writeArray marr i x
111 next marr i
113 extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
115{-# INLINE_NORMAL write #-}
116write :: MonadIO m => Fold m a (Array a)
117write = Fold step initial extract
118 where
119 initial = do
120 marr <- liftIO $ newArray 0 bottomElement
121 return $ FL.Partial (Tuple3' marr 0 0)
122 step (Tuple3' marr i capacity) x
123 | i == capacity =
124 let newCapacity = max (capacity * 2) 1
125 in do newMarr <- liftIO $ newArray newCapacity bottomElement
126 liftIO $ copyMutableArray newMarr 0 marr 0 i
127 liftIO $ writeArray newMarr i x
128 return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
129 | otherwise = do
130 liftIO $ writeArray marr i x
131 return $ FL.Partial $ Tuple3' marr (i + 1) capacity
132 extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
134-------------------------------------------------------------------------------
135-- Construction - from streams
136-------------------------------------------------------------------------------
138{-# INLINE_NORMAL fromStreamDN #-}
139fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
140fromStreamDN limit str = do
141 marr <- liftIO $ newArray (max limit 0) bottomElement
142 i <-
143 D.foldlM'
144 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
145 (return 0) $
146 D.take limit str
147 liftIO $ freezeArray marr 0 i
149{-# INLINE fromStreamD #-}
150fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
151fromStreamD = D.fold write
153{-# INLINE fromStreamN #-}
154fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
155fromStreamN n (SerialT m) = do
156 when (n < 0) $ error "fromStreamN: negative write count specified"
157 fromStreamDN n $ D.fromStreamK m
159{-# INLINE fromStream #-}
160fromStream :: MonadIO m => SerialT m a -> m (Array a)
161fromStream (SerialT m) = fromStreamD $ D.fromStreamK m
163{-# INLINABLE fromListN #-}
164fromListN :: Int -> [a] -> Array a
165fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
167{-# INLINABLE fromList #-}
168fromList :: [a] -> Array a
169fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
171-------------------------------------------------------------------------------
172-- Elimination - Unfolds
173-------------------------------------------------------------------------------
175{-# INLINE length #-}
176length :: Array a -> Int
177length = sizeofArray
179{-# INLINE_NORMAL read #-}
180read :: Monad m => Unfold m (Array a) a
181read = Unfold step inject
182 where
183 inject arr = return (arr, 0)
184 step (arr, i)
185 | i == length arr = return D.Stop
186 step (arr, I# i) =
187 return $
188 case Exts.indexArray# (array# arr) i of
189 (# x #) -> D.Yield x (arr, I# i + 1)
191-------------------------------------------------------------------------------
192-- Elimination - to streams
193-------------------------------------------------------------------------------
195{-# INLINE_NORMAL toStreamD #-}
196toStreamD :: Monad m => Array a -> D.Stream m a
197toStreamD arr = D.Stream step 0
198 where
199 {-# INLINE_LATE step #-}
200 step _ i
201 | i == length arr = return D.Stop
202 step _ (I# i) =
203 return $
204 case Exts.indexArray# (array# arr) i of
205 (# x #) -> D.Yield x (I# i + 1)
207{-# INLINE_NORMAL toStreamDRev #-}
208toStreamDRev :: Monad m => Array a -> D.Stream m a
209toStreamDRev arr = D.Stream step (length arr - 1)
210 where
211 {-# INLINE_LATE step #-}
212 step _ i
213 | i < 0 = return D.Stop
214 step _ (I# i) =
215 return $
216 case Exts.indexArray# (array# arr) i of
217 (# x #) -> D.Yield x (I# i - 1)
219{-# INLINE_EARLY toStream #-}
220toStream :: Monad m => Array a -> SerialT m a
221toStream = SerialT . D.toStreamK . toStreamD
223{-# INLINE_EARLY toStreamRev #-}
224toStreamRev :: Monad m => Array a -> SerialT m a
225toStreamRev = SerialT . D.toStreamK . toStreamDRev
227-------------------------------------------------------------------------------
228-- Elimination - using Folds
229-------------------------------------------------------------------------------
231{-# INLINE_NORMAL foldl' #-}
232foldl' :: (b -> a -> b) -> b -> Array a -> b
233foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
235{-# INLINE_NORMAL foldr #-}
236foldr :: (a -> b -> b) -> b -> Array a -> b
237foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
239#if !MIN_VERSION_primitive(0,7,1)
240instance NFData a => NFData (Array a) where
241 {-# INLINE rnf #-}
242 rnf = foldl' (\_ x -> rnf x) ()
243#endif
245{-# INLINE fold #-}
246fold :: Monad m => Fold m a b -> Array a -> m b
247fold f arr = D.fold f (toStreamD arr)
249{-# INLINE streamFold #-}
250streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
251streamFold f arr = f (toStream arr)
253-------------------------------------------------------------------------------
254-- Random reads and writes
255-------------------------------------------------------------------------------
257-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
258-- not check the bounds.
259--
260-- @since 0.8.0
261{-# INLINE getIndexUnsafe #-}
262getIndexUnsafe :: Array a -> Int -> a
263getIndexUnsafe = indexArray
265{-# INLINE writeLastN #-}
266writeLastN :: MonadIO m => Int -> Fold m a (Array a)
267writeLastN n
268 | n <= 0 = fmap (const mempty) FL.drain
269 | otherwise = Fold step initial done
271 where
273 initial = do
274 rb <- liftIO $ RB.createRing n
275 return $ FL.Partial $ Tuple' rb (0 :: Int)
277 step (Tuple' rb rh) x = do
278 liftIO $ RB.unsafeInsertRing rb rh x
279 return $ FL.Partial $ Tuple' rb (rh + 1)
281 done (Tuple' rb rh) = do
282 arr' <- liftIO $ newArray (min rh n) (undefined :: a)
283 ref <- liftIO $ readIORef $ RB.ringHead rb
284 if rh < n
285 then
286 liftIO $ copyMutableArray arr' 0 (RB.arr rb) 0 ref
287 else do
288 liftIO $ copyMutableArray arr' 0 (RB.arr rb) ref (n - ref)
289 liftIO $ copyMutableArray arr' (n - ref) (RB.arr rb) 0 ref
290 liftIO $ unsafeFreezeArray arr'
292-- XXX This is not efficient as it copies the array. We should support array
293-- slicing so that we can just refer to the underlying array memory instead of
294-- copying.
295--
296-- | Truncate the array at the beginning and end as long as the predicate
297-- holds true.
298strip :: (a -> Bool) -> Array a -> Array a
299strip p arr =
300 let lastIndex = length arr - 1
301 indexR = getIndexR lastIndex -- last predicate failing index
302 in if indexR == -1
303 then nil
304 else
305 let indexL = getIndexL 0 -- first predicate failing index
306 in if indexL == 0 && indexR == lastIndex
307 then arr
308 else cloneArray arr indexL (indexR - indexL + 1)
310 where
312 getIndexR idx
313 | idx < 0 = idx
314 | otherwise =
315 if p (indexArray arr idx) then getIndexR (idx - 1) else idx
317 getIndexL idx = if p (indexArray arr idx) then getIndexL (idx + 1) else idx
319-- | Write an input stream of (index, value) pairs to an array. Throws an
320-- error if any index is out of bounds.
321--
322-- /Pre-release/
323{-# INLINE putIndices #-}
324putIndices :: PrimMonad m
325 => Array a -> Fold m (Int, a) ()
326putIndices arr = FL.rmapM (\ _ -> return ()) (FL.foldlM' step initial)
328 where
330 initial = unsafeThawArray arr
332 step marr (i, x) = do
333 writeArray marr i x
334 return marr