git.iko.soy/streamly
on d28e231 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
241 lines 7.3 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 , writeN
20 , write
22 , fromStreamDN
23 , fromStreamD
24 , fromStreamN
25 , fromStream
27 , fromListN
28 , fromList
30 -- * Elimination
31 , length
32 , read
34 , toStreamD
35 , toStreamDRev
36 , toStream
37 , toStreamRev
39 , foldl'
40 , foldr
41 , streamFold
42 , fold
44 -- * Random Access
45 , getIndexUnsafe
46 )
47where
49#if !MIN_VERSION_primitive(0,7,1)
50import Control.DeepSeq (NFData(..))
51#endif
52import Control.Monad (when)
53import Control.Monad.IO.Class (liftIO, MonadIO)
54import Data.Functor.Identity (runIdentity)
55import GHC.Base (Int(..))
56import GHC.IO (unsafePerformIO)
58import Streamly.Internal.Data.Fold.Type (Fold(..))
59import Streamly.Internal.Data.Stream.Serial (SerialT(..))
60import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
61import Streamly.Internal.Data.Unfold.Type (Unfold(..))
63import qualified GHC.Exts as Exts
64import qualified Streamly.Internal.Data.Fold.Type as FL
65import qualified Streamly.Internal.Data.Stream.StreamD as D
67import Data.Primitive.Array hiding (fromList, fromListN)
68import Prelude hiding (foldr, length, read)
70{-# NOINLINE bottomElement #-}
71bottomElement :: a
72bottomElement = undefined
74-------------------------------------------------------------------------------
75-- Construction - Folds
76-------------------------------------------------------------------------------
78-- writeN n = S.evertM (fromStreamDN n)
79{-# INLINE_NORMAL writeN #-}
80writeN :: MonadIO m => Int -> Fold m a (Array a)
81writeN limit = Fold step initial extract
82 where
83 initial = do
84 marr <- liftIO $ newArray limit bottomElement
85 return $ FL.Partial (Tuple' marr 0)
86 step st@(Tuple' marr i) x
87 | i == limit = fmap FL.Done $ extract st
88 | otherwise = do
89 liftIO $ writeArray marr i x
90 return $ FL.Partial $ Tuple' marr (i + 1)
91 extract (Tuple' marr len) = liftIO $ freezeArray marr 0 len
93{-# INLINE_NORMAL write #-}
94write :: MonadIO m => Fold m a (Array a)
95write = Fold step initial extract
96 where
97 initial = do
98 marr <- liftIO $ newArray 0 bottomElement
99 return $ FL.Partial (Tuple3' marr 0 0)
100 step (Tuple3' marr i capacity) x
101 | i == capacity =
102 let newCapacity = max (capacity * 2) 1
103 in do newMarr <- liftIO $ newArray newCapacity bottomElement
104 liftIO $ copyMutableArray newMarr 0 marr 0 i
105 liftIO $ writeArray newMarr i x
106 return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
107 | otherwise = do
108 liftIO $ writeArray marr i x
109 return $ FL.Partial $ Tuple3' marr (i + 1) capacity
110 extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
112-------------------------------------------------------------------------------
113-- Construction - from streams
114-------------------------------------------------------------------------------
116{-# INLINE_NORMAL fromStreamDN #-}
117fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
118fromStreamDN limit str = do
119 marr <- liftIO $ newArray (max limit 0) bottomElement
120 i <-
121 D.foldlM'
122 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
123 (return 0) $
124 D.take limit str
125 liftIO $ freezeArray marr 0 i
127{-# INLINE fromStreamD #-}
128fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
129fromStreamD = D.fold write
131{-# INLINE fromStreamN #-}
132fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
133fromStreamN n (SerialT m) = do
134 when (n < 0) $ error "fromStreamN: negative write count specified"
135 fromStreamDN n $ D.fromStreamK m
137{-# INLINE fromStream #-}
138fromStream :: MonadIO m => SerialT m a -> m (Array a)
139fromStream (SerialT m) = fromStreamD $ D.fromStreamK m
141{-# INLINABLE fromListN #-}
142fromListN :: Int -> [a] -> Array a
143fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
145{-# INLINABLE fromList #-}
146fromList :: [a] -> Array a
147fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
149-------------------------------------------------------------------------------
150-- Elimination - Unfolds
151-------------------------------------------------------------------------------
153{-# INLINE length #-}
154length :: Array a -> Int
155length = sizeofArray
157{-# INLINE_NORMAL read #-}
158read :: Monad m => Unfold m (Array a) a
159read = Unfold step inject
160 where
161 inject arr = return (arr, 0)
162 step (arr, i)
163 | i == length arr = return D.Stop
164 step (arr, I# i) =
165 return $
166 case Exts.indexArray# (array# arr) i of
167 (# x #) -> D.Yield x (arr, I# i + 1)
169-------------------------------------------------------------------------------
170-- Elimination - to streams
171-------------------------------------------------------------------------------
173{-# INLINE_NORMAL toStreamD #-}
174toStreamD :: Monad m => Array a -> D.Stream m a
175toStreamD arr = D.Stream step 0
176 where
177 {-# INLINE_LATE step #-}
178 step _ i
179 | i == length arr = return D.Stop
180 step _ (I# i) =
181 return $
182 case Exts.indexArray# (array# arr) i of
183 (# x #) -> D.Yield x (I# i + 1)
185{-# INLINE_NORMAL toStreamDRev #-}
186toStreamDRev :: Monad m => Array a -> D.Stream m a
187toStreamDRev arr = D.Stream step (length arr - 1)
188 where
189 {-# INLINE_LATE step #-}
190 step _ i
191 | i < 0 = return D.Stop
192 step _ (I# i) =
193 return $
194 case Exts.indexArray# (array# arr) i of
195 (# x #) -> D.Yield x (I# i - 1)
197{-# INLINE_EARLY toStream #-}
198toStream :: Monad m => Array a -> SerialT m a
199toStream = SerialT . D.toStreamK . toStreamD
201{-# INLINE_EARLY toStreamRev #-}
202toStreamRev :: Monad m => Array a -> SerialT m a
203toStreamRev = SerialT . D.toStreamK . toStreamDRev
205-------------------------------------------------------------------------------
206-- Elimination - using Folds
207-------------------------------------------------------------------------------
209{-# INLINE_NORMAL foldl' #-}
210foldl' :: (b -> a -> b) -> b -> Array a -> b
211foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
213{-# INLINE_NORMAL foldr #-}
214foldr :: (a -> b -> b) -> b -> Array a -> b
215foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
217#if !MIN_VERSION_primitive(0,7,1)
218instance NFData a => NFData (Array a) where
219 {-# INLINE rnf #-}
220 rnf = foldl' (\_ x -> rnf x) ()
221#endif
223{-# INLINE fold #-}
224fold :: Monad m => Fold m a b -> Array a -> m b
225fold f arr = D.fold f (toStreamD arr)
227{-# INLINE streamFold #-}
228streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
229streamFold f arr = f (toStream arr)
231-------------------------------------------------------------------------------
232-- Random reads and writes
233-------------------------------------------------------------------------------
235-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
236-- not check the bounds.
237--
238-- @since 0.8.0
239{-# INLINE getIndexUnsafe #-}
240getIndexUnsafe :: Array a -> Int -> a
241getIndexUnsafe = indexArray