git.iko.soy/streamly
on ae2dce8 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
229 lines 6.7 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE MagicHash #-}
3{-# LANGUAGE UnboxedTuples #-}
5#include "inline.hs"
7-- |
8-- Module : Streamly.Internal.Data.Array
9-- Copyright : (c) 2019 Composewell Technologies
10--
11-- License : BSD-3-Clause
12-- Maintainer : streamly@composewell.com
13-- Stability : experimental
14-- Portability : GHC
15--
16module Streamly.Internal.Data.Array
17 ( Array(..)
19 , foldl'
20 , foldr
22 , length
24 , writeN
25 , write
27 , toStreamD
28 , toStreamDRev
30 , fromListN
31 , fromList
32 , fromStreamDN
33 , fromStreamD
34 )
35where
37import Prelude hiding (foldr, length)
38import Control.DeepSeq (NFData(..))
39import Control.Monad.ST (stToIO, RealWorld)
40import Control.Monad.IO.Class (liftIO, MonadIO)
41import GHC.ST (ST(..))
42import GHC.IO (unsafePerformIO)
43import GHC.Base (Int(..))
44import Data.Functor.Identity (runIdentity)
45import qualified GHC.Exts as Exts
47import Streamly.Internal.Data.Fold.Types (Fold(..))
49import qualified Streamly.Streams.StreamD as D
51data Array a =
52 Array
53 { array# :: Exts.Array# a
54 }
56{-# NOINLINE bottomElement #-}
57bottomElement :: a
58bottomElement = undefined
60data MutableArray s a =
61 MutableArray
62 { marray# :: Exts.MutableArray# s a
63 }
65{-# INLINE newArray #-}
66newArray :: Int -> a -> ST s (MutableArray s a)
67newArray (I# n#) x = ST
68 (\s# -> case Exts.newArray# n# x s# of
69 (# s'#, marr# #) ->
70 let ma = MutableArray marr#
71 in (# s'# , ma #))
73{-# INLINE writeArray #-}
74writeArray :: MutableArray s a -> Int -> a -> ST s ()
75writeArray marr (I# i#) x =
76 ST (\s# ->
77 case Exts.writeArray# (marray# marr) i# x s# of
78 s'# -> (# s'#, () #))
80{-# INLINE freezeArray #-}
81freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
82freezeArray marr (I# start#) (I# len#) =
83 ST (\s# ->
84 case Exts.freezeArray# (marray# marr) start# len# s# of
85 (# s'#, arr# #) -> (# s'#, Array arr# #))
87{-# INLINE copyMutableArray #-}
88copyMutableArray :: MutableArray s a
89 -- ^ The source array
90 -> Int
91 -- ^ Offset into the source array
92 -> MutableArray s a
93 -- ^ Destination array to copy to
94 -> Int
95 -- ^ Offset into the destination array
96 -> Int
97 -- ^ Number of elements to copy
98 -> ST s ()
99copyMutableArray sourceMarr (I# sourceOffset#) destMarr (I# destOffset#) (I# numElements#) =
100 ST (\s# ->
101 case Exts.copyMutableArray#
102 (marray# sourceMarr)
103 sourceOffset#
104 (marray# destMarr)
105 destOffset#
106 numElements#
107 s# of
108 s'# -> (# s'#, () #))
110{-# INLINE newArrayIO #-}
111newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
112newArrayIO i x = stToIO $ newArray i x
114{-# INLINE writeArrayIO #-}
115writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
116writeArrayIO marr i x = stToIO $ writeArray marr i x
118{-# INLINE freezeArrayIO #-}
119freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
120freezeArrayIO marr i len = stToIO $ freezeArray marr i len
122{-# INLINE length #-}
123length :: Array a -> Int
124length arr = I# (Exts.sizeofArray# (array# arr))
126{-# INLINE copyMutableArrayIO #-}
127copyMutableArrayIO ::
128 MutableArray RealWorld a
129 -> Int
130 -> MutableArray RealWorld a
131 -> Int
132 -> Int
133 -> IO ()
134copyMutableArrayIO sourceMarr sourceOffset destMarr destOffset numElements =
135 stToIO $
136 copyMutableArray sourceMarr sourceOffset destMarr destOffset numElements
138{-# INLINE_NORMAL toStreamD #-}
139toStreamD :: Monad m => Array a -> D.Stream m a
140toStreamD arr = D.Stream step 0
141 where
142 {-# INLINE_LATE step #-}
143 step _ i
144 | i == length arr = return D.Stop
145 step _ (I# i) =
146 return $
147 case Exts.indexArray# (array# arr) i of
148 (# x #) -> D.Yield x ((I# i) + 1)
150{-# INLINE_NORMAL toStreamDRev #-}
151toStreamDRev :: Monad m => Array a -> D.Stream m a
152toStreamDRev arr = D.Stream step (length arr - 1)
153 where
154 {-# INLINE_LATE step #-}
155 step _ i
156 | i < 0 = return D.Stop
157 step _ (I# i) =
158 return $
159 case Exts.indexArray# (array# arr) i of
160 (# x #) -> D.Yield x ((I# i) - 1)
162{-# INLINE_NORMAL foldl' #-}
163foldl' :: (b -> a -> b) -> b -> Array a -> b
164foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
166{-# INLINE_NORMAL foldr #-}
167foldr :: (a -> b -> b) -> b -> Array a -> b
168foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
170-- writeN n = S.evertM (fromStreamDN n)
171{-# INLINE_NORMAL writeN #-}
172writeN :: MonadIO m => Int -> Fold m a (Array a)
173writeN limit = Fold step initial extract
174 where
175 initial = do
176 marr <- liftIO $ newArrayIO limit bottomElement
177 return (marr, 0)
178 step (marr, i) x
179 | i == limit = return (marr, i)
180 | otherwise = do
181 liftIO $ writeArrayIO marr i x
182 return (marr, i + 1)
183 extract (marr, len) = liftIO $ freezeArrayIO marr 0 len
185{-# INLINE_NORMAL write #-}
186write :: MonadIO m => Fold m a (Array a)
187write = Fold step initial extract
188 where
189 initial = do
190 marr <- liftIO $ newArrayIO 0 bottomElement
191 return (marr, 0, 0)
192 step (marr, i, capacity) x
193 | i == capacity =
194 let newCapacity = max (capacity * 2) 1
195 in do newMarr <- liftIO $ newArrayIO newCapacity bottomElement
196 liftIO $ copyMutableArrayIO marr 0 newMarr 0 i
197 liftIO $ writeArrayIO newMarr i x
198 return (newMarr, i + 1, newCapacity)
199 | otherwise = do
200 liftIO $ writeArrayIO marr i x
201 return (marr, i + 1, capacity)
202 extract (marr, len, _) = liftIO $ freezeArrayIO marr 0 len
204{-# INLINE_NORMAL fromStreamDN #-}
205fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
206fromStreamDN limit str = do
207 marr <- liftIO $ newArrayIO (max limit 0) bottomElement
208 i <-
209 D.foldlM'
210 (\i x -> i `seq` (liftIO $ writeArrayIO marr i x) >> return (i + 1))
211 0 $
212 D.take limit str
213 liftIO $ freezeArrayIO marr 0 i
215{-# INLINE fromStreamD #-}
216fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
217fromStreamD str = D.runFold write str
219{-# INLINABLE fromListN #-}
220fromListN :: Int -> [a] -> Array a
221fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
223{-# INLINABLE fromList #-}
224fromList :: [a] -> Array a
225fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
227instance NFData a => NFData (Array a) where
228 {-# INLINE rnf #-}
229 rnf = foldl' (\_ x -> rnf x) ()