git.iko.soy/streamly
on a05abfc help sign in
streamly/src/Streamly/Internal/Data/Array.hs
315 lines 9.4 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 )
50where
52#if !MIN_VERSION_primitive(0,7,1)
53import Control.DeepSeq (NFData(..))
54#endif
55import Control.Monad (when)
56import Control.Monad.IO.Class (liftIO, MonadIO)
57import Data.Functor.Identity (runIdentity)
58import Data.IORef
59import GHC.Base (Int(..))
60import GHC.IO (unsafePerformIO)
62import Streamly.Internal.Data.Fold.Type (Fold(..))
63import Streamly.Internal.Data.Stream.Serial (SerialT(..))
64import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
65import Streamly.Internal.Data.Unfold.Type (Unfold(..))
67import qualified GHC.Exts as Exts
68import qualified Streamly.Internal.Data.Fold.Type as FL
69import qualified Streamly.Internal.Data.Ring as RB
70import qualified Streamly.Internal.Data.Stream.StreamD as D
72import Data.Primitive.Array hiding (fromList, fromListN)
73import Prelude hiding (foldr, length, read)
75{-# NOINLINE bottomElement #-}
76bottomElement :: a
77bottomElement = undefined
79{-# NOINLINE nil #-}
80nil :: Array a
81nil = unsafePerformIO $ do
82 marr <- liftIO $ newArray 0 bottomElement
83 liftIO $ freezeArray marr 0 0
85-------------------------------------------------------------------------------
86-- Construction - Folds
87-------------------------------------------------------------------------------
89{-# INLINE_NORMAL writeN #-}
90writeN :: MonadIO m => Int -> Fold m a (Array a)
91writeN len = Fold step initial extract
93 where
95 {-# INLINE next #-}
96 next marr i = do
97 let i1 = i + 1
98 st = Tuple' marr i1
99 if len > i1
100 then return $ FL.Partial st
101 else fmap FL.Done $ extract st
103 initial = do
104 marr <- liftIO $ newArray len bottomElement
105 next marr (-1)
107 step (Tuple' marr i) x = do
108 liftIO $ writeArray marr i x
109 next marr i
111 extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
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 $ newArray 0 bottomElement
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 $ newArray newCapacity bottomElement
124 liftIO $ copyMutableArray newMarr 0 marr 0 i
125 liftIO $ writeArray newMarr i x
126 return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
127 | otherwise = do
128 liftIO $ writeArray marr i x
129 return $ FL.Partial $ Tuple3' marr (i + 1) capacity
130 extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
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 $ newArray (max limit 0) bottomElement
140 i <-
141 D.foldlM'
142 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
143 (return 0) $
144 D.take limit str
145 liftIO $ freezeArray marr 0 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 -> SerialT m a -> m (Array a)
153fromStreamN n (SerialT m) = do
154 when (n < 0) $ error "fromStreamN: negative write count specified"
155 fromStreamDN n $ D.fromStreamK m
157{-# INLINE fromStream #-}
158fromStream :: MonadIO m => SerialT m a -> m (Array a)
159fromStream (SerialT m) = fromStreamD $ D.fromStreamK m
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 = sizeofArray
177{-# INLINE_NORMAL read #-}
178read :: Monad m => Unfold m (Array a) a
179read = Unfold step inject
180 where
181 inject arr = return (arr, 0)
182 step (arr, i)
183 | i == length arr = return D.Stop
184 step (arr, I# i) =
185 return $
186 case Exts.indexArray# (array# arr) i of
187 (# x #) -> D.Yield x (arr, I# i + 1)
189-------------------------------------------------------------------------------
190-- Elimination - to streams
191-------------------------------------------------------------------------------
193{-# INLINE_NORMAL toStreamD #-}
194toStreamD :: Monad m => Array a -> D.Stream m a
195toStreamD arr = D.Stream step 0
196 where
197 {-# INLINE_LATE step #-}
198 step _ i
199 | i == length arr = return D.Stop
200 step _ (I# i) =
201 return $
202 case Exts.indexArray# (array# arr) i of
203 (# x #) -> D.Yield x (I# i + 1)
205{-# INLINE_NORMAL toStreamDRev #-}
206toStreamDRev :: Monad m => Array a -> D.Stream m a
207toStreamDRev arr = D.Stream step (length arr - 1)
208 where
209 {-# INLINE_LATE step #-}
210 step _ i
211 | i < 0 = return D.Stop
212 step _ (I# i) =
213 return $
214 case Exts.indexArray# (array# arr) i of
215 (# x #) -> D.Yield x (I# i - 1)
217{-# INLINE_EARLY toStream #-}
218toStream :: Monad m => Array a -> SerialT m a
219toStream = SerialT . D.toStreamK . toStreamD
221{-# INLINE_EARLY toStreamRev #-}
222toStreamRev :: Monad m => Array a -> SerialT m a
223toStreamRev = SerialT . D.toStreamK . toStreamDRev
225-------------------------------------------------------------------------------
226-- Elimination - using Folds
227-------------------------------------------------------------------------------
229{-# INLINE_NORMAL foldl' #-}
230foldl' :: (b -> a -> b) -> b -> Array a -> b
231foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
233{-# INLINE_NORMAL foldr #-}
234foldr :: (a -> b -> b) -> b -> Array a -> b
235foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
237#if !MIN_VERSION_primitive(0,7,1)
238instance NFData a => NFData (Array a) where
239 {-# INLINE rnf #-}
240 rnf = foldl' (\_ x -> rnf x) ()
241#endif
243{-# INLINE fold #-}
244fold :: Monad m => Fold m a b -> Array a -> m b
245fold f arr = D.fold f (toStreamD arr)
247{-# INLINE streamFold #-}
248streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
249streamFold f arr = f (toStream arr)
251-------------------------------------------------------------------------------
252-- Random reads and writes
253-------------------------------------------------------------------------------
255-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
256-- not check the bounds.
257--
258-- @since 0.8.0
259{-# INLINE getIndexUnsafe #-}
260getIndexUnsafe :: Array a -> Int -> a
261getIndexUnsafe = indexArray
263{-# INLINE writeLastN #-}
264writeLastN :: MonadIO m => Int -> Fold m a (Array a)
265writeLastN n
266 | n <= 0 = fmap (const mempty) FL.drain
267 | otherwise = Fold step initial done
269 where
271 initial = do
272 rb <- liftIO $ RB.createRing n
273 return $ FL.Partial $ Tuple' rb (0 :: Int)
275 step (Tuple' rb rh) x = do
276 liftIO $ RB.unsafeInsertRing rb rh x
277 return $ FL.Partial $ Tuple' rb (rh + 1)
279 done (Tuple' rb rh) = do
280 arr' <- liftIO $ newArray (min rh n) (undefined :: a)
281 ref <- liftIO $ readIORef $ RB.ringHead rb
282 if rh < n
283 then
284 liftIO $ copyMutableArray arr' 0 (RB.arr rb) 0 ref
285 else do
286 liftIO $ copyMutableArray arr' 0 (RB.arr rb) ref (n - ref)
287 liftIO $ copyMutableArray arr' (n - ref) (RB.arr rb) 0 ref
288 liftIO $ unsafeFreezeArray arr'
290-- XXX This is not efficient as it copies the array. We should support array
291-- slicing so that we can just refer to the underlying array memory instead of
292-- copying.
293--
294-- | Truncate the array at the beginning and end as long as the predicate
295-- holds true.
296strip :: (a -> Bool) -> Array a -> Array a
297strip p arr =
298 let lastIndex = length arr - 1
299 indexR = getIndexR lastIndex -- last predicate failing index
300 in if indexR == -1
301 then nil
302 else
303 let indexL = getIndexL 0 -- first predicate failing index
304 in if indexL == 0 && indexR == lastIndex
305 then arr
306 else cloneArray arr indexL (indexR - indexL + 1)
308 where
310 getIndexR idx
311 | idx < 0 = idx
312 | otherwise =
313 if p (indexArray arr idx) then getIndexR (idx - 1) else idx
315 getIndexL idx = if p (indexArray arr idx) then getIndexL (idx + 1) else idx