@@ -1,3 +1,5 @@
1+{-# OPTIONS_GHC -fno-warn-orphans #-}
2+
13{-# LANGUAGE CPP #-}
24{-# LANGUAGE MagicHash #-}
35{-# LANGUAGE UnboxedTuples #-}
@@ -36,105 +38,21 @@ where
3638
3739import Prelude hiding (foldr, length)
3840import Control.DeepSeq (NFData(..))
39−import Control.Monad.ST (stToIO, RealWorld)
4041import Control.Monad.IO.Class (liftIO, MonadIO)
41−import GHC.ST (ST(..))
4242import GHC.IO (unsafePerformIO)
4343import GHC.Base (Int(..))
4444import Data.Functor.Identity (runIdentity)
45+import Data.Primitive.Array hiding (fromList, fromListN)
4546import qualified GHC.Exts as Exts
4647
4748import Streamly.Internal.Data.Fold.Types (Fold(..))
4849
4950import qualified Streamly.Streams.StreamD as D
5051
51−data Array a =
52− Array
53− { array# :: Exts.Array# a
54− }
55−
5652{-# NOINLINE bottomElement #-}
5753bottomElement :: a
5854bottomElement = undefined
5955
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−
13856{-# INLINE_NORMAL toStreamD #-}
13957toStreamD :: Monad m => Array a -> D.Stream m a
14058toStreamD arr = D.Stream step 0
@@ -147,6 +65,10 @@ toStreamD arr = D.Stream step 0
14765 case Exts.indexArray# (array# arr) i of
14866 (# x #) -> D.Yield x ((I# i) + 1)
14967
68+{-# INLINE length #-}
69+length :: Array a -> Int
70+length arr = sizeofArray arr
71+
15072{-# INLINE_NORMAL toStreamDRev #-}
15173toStreamDRev :: Monad m => Array a -> D.Stream m a
15274toStreamDRev arr = D.Stream step (length arr - 1)
@@ -173,44 +95,44 @@ writeN :: MonadIO m => Int -> Fold m a (Array a)
17395writeN limit = Fold step initial extract
17496 where
17597 initial = do
176− marr <- liftIO $ newArrayIO limit bottomElement
98+ marr <- liftIO $ newArray limit bottomElement
17799 return (marr, 0)
178100 step (marr, i) x
179101 | i == limit = return (marr, i)
180102 | otherwise = do
181− liftIO $ writeArrayIO marr i x
103+ liftIO $ writeArray marr i x
182104 return (marr, i + 1)
183− extract (marr, len) = liftIO $ freezeArrayIO marr 0 len
105+ extract (marr, len) = liftIO $ freezeArray marr 0 len
184106
185107{-# INLINE_NORMAL write #-}
186108write :: MonadIO m => Fold m a (Array a)
187109write = Fold step initial extract
188110 where
189111 initial = do
190− marr <- liftIO $ newArrayIO 0 bottomElement
112+ marr <- liftIO $ newArray 0 bottomElement
191113 return (marr, 0, 0)
192114 step (marr, i, capacity) x
193115 | i == capacity =
194116 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
117+ in do newMarr <- liftIO $ newArray newCapacity bottomElement
118+ liftIO $ copyMutableArray newMarr 0 marr 0 i
119+ liftIO $ writeArray newMarr i x
198120 return (newMarr, i + 1, newCapacity)
199121 | otherwise = do
200− liftIO $ writeArrayIO marr i x
122+ liftIO $ writeArray marr i x
201123 return (marr, i + 1, capacity)
202− extract (marr, len, _) = liftIO $ freezeArrayIO marr 0 len
124+ extract (marr, len, _) = liftIO $ freezeArray marr 0 len
203125
204126{-# INLINE_NORMAL fromStreamDN #-}
205127fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
206128fromStreamDN limit str = do
207− marr <- liftIO $ newArrayIO (max limit 0) bottomElement
129+ marr <- liftIO $ newArray (max limit 0) bottomElement
208130 i <-
209131 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))
211133 0 $
212134 D.take limit str
213− liftIO $ freezeArrayIO marr 0 i
135+ liftIO $ freezeArray marr 0 i
214136
215137{-# INLINE fromStreamD #-}
216138fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)