on master help sign in
Add support for unpinned Arrays.
94bd0faPranay Sashank
src/Streamly/Data/Array.hsnew file+79
@@ -0,0 +1,79 @@
1+{-# LANGUAGE CPP #-}
2+{-# LANGUAGE RankNTypes #-}
3+{-# LANGUAGE RecordWildCards #-}
4+{-# LANGUAGE ScopedTypeVariables #-}
5+
6+#include "inline.hs"
7+
8+-- |
9+-- Module : Streamly.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+--
17+module Streamly.Data.Array
18+ ( Array
19+
20+ -- * Construction
21+
22+ -- Stream Folds
23+ , fromStreamN
24+
25+ -- MonadicAPIs
26+ , A.writeN
27+
28+ -- * Elimination
29+
30+ , toStream
31+ , toStreamRev
32+
33+ -- * Random Access
34+ , length
35+
36+ -- * Folding Arrays
37+ , streamFold
38+ , fold
39+ )
40+where
41+
42+import Prelude hiding (length)
43+import Control.Monad (when)
44+import Control.Monad.IO.Class (MonadIO)
45+
46+import Streamly.Streams.Serial (SerialT)
47+import Streamly.Streams.StreamK.Type (IsStream)
48+import Streamly.Internal.Data.Fold.Types (Fold(..))
49+import Streamly.Internal.Data.Array (Array(..))
50+
51+import qualified Streamly.Internal.Data.Array as A
52+import qualified Streamly.Streams.Prelude as P
53+import qualified Streamly.Streams.StreamD as D
54+
55+{-# INLINE fromStreamN #-}
56+fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
57+fromStreamN n m = do
58+ when (n < 0) $ error "fromStreamN: negative write count specified"
59+ A.fromStreamDN n $ D.toStreamD m
60+
61+{-# INLINE_EARLY toStream #-}
62+toStream :: (Monad m, IsStream t) => Array a -> t m a
63+toStream = D.fromStreamD . A.toStreamD
64+
65+{-# INLINE_EARLY toStreamRev #-}
66+toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
67+toStreamRev = D.fromStreamD . A.toStreamDRev
68+
69+{-# INLINE length #-}
70+length :: Array a -> Int
71+length Array {..} = aLen
72+
73+{-# INLINE fold #-}
74+fold :: forall m a b . Monad m => Fold m a b -> Array a -> m b
75+fold f arr = P.runFold f (toStream arr :: SerialT m a)
76+
77+{-# INLINE streamFold #-}
78+streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
79+streamFold f arr = f (toStream arr)
src/Streamly/Internal/Data/Array.hsnew file+147
@@ -0,0 +1,147 @@
1+{-# LANGUAGE CPP #-}
2+{-# LANGUAGE MagicHash #-}
3+{-# LANGUAGE RecordWildCards #-}
4+{-# LANGUAGE UnboxedTuples #-}
5+
6+#include "inline.hs"
7+
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+--
17+module Streamly.Internal.Data.Array
18+ ( Array(..)
19+
20+ , foldl'
21+ , foldr
22+
23+ , writeN
24+
25+ , toStreamD
26+ , toStreamDRev
27+ , fromStreamDN
28+ )
29+where
30+
31+import Prelude hiding (foldr)
32+import Control.Monad.ST (stToIO, RealWorld)
33+import Control.Monad.IO.Class (liftIO, MonadIO)
34+import GHC.ST (ST(..))
35+import GHC.Base (Int(..))
36+import Data.Functor.Identity (runIdentity)
37+import qualified GHC.Exts as Exts
38+
39+import Streamly.Internal.Data.Fold.Types (Fold(..))
40+
41+import qualified Streamly.Streams.StreamD as D
42+
43+data Array a =
44+ Array
45+ { aLen :: {-# UNPACK #-} !Int -- Length of the array
46+ , aA :: Exts.Array# a
47+ }
48+
49+{-# NOINLINE bottomElement #-}
50+bottomElement :: a
51+bottomElement = undefined
52+
53+data MutableArray s a =
54+ MutableArray
55+ { marray# :: Exts.MutableArray# s a
56+ }
57+
58+{-# INLINE newArray #-}
59+newArray :: Int -> a -> ST s (MutableArray s a)
60+newArray (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 #))
65+
66+{-# INLINE writeArray #-}
67+writeArray :: MutableArray s a -> Int -> a -> ST s ()
68+writeArray marr (I# i#) x =
69+ ST (\s# ->
70+ case Exts.writeArray# (marray# marr) i# x s# of
71+ s'# -> (# s'#, () #))
72+
73+{-# INLINE freezeArray #-}
74+freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
75+freezeArray 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# #))
79+
80+{-# INLINE newArrayIO #-}
81+newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
82+newArrayIO i x = stToIO $ newArray i x
83+
84+{-# INLINE writeArrayIO #-}
85+writeArrayIO :: MutableArray RealWorld a -> Int -> a -> IO ()
86+writeArrayIO marr i x = stToIO $ writeArray marr i x
87+
88+{-# INLINE freezeArrayIO #-}
89+freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
90+freezeArrayIO marr i len = stToIO $ freezeArray marr i len
91+
92+{-# INLINE_NORMAL toStreamD #-}
93+toStreamD :: Monad m => Array a -> D.Stream m a
94+toStreamD 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)
103+
104+{-# INLINE_NORMAL toStreamDRev #-}
105+toStreamDRev :: Monad m => Array a -> D.Stream m a
106+toStreamDRev 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)
115+
116+{-# INLINE_NORMAL foldl' #-}
117+foldl' :: (b -> a -> b) -> b -> Array a -> b
118+foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
119+
120+{-# INLINE_NORMAL foldr #-}
121+foldr :: (a -> b -> b) -> b -> Array a -> b
122+foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
123+
124+-- writeN n = S.evertM (fromStreamDN n)
125+{-# INLINE_NORMAL writeN #-}
126+writeN :: MonadIO m => Int -> Fold m a (Array a)
127+writeN 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
137+
138+{-# INLINE_NORMAL fromStreamDN #-}
139+fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
140+fromStreamDN 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
streamly.cabal+2
@@ -301,6 +301,7 @@ library
301
301
302 exposed-modules: Streamly.Prelude
302 exposed-modules: Streamly.Prelude
303 , Streamly
303 , Streamly
304+ , Streamly.Data.Array
304 , Streamly.Data.Fold
305 , Streamly.Data.Fold
305 , Streamly.Data.Unfold
306 , Streamly.Data.Unfold
306 , Streamly.Data.Unicode.Stream
307 , Streamly.Data.Unicode.Stream
@@ -320,6 +321,7 @@ library
320 , Streamly.Internal.Data.Time.Clock
321 , Streamly.Internal.Data.Time.Clock
321 , Streamly.Internal.Data.Stream.StreamD.Type
322 , Streamly.Internal.Data.Stream.StreamD.Type
322 , Streamly.Internal.Data.SVar
323 , Streamly.Internal.Data.SVar
324+ , Streamly.Internal.Data.Array
323 , Streamly.Internal.Memory.Array.Types
325 , Streamly.Internal.Memory.Array.Types
324 , Streamly.Internal.Memory.Array
326 , Streamly.Internal.Memory.Array
325 , Streamly.Internal.Memory.ArrayStream
327 , Streamly.Internal.Memory.ArrayStream