git.iko.soy/streamly
on master help sign in
Add write and fromStream Data.Array apis and corresponding tests.
ae2dce8Pranay Sashank
src/Streamly/Data/Array.hs+13−8
compare with master
unifiedsplit
@@ -20,12 +20,16 @@ module Streamly.Data.Array
20 ( Array
20 ( Array
21
21
22 -- * Construction
22 -- * Construction
23+ , A.fromListN
24+ , A.fromList
23
25
24 -- Stream Folds
26 -- Stream Folds
25 , fromStreamN
27 , fromStreamN
28+ , fromStream
26
29
27 -- MonadicAPIs
30 -- MonadicAPIs
28 , A.writeN
31 , A.writeN
32+ , A.write
29
33
30 -- * Elimination
34 -- * Elimination
31
35
@@ -34,11 +38,12 @@ module Streamly.Data.Array
34 , read
38 , read
35
39
36 -- * Random Access
40 -- * Random Access
37− , length
38
41
39 -- * Folding Arrays
42 -- * Folding Arrays
40 , streamFold
43 , streamFold
41 , fold
44 , fold
45+
46+ , A.length
42 )
47 )
43where
48where
44
49
@@ -64,6 +69,10 @@ fromStreamN n m = do
64 when (n < 0) $ error "fromStreamN: negative write count specified"
69 when (n < 0) $ error "fromStreamN: negative write count specified"
65 A.fromStreamDN n $ D.toStreamD m
70 A.fromStreamDN n $ D.toStreamD m
66
71
72+{-# INLINE fromStream #-}
73+fromStream :: MonadIO m => SerialT m a -> m (Array a)
74+fromStream m = A.fromStreamD $ D.toStreamD m
75+
67{-# INLINE_EARLY toStream #-}
76{-# INLINE_EARLY toStream #-}
68toStream :: (Monad m, IsStream t) => Array a -> t m a
77toStream :: (Monad m, IsStream t) => Array a -> t m a
69toStream = D.fromStreamD . A.toStreamD
78toStream = D.fromStreamD . A.toStreamD
@@ -72,10 +81,6 @@ toStream = D.fromStreamD . A.toStreamD
72toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
81toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
73toStreamRev = D.fromStreamD . A.toStreamDRev
82toStreamRev = D.fromStreamD . A.toStreamDRev
74
83
75−{-# INLINE length #-}
76−length :: Array a -> Int
77−length Array {..} = aLen
78−
79{-# INLINE fold #-}
84{-# INLINE fold #-}
80fold :: forall m a b . Monad m => Fold m a b -> Array a -> m b
85fold :: forall m a b . Monad m => Fold m a b -> Array a -> m b
81fold f arr = P.runFold f (toStream arr :: SerialT m a)
86fold f arr = P.runFold f (toStream arr :: SerialT m a)
@@ -89,9 +94,9 @@ read :: Monad m => Unfold m (Array a) a
89read = Unfold step inject
94read = Unfold step inject
90 where
95 where
91 inject arr = return (arr, 0)
96 inject arr = return (arr, 0)
92− step (Array {..}, i)
97+ step (arr, i)
93− | i == aLen = return D.Stop
98+ | i == A.length arr = return D.Stop
94 step (arr, (I# i)) =
99 step (arr, (I# i)) =
95 return $
100 return $
96− case Exts.indexArray# (aA arr) i of
101+ case Exts.indexArray# (array# arr) i of
97 (# x #) -> D.Yield x (arr, I# i + 1)
102 (# x #) -> D.Yield x (arr, I# i + 1)
src/Streamly/Internal/Data/Array.hs+99−17
@@ -1,6 +1,5 @@
1{-# LANGUAGE CPP #-}
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE MagicHash #-}
2{-# LANGUAGE MagicHash #-}
3−{-# LANGUAGE RecordWildCards #-}
4{-# LANGUAGE UnboxedTuples #-}
3{-# LANGUAGE UnboxedTuples #-}
5
4
6#include "inline.hs"
5#include "inline.hs"
@@ -20,18 +19,27 @@ module Streamly.Internal.Data.Array
20 , foldl'
19 , foldl'
21 , foldr
20 , foldr
22
21
22+ , length
23+
23 , writeN
24 , writeN
25+ , write
24
26
25 , toStreamD
27 , toStreamD
26 , toStreamDRev
28 , toStreamDRev
29+
30+ , fromListN
31+ , fromList
27 , fromStreamDN
32 , fromStreamDN
33+ , fromStreamD
28 )
34 )
29where
35where
30
36
31−import Prelude hiding (foldr)
37+import Prelude hiding (foldr, length)
38+import Control.DeepSeq (NFData(..))
32import Control.Monad.ST (stToIO, RealWorld)
39import Control.Monad.ST (stToIO, RealWorld)
33import Control.Monad.IO.Class (liftIO, MonadIO)
40import Control.Monad.IO.Class (liftIO, MonadIO)
34import GHC.ST (ST(..))
41import GHC.ST (ST(..))
42+import GHC.IO (unsafePerformIO)
35import GHC.Base (Int(..))
43import GHC.Base (Int(..))
36import Data.Functor.Identity (runIdentity)
44import Data.Functor.Identity (runIdentity)
37import qualified GHC.Exts as Exts
45import qualified GHC.Exts as Exts
@@ -42,8 +50,7 @@ import qualified Streamly.Streams.StreamD as D
42
50
43data Array a =
51data Array a =
44 Array
52 Array
45− { aLen :: {-# UNPACK #-} !Int -- Length of the array
46− , aA :: Exts.Array# a
53+ { array# :: Exts.Array# a
47 }
54 }
48
55
49{-# NOINLINE bottomElement #-}
56{-# NOINLINE bottomElement #-}
@@ -75,7 +82,30 @@ freezeArray :: MutableArray s a -> Int -> Int -> ST s (Array a)
75freezeArray marr (I# start#) (I# len#) =
82freezeArray marr (I# start#) (I# len#) =
76 ST (\s# ->
83 ST (\s# ->
77 case Exts.freezeArray# (marray# marr) start# len# s# of
84 case Exts.freezeArray# (marray# marr) start# len# s# of
78− (# s'#, arr# #) -> (# s'#, Array (I# len#) arr# #))
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'#, () #))
79
109
80{-# INLINE newArrayIO #-}
110{-# INLINE newArrayIO #-}
81newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
111newArrayIO :: Int -> a -> IO (MutableArray RealWorld a)
@@ -89,28 +119,44 @@ writeArrayIO marr i x = stToIO $ writeArray marr i x
89freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
119freezeArrayIO :: MutableArray RealWorld a -> Int -> Int -> IO (Array a)
90freezeArrayIO marr i len = stToIO $ freezeArray marr i len
120freezeArrayIO marr i len = stToIO $ freezeArray marr i len
91
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+
92{-# INLINE_NORMAL toStreamD #-}
138{-# INLINE_NORMAL toStreamD #-}
93toStreamD :: Monad m => Array a -> D.Stream m a
139toStreamD :: Monad m => Array a -> D.Stream m a
94−toStreamD Array {..} = D.Stream step 0
140+toStreamD arr = D.Stream step 0
95 where
141 where
96 {-# INLINE_LATE step #-}
142 {-# INLINE_LATE step #-}
97 step _ i
143 step _ i
98− | i == aLen = return D.Stop
144+ | i == length arr = return D.Stop
99 step _ (I# i) =
145 step _ (I# i) =
100 return $
146 return $
101− case Exts.indexArray# aA i of
147+ case Exts.indexArray# (array# arr) i of
102 (# x #) -> D.Yield x ((I# i) + 1)
148 (# x #) -> D.Yield x ((I# i) + 1)
103
149
104{-# INLINE_NORMAL toStreamDRev #-}
150{-# INLINE_NORMAL toStreamDRev #-}
105toStreamDRev :: Monad m => Array a -> D.Stream m a
151toStreamDRev :: Monad m => Array a -> D.Stream m a
106−toStreamDRev Array {..} = D.Stream step (aLen - 1)
152+toStreamDRev arr = D.Stream step (length arr - 1)
107 where
153 where
108 {-# INLINE_LATE step #-}
154 {-# INLINE_LATE step #-}
109 step _ i
155 step _ i
110 | i < 0 = return D.Stop
156 | i < 0 = return D.Stop
111 step _ (I# i) =
157 step _ (I# i) =
112 return $
158 return $
113− case Exts.indexArray# aA i of
159+ case Exts.indexArray# (array# arr) i of
114 (# x #) -> D.Yield x ((I# i) - 1)
160 (# x #) -> D.Yield x ((I# i) - 1)
115
161
116{-# INLINE_NORMAL foldl' #-}
162{-# INLINE_NORMAL foldl' #-}
@@ -127,13 +173,33 @@ writeN :: MonadIO m => Int -> Fold m a (Array a)
127writeN limit = Fold step initial extract
173writeN limit = Fold step initial extract
128 where
174 where
129 initial = do
175 initial = do
130− marr <- liftIO $ newArrayIO limit bottomElement
176+ marr <- liftIO $ newArrayIO limit bottomElement
131− return (marr, 0)
177+ return (marr, 0)
178+ step (marr, i) x
132− step (marr,i) x | i == limit = return (marr, i)
179+ | i == limit = return (marr, i)
133− | otherwise = do
180+ | otherwise = do
134− liftIO $ writeArrayIO marr i x
181+ liftIO $ writeArrayIO marr i x
135− return (marr, i+1)
182+ return (marr, i + 1)
136− extract (marr,len) = liftIO $ freezeArrayIO marr 0 len
183+ extract (marr, len) = liftIO $ freezeArrayIO marr 0 len
184+
185+{-# INLINE_NORMAL write #-}
186+write :: MonadIO m => Fold m a (Array a)
187+write = Fold step initial extract
188+ where
189+ initial = do
190+ marr <- liftIO $ newArrayIO 0 bottomElement
191+ return (marr, 0, 0)
192+ step (marr, i, capacity) x
193+ | i == capacity =
194+ 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
198+ return (newMarr, i + 1, newCapacity)
199+ | otherwise = do
200+ liftIO $ writeArrayIO marr i x
201+ return (marr, i + 1, capacity)
202+ extract (marr, len, _) = liftIO $ freezeArrayIO marr 0 len
137
203
138{-# INLINE_NORMAL fromStreamDN #-}
204{-# INLINE_NORMAL fromStreamDN #-}
139fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
205fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
@@ -145,3 +211,19 @@ fromStreamDN limit str = do
145 0 $
211 0 $
146 D.take limit str
212 D.take limit str
147 liftIO $ freezeArrayIO marr 0 i
213 liftIO $ freezeArrayIO marr 0 i
214+
215+{-# INLINE fromStreamD #-}
216+fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
217+fromStreamD str = D.runFold write str
218+
219+{-# INLINABLE fromListN #-}
220+fromListN :: Int -> [a] -> Array a
221+fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
222+
223+{-# INLINABLE fromList #-}
224+fromList :: [a] -> Array a
225+fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
226+
227+instance NFData a => NFData (Array a) where
228+ {-# INLINE rnf #-}
229+ rnf = foldl' (\_ x -> rnf x) ()
test/Streamly/Test/Data/Array.hs+58−66
@@ -15,10 +15,13 @@ import Foreign.Storable (Storable(..))
15
15
16import Test.Hspec.QuickCheck
16import Test.Hspec.QuickCheck
17import Test.QuickCheck (Property, forAll, Gen, vectorOf, arbitrary, choose)
17import Test.QuickCheck (Property, forAll, Gen, vectorOf, arbitrary, choose)
18−import Test.QuickCheck.Monadic (monadicIO, assert)
18+import Test.QuickCheck.Monadic (monadicIO, assert, run)
19
19
20import Test.Hspec as H
20import Test.Hspec as H
21
21
22+import Streamly (SerialT)
23+import Streamly.Data.Array (Array)
24+
22import qualified Streamly.Data.Array as A
25import qualified Streamly.Data.Array as A
23import qualified Streamly.Prelude as S
26import qualified Streamly.Prelude as S
24
27
@@ -44,78 +47,64 @@ defaultChunkSize = 32 * k - allocOverhead
44maxArrLen :: Int
47maxArrLen :: Int
45maxArrLen = defaultChunkSize * 8
48maxArrLen = defaultChunkSize * 8
46
49
50+genericTestFrom ::
51+ (Int -> SerialT IO Int -> IO (Array Int))
47−testLength :: Property
52+ -> Property
48−testLength =
53+genericTestFrom arrFold =
49 forAll (choose (0, maxArrLen)) $ \len ->
54 forAll (choose (0, maxArrLen)) $ \len ->
50 forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
55 forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
51 monadicIO $ do
56 monadicIO $ do
52− arr <- S.fold (A.writeN len)
53− $ S.fromList list
57+ arr <- run $ arrFold len $ S.fromList list
54 assert (A.length arr == len)
58 assert (A.length arr == len)
55
59
60+testLength :: Property
61+testLength = genericTestFrom (\n -> S.fold (A.writeN n))
62+
56testLengthFromStreamN :: Property
63testLengthFromStreamN :: Property
57−testLengthFromStreamN =
64+testLengthFromStreamN = genericTestFrom A.fromStreamN
58− forAll (choose (0, maxArrLen)) $ \len ->
59− forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
60− monadicIO $ do
61− arr <- A.fromStreamN len
62− $ S.fromList list
63− assert (A.length arr == len)
64
65
65−testFoldUnfold :: Property
66+testLengthFromStream :: Property
66−testFoldUnfold =
67+testLengthFromStream = genericTestFrom (const A.fromStream)
67− forAll (choose (0, maxArrLen)) $ \len ->
68+
68− forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
69+genericTestFromTo ::
69− monadicIO $ do
70+ (Int -> SerialT IO Int -> IO (Array Int))
70− arr <- S.fold (A.writeN len)
71+ -> (Array Int -> SerialT IO Int)
71− $ S.fromList list
72+ -> ([Int] -> [Int] -> Bool)
72− xs <- S.toList
73+ -> Property
73− $ S.unfold A.read arr
74+genericTestFromTo arrFold arrUnfold listEq =
74− assert (xs == list)
75−
76−testFoldToStream :: Property
77−testFoldToStream =
78− forAll (choose (0, maxArrLen)) $ \len ->
79− forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
80− monadicIO $ do
81− arr <- S.fold (A.writeN len)
82− $ S.fromList list
83− xs <- S.toList
84− $ A.toStream arr
85− assert (xs == list)
86−
87−testFoldToStreamRev :: Property
88−testFoldToStreamRev =
89− forAll (choose (0, maxArrLen)) $ \len ->
90− forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
91− monadicIO $ do
92− arr <- S.fold (A.writeN len)
93− $ S.fromList list
94− xs <- S.toList
95− $ A.toStreamRev arr
96− assert (xs == reverse list)
97−
98−testFromStreamUnfold :: Property
99−testFromStreamUnfold =
100 forAll (choose (0, maxArrLen)) $ \len ->
75 forAll (choose (0, maxArrLen)) $ \len ->
101 forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
76 forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
102 monadicIO $ do
77 monadicIO $ do
103− arr <- A.fromStreamN len
104− $ S.fromList list
78+ arr <- run $ arrFold len $ S.fromList list
105− xs <- S.toList
79+ xs <- run $ S.toList $ arrUnfold arr
106− $ S.unfold A.read arr
107− assert (xs == list)
80+ assert (listEq xs list)
81+
82+testFoldNUnfold :: Property
83+testFoldNUnfold =
84+ genericTestFromTo (\n -> S.fold (A.writeN n)) (S.unfold A.read) (==)
85+
86+testFoldNToStream :: Property
87+testFoldNToStream =
88+ genericTestFromTo (\n -> S.fold (A.writeN n)) A.toStream (==)
89+
90+testFoldNToStreamRev :: Property
91+testFoldNToStreamRev =
92+ genericTestFromTo
93+ (\n -> S.fold (A.writeN n))
94+ A.toStreamRev
95+ (\xs list -> xs == reverse list)
96+
97+testFromStreamNUnfold :: Property
98+testFromStreamNUnfold = genericTestFromTo A.fromStreamN (S.unfold A.read) (==)
99+
100+testFromStreamNToStream :: Property
101+testFromStreamNToStream = genericTestFromTo A.fromStreamN A.toStream (==)
108
102
109testFromStreamToStream :: Property
103testFromStreamToStream :: Property
110−testFromStreamToStream =
104+testFromStreamToStream = genericTestFromTo (const A.fromStream) A.toStream (==)
111− forAll (choose (0, maxArrLen)) $ \len ->
105+
112− forAll (vectorOf len (arbitrary :: Gen Int)) $ \list ->
106+testFoldUnfold :: Property
113− monadicIO $ do
107+testFoldUnfold = genericTestFromTo (const (S.fold A.write)) (S.unfold A.read) (==)
114− arr <- A.fromStreamN len
115− $ S.fromList list
116− xs <- S.toList
117− $ A.toStream arr
118− assert (xs == list)
119
108
120main :: IO ()
109main :: IO ()
121main =
110main =
@@ -125,8 +114,11 @@ main =
125 describe "Construction" $ do
114 describe "Construction" $ do
126 prop "length . writeN n === n" testLength
115 prop "length . writeN n === n" testLength
127 prop "length . fromStreamN n === n" testLengthFromStreamN
116 prop "length . fromStreamN n === n" testLengthFromStreamN
128− prop "read . writeN === id " testFoldUnfold
117+ prop "length . fromStream === n" testLengthFromStream
118+ prop "read . writeN === id " testFoldNUnfold
129− prop "toStream . writeN === id" testFoldToStream
119+ prop "toStream . writeN === id" testFoldNToStream
130− prop "toStreamRev . writeN === reverse" testFoldToStreamRev
120+ prop "toStreamRev . writeN === reverse" testFoldNToStreamRev
131− prop "read . fromStreamN === id" testFromStreamUnfold
121+ prop "read . fromStreamN === id" testFromStreamNUnfold
132− prop "toStream . fromStreamN === reverse" testFromStreamToStream
122+ prop "toStream . fromStreamN === id" testFromStreamNToStream
123+ prop "toStream . fromStream === id" testFromStreamToStream
124+ prop "read . write === id" testFoldUnfold