git.iko.soy/streamly
on 175b0b6 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
147 lines 4.2 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE MagicHash #-}
3{-# LANGUAGE RecordWildCards #-}
4{-# LANGUAGE UnboxedTuples #-}
6#include "inline.hs"
8-- |
9-- Module : Streamly.Internal.Data.Array
10-- Copyright : (c) 2019 Composewell Technologies
11--
12-- License : BSD-3-Clause
13-- Maintainer : streamly@composewell.com
14-- Stability : experimental
15-- Portability : GHC
16--
17module Streamly.Internal.Data.Array
18 ( Array(..)
20 , foldl'
21 , foldr
23 , writeN
25 , toStreamD
26 , toStreamDRev
27 , fromStreamDN
28 )
29where
31import Prelude hiding (foldr)
32import Control.Monad.ST (stToIO, RealWorld)
33import Control.Monad.IO.Class (liftIO, MonadIO)
34import GHC.ST (ST(..))
35import GHC.Base (Int(..))
36import Data.Functor.Identity (runIdentity)
37import qualified GHC.Exts as Exts
39import Streamly.Internal.Data.Fold.Types (Fold(..))
41import qualified Streamly.Streams.StreamD as D
43data Array a =
44 Array
45 { aLen :: {-# UNPACK #-} !Int -- Length of the array
46 , aA :: Exts.Array# a
47 }
49{-# NOINLINE bottomElement #-}
50bottomElement :: a
51bottomElement = undefined
53data MutableArray s a =
54 MutableArray
55 { marray# :: Exts.MutableArray# s a
56 }
58{-# INLINE newArray #-}
59newArray :: Int -> a -> ST s (MutableArray s a)
60newArray (I# n#) x = ST
61 (\s# -> case Exts.newArray# n# x s# of
62 (# s'#, marr# #) ->
63 let ma = MutableArray marr#
64 in (# s'# , ma #))
66{-# INLINE writeArray #-}
67writeArray :: MutableArray s a -> Int -> a -> ST s ()
68writeArray marr (I# i#) x =
69 ST (\s# ->
70 case Exts.writeArray# (marray# marr) i# x s# of
71 s'# -> (# s'#, () #))
73{-# INLINE freezeArray #-}
74freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
75freezeArray marr (I# start#) (I# len#) =
76 ST (\s# ->
77 case Exts.freezeArray# (marray# marr) start# len# s# of
78 (# s'#, arr# #) -> (# s'#, Array (I# len#) arr# #))
80{-# INLINE newArrayIO #-}
81newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
82newArrayIO i x = stToIO $ newArray i x
84{-# INLINE writeArrayIO #-}
85writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
86writeArrayIO marr i x = stToIO $ writeArray marr i x
88{-# INLINE freezeArrayIO #-}
89freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
90freezeArrayIO marr i len = stToIO $ freezeArray marr i len
92{-# INLINE_NORMAL toStreamD #-}
93toStreamD :: Monad m => Array a -> D.Stream m a
94toStreamD Array {..} = D.Stream step 0
95 where
96 {-# INLINE_LATE step #-}
97 step _ i
98 | i == aLen = return D.Stop
99 step _ (I# i) =
100 return $
101 case Exts.indexArray# aA i of
102 (# x #) -> D.Yield x ((I# i) + 1)
104{-# INLINE_NORMAL toStreamDRev #-}
105toStreamDRev :: Monad m => Array a -> D.Stream m a
106toStreamDRev Array {..} = D.Stream step (aLen - 1)
107 where
108 {-# INLINE_LATE step #-}
109 step _ i
110 | i < 0 = return D.Stop
111 step _ (I# i) =
112 return $
113 case Exts.indexArray# aA i of
114 (# x #) -> D.Yield x ((I# i) - 1)
116{-# INLINE_NORMAL foldl' #-}
117foldl' :: (b -> a -> b) -> b -> Array a -> b
118foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
120{-# INLINE_NORMAL foldr #-}
121foldr :: (a -> b -> b) -> b -> Array a -> b
122foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
124-- writeN n = S.evertM (fromStreamDN n)
125{-# INLINE_NORMAL writeN #-}
126writeN :: MonadIO m => Int -> Fold m a (Array a)
127writeN limit = Fold step initial extract
128 where
129 initial = do
130 marr <- liftIO $ newArrayIO limit bottomElement
131 return (marr, 0)
132 step (marr,i) x | i == limit = return (marr, i)
133 | otherwise = do
134 liftIO $ writeArrayIO marr i x
135 return (marr, i+1)
136 extract (marr,len) = liftIO $ freezeArrayIO marr 0 len
138{-# INLINE_NORMAL fromStreamDN #-}
139fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
140fromStreamDN limit str = do
141 marr <- liftIO $ newArrayIO (max limit 0) bottomElement
142 i <-
143 D.foldlM'
144 (\i x -> i `seq` (liftIO $ writeArrayIO marr i x) >> return (i + 1))
145 0 $
146 D.take limit str
147 liftIO $ freezeArrayIO marr 0 i