@@ -1,3 +1,5 @@
1+{-# OPTIONS_GHC -fno-warn-orphans #-}
1{-# LANGUAGE CPP #-}
3{-# LANGUAGE CPP #-}
2{-# LANGUAGE MagicHash #-}
4{-# LANGUAGE MagicHash #-}
3{-# LANGUAGE UnboxedTuples #-}
5{-# LANGUAGE UnboxedTuples #-}
@@ -36,105 +38,21 @@ where
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
47import Streamly.Internal.Data.Fold.Types (Fold(..))
48import Streamly.Internal.Data.Fold.Types (Fold(..))
49import qualified Streamly.Streams.StreamD as D
50import qualified Streamly.Streams.StreamD as D
53− { array# :: Exts.Array# a
56{-# NOINLINE bottomElement #-}
52{-# NOINLINE bottomElement #-}
57bottomElement :: a
53bottomElement :: a
58bottomElement = undefined
54bottomElement = undefined
60−data MutableArray s a =
62− { marray# :: Exts.MutableArray# s a
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
70− let ma = MutableArray marr#
73−{-# INLINE writeArray #-}
74−writeArray :: MutableArray s a -> Int -> a -> ST s ()
75−writeArray marr (I# i#) x =
77− case Exts.writeArray# (marray# marr) i# x s# of
78− s'# -> (# s'#, () #))
80−{-# INLINE freezeArray #-}
81−freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
82−freezeArray marr (I# start#) (I# len#) =
84− case Exts.freezeArray# (marray# marr) start# len# s# of
85− (# s'#, arr# #) -> (# s'#, Array arr# #))
87−{-# INLINE copyMutableArray #-}
88−copyMutableArray :: MutableArray s a
89− -- ^ The source array
91− -- ^ Offset into the source array
93− -- ^ Destination array to copy to
95− -- ^ Offset into the destination array
97− -- ^ Number of elements to copy
99−copyMutableArray sourceMarr (I# sourceOffset#) destMarr (I# destOffset#) (I# numElements#) =
101− case Exts.copyMutableArray#
102− (marray# sourceMarr)
108− s'# -> (# s'#, () #))
110−{-# INLINE newArrayIO #-}
111−newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
112−newArrayIO i x = stToIO $ newArray i x
114−{-# INLINE writeArrayIO #-}
115−writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
116−writeArrayIO marr i x = stToIO $ writeArray marr i x
118−{-# INLINE freezeArrayIO #-}
119−freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
120−freezeArrayIO marr i len = stToIO $ freezeArray marr i len
122−{-# INLINE length #-}
123−length :: Array a -> Int
124−length arr = I# (Exts.sizeofArray# (array# arr))
126−{-# INLINE copyMutableArrayIO #-}
127−copyMutableArrayIO ::
128− MutableArray RealWorld a
130− -> MutableArray RealWorld a
134−copyMutableArrayIO sourceMarr sourceOffset destMarr destOffset numElements =
136− copyMutableArray sourceMarr sourceOffset destMarr destOffset numElements
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)
69+length :: Array a -> Int
70+length arr = sizeofArray arr
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
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
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
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
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
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))
212 D.take limit str
134 D.take limit str
213− liftIO $ freezeArrayIO marr 0 i
135+ liftIO $ freezeArray marr 0 i
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)