@@ -11,10 +11,11 @@ module Main (main) where
1111
1212import Control.DeepSeq (NFData(..))
1313import System.Random (randomRIO)
14+import Streamly.Internal.Data.Unfold (Unfold)
1415
1516import qualified Prelude
17+import qualified Streamly.Internal.Data.Fold as FL
1618import qualified Streamly.Internal.Data.Unfold as UF
17−import qualified Streamly.Internal.Data.Stream.IsStream as S
1819
1920import Streamly.Benchmark.Common
2021import Streamly.Benchmark.Data.NestedUnfoldOps
@@ -25,6 +26,52 @@ import Prelude hiding (concat, take)
2526benchIO :: (NFData b) => String -> (Int -> IO b) -> Benchmark
2627benchIO name f = bench name $ nfIO $ randomRIO (1,1) >>= f
2728
29+-------------------------------------------------------------------------------
30+-- Benchmark helpers
31+-------------------------------------------------------------------------------
32+
33+{-# INLINE drainUnfold #-}
34+drainUnfold :: Monad m => Unfold m a b -> a -> m ()
35+drainUnfold unf seed = UF.fold unf FL.drain seed
36+
37+{-# INLINE drainGeneration #-}
38+drainGeneration :: Monad m => Unfold m a b -> a -> m ()
39+drainGeneration = drainUnfold
40+
41+{-# INLINE drainTransformation #-}
42+drainTransformation ::
43+ Monad m => Unfold m a b -> (Unfold m a b -> Unfold m c d) -> c -> m ()
44+drainTransformation unf f seed = drainUnfold (f unf) seed
45+
46+{-# INLINE drainTransformationDefault #-}
47+drainTransformationDefault ::
48+ Monad m => Int -> (Unfold m Int Int -> Unfold m c d) -> c -> m ()
49+drainTransformationDefault size =
50+ drainTransformation (UF.take size UF.enumerateFromIntegral)
51+
52+{-# INLINE drainProduct #-}
53+drainProduct ::
54+ Monad m
55+ => Unfold m a b
56+ -> Unfold m c d
57+ -> (Unfold m a b -> Unfold m c d -> Unfold m e f)
58+ -> e
59+ -> m ()
60+drainProduct unf1 unf2 f seed = drainUnfold (f unf1 unf2) seed
61+
62+{-# INLINE drainProductDefault #-}
63+drainProductDefault ::
64+ Monad m
65+ => Int
66+ -> (Unfold m Int Int -> Unfold m Int Int -> Unfold m e f)
67+ -> e
68+ -> m ()
69+drainProductDefault size = drainProduct src src
70+
71+ where
72+
73+ src = UF.take size UF.enumerateFromIntegral
74+
2875-------------------------------------------------------------------------------
2976-- Stream generation
3077-------------------------------------------------------------------------------
@@ -32,26 +79,25 @@ benchIO name f = bench name $ nfIO $ randomRIO (1,1) >>= f
3279{-# INLINE fromListM #-}
3380fromListM :: Monad m => Int -> Int -> m ()
3481fromListM size start =
35− S.drain
36− $ S.unfold UF.fromListM (Prelude.map return [start .. start + size])
82+ drainGeneration UF.fromListM (Prelude.map return [start .. start + size])
3783
3884{-# INLINE replicateM #-}
3985replicateM :: Monad m => Int -> Int -> m ()
40−replicateM size start = S.drain $ S.unfold (UF.replicateM size) (return start)
86+replicateM size start = drainGeneration (UF.replicateM size) (return start)
4187
4288{-# INLINE repeatM #-}
4389repeatM :: Monad m => Int -> Int -> m ()
44−repeatM size start = S.drain $ S.take size $ S.unfold UF.repeatM (return start)
90+repeatM size start = drainGeneration (UF.take size UF.repeatM) (return start)
4591
4692{-# INLINE enumerateFromStepIntegral #-}
4793enumerateFromStepIntegral :: Monad m => Int -> Int -> m ()
48−enumerateFromStepIntegral size start = S.drain $
49− S.take size $ S.unfold UF.enumerateFromStepIntegral (start, 1)
94+enumerateFromStepIntegral size start =
95+ drainGeneration (UF.take size UF.enumerateFromStepIntegral) (start, 1)
5096
5197{-# INLINE enumerateFromToIntegral #-}
5298enumerateFromToIntegral :: Monad m => Int -> Int -> m ()
5399enumerateFromToIntegral size start =
54− S.drain $ S.unfold (UF.enumerateFromToIntegral (size + start)) start
100+ drainGeneration (UF.enumerateFromToIntegral (size + start)) start
55101
56102-------------------------------------------------------------------------------
57103-- Stream transformation
@@ -59,55 +105,41 @@ enumerateFromToIntegral size start =
59105
60106{-# INLINE take #-}
61107take :: Monad m => Int -> Int -> m ()
62−take size start =
63− S.drain
64− $ S.unfold
65− (UF.take size (UF.enumerateFromToIntegral (size + start)))
66− start
108+take size start = drainTransformationDefault size (UF.take size) start
67109
68110{-# INLINE takeWhileM #-}
69111takeWhileM :: Monad m => Int -> Int -> m ()
70112takeWhileM size start =
71− S.drain
72− $ S.unfold
73− (UF.takeWhileM
74− (\b -> return (b <= size + start))
75− (UF.enumerateFromToIntegral (size + start)))
76− start
113+ drainTransformationDefault
114+ size
115+ (UF.takeWhileM (\b -> return (b <= size + start)))
116+ start
77117
78118{-# INLINE _dropOne #-}
79119_dropOne :: Monad m => Int -> Int -> m ()
80120_dropOne size start =
81− S.drain
82− $ S.unfold (UF.drop 1 (UF.enumerateFromToIntegral (size + start))) start
121+ drainTransformationDefault size (UF.drop 1) start
83122
84123{-# INLINE dropAll #-}
85124dropAll :: Monad m => Int -> Int -> m ()
86125dropAll size start =
87− S.drain
88− $ S.unfold
89− (UF.drop (size + 1) (UF.enumerateFromToIntegral (size + start)))
90− start
126+ drainTransformationDefault size (UF.drop (size + 1)) start
91127
92128{-# INLINE dropWhileMTrue #-}
93129dropWhileMTrue :: Monad m => Int -> Int -> m ()
94130dropWhileMTrue size start =
95− S.drain
96− $ S.unfold
97− (UF.dropWhileM
98− (\_ -> return True)
99− (UF.enumerateFromToIntegral (size + start)))
100− start
131+ drainTransformationDefault
132+ size
133+ (UF.dropWhileM (\_ -> return True))
134+ start
101135
102136{-# INLINE dropWhileMFalse #-}
103137dropWhileMFalse :: Monad m => Int -> Int -> m ()
104138dropWhileMFalse size start =
105− S.drain
106− $ S.unfold
107− (UF.dropWhileM
108− (\_ -> return False)
109− (UF.enumerateFromToIntegral (size + start)))
110− start
139+ drainTransformationDefault
140+ size
141+ (UF.dropWhileM (\_ -> return False))
142+ start
111143
112144-------------------------------------------------------------------------------
113145-- Stream combination
@@ -116,13 +148,10 @@ dropWhileMFalse size start =
116148{-# INLINE zipWithM #-}
117149zipWithM :: Monad m => Int -> Int -> m ()
118150zipWithM size start =
119− S.drain
120− $ S.unfold
121− (UF.zipWithM
122− (\a b -> return $ a + b)
123− (UF.enumerateFromToIntegral (size + start))
124− (UF.enumerateFromToIntegral (size + start + 1)))
125− (start, start + 1)
151+ drainProductDefault
152+ size
153+ (UF.zipWithM (\a b -> return $ a + b))
154+ (start, start + 1)
126155
127156-------------------------------------------------------------------------------
128157-- Benchmarks