mirror of
https://github.com/ilyakooo0/streamly.git
synced 2024-10-26 09:59:48 +03:00
Use Array from Data.Primitive.Array rather than defining own Array.
This commit is contained in:
Pranay Sashank
parent
72d14d4581
commit
0d91f2ece9
@ -1,3 +1,5 @@
|
||||
{-# OPTIONS_GHC -fno-warn-orphans #-}
|
||||
|
||||
{-# LANGUAGE CPP #-}
|
||||
{-# LANGUAGE MagicHash #-}
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
@ -36,105 +38,21 @@ where
|
||||
|
||||
import Prelude hiding (foldr, length)
|
||||
import Control.DeepSeq (NFData(..))
|
||||
import Control.Monad.ST (stToIO, RealWorld)
|
||||
import Control.Monad.IO.Class (liftIO, MonadIO)
|
||||
import GHC.ST (ST(..))
|
||||
import GHC.IO (unsafePerformIO)
|
||||
import GHC.Base (Int(..))
|
||||
import Data.Functor.Identity (runIdentity)
|
||||
import Data.Primitive.Array hiding (fromList, fromListN)
|
||||
import qualified GHC.Exts as Exts
|
||||
|
||||
import Streamly.Internal.Data.Fold.Types (Fold(..))
|
||||
|
||||
import qualified Streamly.Streams.StreamD as D
|
||||
|
||||
data Array a =
|
||||
Array
|
||||
{ array# :: Exts.Array# a
|
||||
}
|
||||
|
||||
{-# NOINLINE bottomElement #-}
|
||||
bottomElement :: a
|
||||
bottomElement = undefined
|
||||
|
||||
data MutableArray s a =
|
||||
MutableArray
|
||||
{ marray# :: Exts.MutableArray# s a
|
||||
}
|
||||
|
||||
{-# INLINE newArray #-}
|
||||
newArray :: Int -> a -> ST s (MutableArray s a)
|
||||
newArray (I# n#) x = ST
|
||||
(\s# -> case Exts.newArray# n# x s# of
|
||||
(# s'#, marr# #) ->
|
||||
let ma = MutableArray marr#
|
||||
in (# s'# , ma #))
|
||||
|
||||
{-# INLINE writeArray #-}
|
||||
writeArray :: MutableArray s a -> Int -> a -> ST s ()
|
||||
writeArray marr (I# i#) x =
|
||||
ST (\s# ->
|
||||
case Exts.writeArray# (marray# marr) i# x s# of
|
||||
s'# -> (# s'#, () #))
|
||||
|
||||
{-# INLINE freezeArray #-}
|
||||
freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
|
||||
freezeArray marr (I# start#) (I# len#) =
|
||||
ST (\s# ->
|
||||
case Exts.freezeArray# (marray# marr) start# len# s# of
|
||||
(# s'#, arr# #) -> (# s'#, Array arr# #))
|
||||
|
||||
{-# INLINE copyMutableArray #-}
|
||||
copyMutableArray :: MutableArray s a
|
||||
-- ^ The source array
|
||||
-> Int
|
||||
-- ^ Offset into the source array
|
||||
-> MutableArray s a
|
||||
-- ^ Destination array to copy to
|
||||
-> Int
|
||||
-- ^ Offset into the destination array
|
||||
-> Int
|
||||
-- ^ Number of elements to copy
|
||||
-> ST s ()
|
||||
copyMutableArray sourceMarr (I# sourceOffset#) destMarr (I# destOffset#) (I# numElements#) =
|
||||
ST (\s# ->
|
||||
case Exts.copyMutableArray#
|
||||
(marray# sourceMarr)
|
||||
sourceOffset#
|
||||
(marray# destMarr)
|
||||
destOffset#
|
||||
numElements#
|
||||
s# of
|
||||
s'# -> (# s'#, () #))
|
||||
|
||||
{-# INLINE newArrayIO #-}
|
||||
newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
|
||||
newArrayIO i x = stToIO $ newArray i x
|
||||
|
||||
{-# INLINE writeArrayIO #-}
|
||||
writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
|
||||
writeArrayIO marr i x = stToIO $ writeArray marr i x
|
||||
|
||||
{-# INLINE freezeArrayIO #-}
|
||||
freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
|
||||
freezeArrayIO marr i len = stToIO $ freezeArray marr i len
|
||||
|
||||
{-# INLINE length #-}
|
||||
length :: Array a -> Int
|
||||
length arr = I# (Exts.sizeofArray# (array# arr))
|
||||
|
||||
{-# INLINE copyMutableArrayIO #-}
|
||||
copyMutableArrayIO ::
|
||||
MutableArray RealWorld a
|
||||
-> Int
|
||||
-> MutableArray RealWorld a
|
||||
-> Int
|
||||
-> Int
|
||||
-> IO ()
|
||||
copyMutableArrayIO sourceMarr sourceOffset destMarr destOffset numElements =
|
||||
stToIO $
|
||||
copyMutableArray sourceMarr sourceOffset destMarr destOffset numElements
|
||||
|
||||
{-# INLINE_NORMAL toStreamD #-}
|
||||
toStreamD :: Monad m => Array a -> D.Stream m a
|
||||
toStreamD arr = D.Stream step 0
|
||||
@ -147,6 +65,10 @@ toStreamD arr = D.Stream step 0
|
||||
case Exts.indexArray# (array# arr) i of
|
||||
(# x #) -> D.Yield x ((I# i) + 1)
|
||||
|
||||
{-# INLINE length #-}
|
||||
length :: Array a -> Int
|
||||
length arr = sizeofArray arr
|
||||
|
||||
{-# INLINE_NORMAL toStreamDRev #-}
|
||||
toStreamDRev :: Monad m => Array a -> D.Stream m a
|
||||
toStreamDRev arr = D.Stream step (length arr - 1)
|
||||
@ -173,44 +95,44 @@ writeN :: MonadIO m => Int -> Fold m a (Array a)
|
||||
writeN limit = Fold step initial extract
|
||||
where
|
||||
initial = do
|
||||
marr <- liftIO $ newArrayIO limit bottomElement
|
||||
marr <- liftIO $ newArray limit bottomElement
|
||||
return (marr, 0)
|
||||
step (marr, i) x
|
||||
| i == limit = return (marr, i)
|
||||
| otherwise = do
|
||||
liftIO $ writeArrayIO marr i x
|
||||
liftIO $ writeArray marr i x
|
||||
return (marr, i + 1)
|
||||
extract (marr, len) = liftIO $ freezeArrayIO marr 0 len
|
||||
extract (marr, len) = liftIO $ freezeArray marr 0 len
|
||||
|
||||
{-# INLINE_NORMAL write #-}
|
||||
write :: MonadIO m => Fold m a (Array a)
|
||||
write = Fold step initial extract
|
||||
where
|
||||
initial = do
|
||||
marr <- liftIO $ newArrayIO 0 bottomElement
|
||||
marr <- liftIO $ newArray 0 bottomElement
|
||||
return (marr, 0, 0)
|
||||
step (marr, i, capacity) x
|
||||
| i == capacity =
|
||||
let newCapacity = max (capacity * 2) 1
|
||||
in do newMarr <- liftIO $ newArrayIO newCapacity bottomElement
|
||||
liftIO $ copyMutableArrayIO marr 0 newMarr 0 i
|
||||
liftIO $ writeArrayIO newMarr i x
|
||||
in do newMarr <- liftIO $ newArray newCapacity bottomElement
|
||||
liftIO $ copyMutableArray newMarr 0 marr 0 i
|
||||
liftIO $ writeArray newMarr i x
|
||||
return (newMarr, i + 1, newCapacity)
|
||||
| otherwise = do
|
||||
liftIO $ writeArrayIO marr i x
|
||||
liftIO $ writeArray marr i x
|
||||
return (marr, i + 1, capacity)
|
||||
extract (marr, len, _) = liftIO $ freezeArrayIO marr 0 len
|
||||
extract (marr, len, _) = liftIO $ freezeArray marr 0 len
|
||||
|
||||
{-# INLINE_NORMAL fromStreamDN #-}
|
||||
fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
|
||||
fromStreamDN limit str = do
|
||||
marr <- liftIO $ newArrayIO (max limit 0) bottomElement
|
||||
marr <- liftIO $ newArray (max limit 0) bottomElement
|
||||
i <-
|
||||
D.foldlM'
|
||||
(\i x -> i `seq` (liftIO $ writeArrayIO marr i x) >> return (i + 1))
|
||||
(\i x -> i `seq` (liftIO $ writeArray marr i x) >> return (i + 1))
|
||||
0 $
|
||||
D.take limit str
|
||||
liftIO $ freezeArrayIO marr 0 i
|
||||
liftIO $ freezeArray marr 0 i
|
||||
|
||||
{-# INLINE fromStreamD #-}
|
||||
fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
|
||||
|
||||
@ -361,6 +361,7 @@ library
|
||||
, containers >= 0.5.8.2 && < 0.7
|
||||
, heaps >= 0.3 && < 0.4
|
||||
, directory >= 1.3 && < 1.4
|
||||
, primitive >= 0.6.4.0 && < 0.8
|
||||
|
||||
-- concurrency
|
||||
, atomic-primops >= 0.8 && < 0.9
|
||||
|
||||
Loading…
Reference in New Issue
Block a user