on master help sign in
Use Array from Data.Primitive.Array rather than defining own Array.
0d91f2ePranay Sashank
src/Streamly/Internal/Data/Array.hs+19−97
@@ -1,3 +1,5 @@
1+{-# OPTIONS_GHC -fno-warn-orphans #-}
2+
1{-# LANGUAGE CPP #-}
3{-# LANGUAGE CPP #-}
2{-# LANGUAGE MagicHash #-}
4{-# LANGUAGE MagicHash #-}
3{-# LANGUAGE UnboxedTuples #-}
5{-# LANGUAGE UnboxedTuples #-}
@@ -36,105 +38,21 @@ where
36
38
37import Prelude hiding (foldr, length)
39import Prelude hiding (foldr, length)
38import Control.DeepSeq (NFData(..))
40import Control.DeepSeq (NFData(..))
39−import Control.Monad.ST (stToIO, RealWorld)
40import Control.Monad.IO.Class (liftIO, MonadIO)
41import Control.Monad.IO.Class (liftIO, MonadIO)
41−import GHC.ST (ST(..))
42import GHC.IO (unsafePerformIO)
42import GHC.IO (unsafePerformIO)
43import GHC.Base (Int(..))
43import GHC.Base (Int(..))
44import Data.Functor.Identity (runIdentity)
44import Data.Functor.Identity (runIdentity)
45+import Data.Primitive.Array hiding (fromList, fromListN)
45import qualified GHC.Exts as Exts
46import qualified GHC.Exts as Exts
46
47
47import Streamly.Internal.Data.Fold.Types (Fold(..))
48import Streamly.Internal.Data.Fold.Types (Fold(..))
48
49
49import qualified Streamly.Streams.StreamD as D
50import qualified Streamly.Streams.StreamD as D
50
51
51−data Array a =
52− Array
53− { array# :: Exts.Array# a
54− }
55−
56{-# NOINLINE bottomElement #-}
52{-# NOINLINE bottomElement #-}
57bottomElement :: a
53bottomElement :: a
58bottomElement = undefined
54bottomElement = undefined
59
55
60−data MutableArray s a =
61− MutableArray
62− { marray# :: Exts.MutableArray# s a
63− }
64−
65−{-# INLINE newArray #-}
66−newArray :: Int -> a -> ST s (MutableArray s a)
67−newArray (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 #))
72−
73−{-# INLINE writeArray #-}
74−writeArray :: MutableArray s a -> Int -> a -> ST s ()
75−writeArray marr (I# i#) x =
76− ST (\s# ->
77− case Exts.writeArray# (marray# marr) i# x s# of
78− s'# -> (# s'#, () #))
79−
80−{-# INLINE freezeArray #-}
81−freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
82−freezeArray marr (I# start#) (I# len#) =
83− ST (\s# ->
84− case Exts.freezeArray# (marray# marr) start# len# s# of
85− (# s'#, arr# #) -> (# s'#, Array arr# #))
86−
87−{-# INLINE copyMutableArray #-}
88−copyMutableArray :: 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 ()
99−copyMutableArray 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'#, () #))
109−
110−{-# INLINE newArrayIO #-}
111−newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
112−newArrayIO i x = stToIO $ newArray i x
113−
114−{-# INLINE writeArrayIO #-}
115−writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
116−writeArrayIO marr i x = stToIO $ writeArray marr i x
117−
118−{-# INLINE freezeArrayIO #-}
119−freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
120−freezeArrayIO marr i len = stToIO $ freezeArray marr i len
121−
122−{-# INLINE length #-}
123−length :: Array a -> Int
124−length arr = I# (Exts.sizeofArray# (array# arr))
125−
126−{-# INLINE copyMutableArrayIO #-}
127−copyMutableArrayIO ::
128− MutableArray RealWorld a
129− -> Int
130− -> MutableArray RealWorld a
131− -> Int
132− -> Int
133− -> IO ()
134−copyMutableArrayIO sourceMarr sourceOffset destMarr destOffset numElements =
135− stToIO $
136− copyMutableArray sourceMarr sourceOffset destMarr destOffset numElements
137−
138{-# INLINE_NORMAL toStreamD #-}
56{-# INLINE_NORMAL toStreamD #-}
139toStreamD :: Monad m => Array a -> D.Stream m a
57toStreamD :: Monad m => Array a -> D.Stream m a
140toStreamD arr = D.Stream step 0
58toStreamD arr = D.Stream step 0
@@ -147,6 +65,10 @@ toStreamD arr = D.Stream step 0
147 case Exts.indexArray# (array# arr) i of
65 case Exts.indexArray# (array# arr) i of
148 (# x #) -> D.Yield x ((I# i) + 1)
66 (# x #) -> D.Yield x ((I# i) + 1)
149
67
68+{-# INLINE length #-}
69+length :: Array a -> Int
70+length arr = sizeofArray arr
71+
150{-# INLINE_NORMAL toStreamDRev #-}
72{-# INLINE_NORMAL toStreamDRev #-}
151toStreamDRev :: Monad m => Array a -> D.Stream m a
73toStreamDRev :: Monad m => Array a -> D.Stream m a
152toStreamDRev arr = D.Stream step (length arr - 1)
74toStreamDRev arr = D.Stream step (length arr - 1)
@@ -173,44 +95,44 @@ writeN :: MonadIO m => Int -> Fold m a (Array a)
173writeN limit = Fold step initial extract
95writeN limit = Fold step initial extract
174 where
96 where
175 initial = do
97 initial = do
176− marr <- liftIO $ newArrayIO limit bottomElement
98+ marr <- liftIO $ newArray limit bottomElement
177 return (marr, 0)
99 return (marr, 0)
178 step (marr, i) x
100 step (marr, i) x
179 | i == limit = return (marr, i)
101 | i == limit = return (marr, i)
180 | otherwise = do
102 | otherwise = do
181− liftIO $ writeArrayIO marr i x
103+ liftIO $ writeArray marr i x
182 return (marr, i + 1)
104 return (marr, i + 1)
183− extract (marr, len) = liftIO $ freezeArrayIO marr 0 len
105+ extract (marr, len) = liftIO $ freezeArray marr 0 len
184
106
185{-# INLINE_NORMAL write #-}
107{-# INLINE_NORMAL write #-}
186write :: MonadIO m => Fold m a (Array a)
108write :: MonadIO m => Fold m a (Array a)
187write = Fold step initial extract
109write = Fold step initial extract
188 where
110 where
189 initial = do
111 initial = do
190− marr <- liftIO $ newArrayIO 0 bottomElement
112+ marr <- liftIO $ newArray 0 bottomElement
191 return (marr, 0, 0)
113 return (marr, 0, 0)
192 step (marr, i, capacity) x
114 step (marr, i, capacity) x
193 | i == capacity =
115 | i == capacity =
194 let newCapacity = max (capacity * 2) 1
116 let newCapacity = max (capacity * 2) 1
195− in do newMarr <- liftIO $ newArrayIO newCapacity bottomElement
117+ in do newMarr <- liftIO $ newArray newCapacity bottomElement
196− liftIO $ copyMutableArrayIO marr 0 newMarr 0 i
118+ liftIO $ copyMutableArray newMarr 0 marr 0 i
197− liftIO $ writeArrayIO newMarr i x
119+ liftIO $ writeArray newMarr i x
198 return (newMarr, i + 1, newCapacity)
120 return (newMarr, i + 1, newCapacity)
199 | otherwise = do
121 | otherwise = do
200− liftIO $ writeArrayIO marr i x
122+ liftIO $ writeArray marr i x
201 return (marr, i + 1, capacity)
123 return (marr, i + 1, capacity)
202− extract (marr, len, _) = liftIO $ freezeArrayIO marr 0 len
124+ extract (marr, len, _) = liftIO $ freezeArray marr 0 len
203
125
204{-# INLINE_NORMAL fromStreamDN #-}
126{-# INLINE_NORMAL fromStreamDN #-}
205fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
127fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
206fromStreamDN limit str = do
128fromStreamDN limit str = do
207− marr <- liftIO $ newArrayIO (max limit 0) bottomElement
129+ marr <- liftIO $ newArray (max limit 0) bottomElement
208 i <-
130 i <-
209 D.foldlM'
131 D.foldlM'
210− (\i x -> i `seq` (liftIO $ writeArrayIO marr i x) >> return (i + 1))
132+ (\i x -> i `seq` (liftIO $ writeArray marr i x) >> return (i + 1))
211 0 $
133 0 $
212 D.take limit str
134 D.take limit str
213− liftIO $ freezeArrayIO marr 0 i
135+ liftIO $ freezeArray marr 0 i
214
136
215{-# INLINE fromStreamD #-}
137{-# INLINE fromStreamD #-}
216fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
138fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
streamly.cabal+1
@@ -361,6 +361,7 @@ library
361 , containers >= 0.5.8.2 && < 0.7
361 , containers >= 0.5.8.2 && < 0.7
362 , heaps >= 0.3 && < 0.4
362 , heaps >= 0.3 && < 0.4
363 , directory >= 1.3 && < 1.4
363 , directory >= 1.3 && < 1.4
364+ , primitive >= 0.6.4.0 && < 0.8
364
365
365 -- concurrency
366 -- concurrency
366 , atomic-primops >= 0.8 && < 0.9
367 , atomic-primops >= 0.8 && < 0.9