on master help sign in
Add a fold operation for unfolds
e1ba45aHarendra Kumar
src/Streamly/Internal/Data/Unfold.hs+28
@@ -55,6 +55,7 @@ module Streamly.Internal.Data.Unfold
55
55
56 Unfold
56 Unfold
57 , unfold
57 , unfold
58+ , fold
58
59
59 -- ** Unfolds
60 -- ** Unfolds
60 , singleton
61 , singleton
@@ -63,16 +64,21 @@ module Streamly.Internal.Data.Unfold
63 )
64 )
64where
65where
65
66
67+import GHC.Types (SPEC(..))
66import Streamly.Streams.StreamD.Type (Stream(..), Step(..))
68import Streamly.Streams.StreamD.Type (Stream(..), Step(..))
67#if __GLASGOW_HASKELL__ < 800
69#if __GLASGOW_HASKELL__ < 800
68import Streamly.Streams.StreamD.Type (pattern Stream)
70import Streamly.Streams.StreamD.Type (pattern Stream)
69#endif
71#endif
70import Streamly.Internal.Data.Unfold.Types (Unfold(..))
72import Streamly.Internal.Data.Unfold.Types (Unfold(..))
73+import Streamly.Internal.Data.Fold.Types (Fold(..))
71
74
72-------------------------------------------------------------------------------
75-------------------------------------------------------------------------------
73-- Running unfolds
76-- Running unfolds
74-------------------------------------------------------------------------------
77-------------------------------------------------------------------------------
75
78
79+-- XXX flip the first two arguments to make it similar to unfold i.e. returning
80+-- a -> Stream m b?
81+--
76-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
82-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
77--
83--
78{-# INLINE_NORMAL unfold #-}
84{-# INLINE_NORMAL unfold #-}
@@ -88,6 +94,28 @@ unfold seed (Unfold ustep inject) = Stream step Nothing
88 Skip s -> Skip (Just s)
94 Skip s -> Skip (Just s)
89 Stop -> Stop
95 Stop -> Stop
90
96
97+-- | Compose an 'Unfold' and a 'Fold'. Given an @Unfold m a b@ and a
98+-- @Fold m b c@, returns a monadic action @a -> m c@ representing the
99+-- application of the fold on the unfolded stream.
100+--
101+-- /Internal/
102+--
103+{-# INLINE_NORMAL fold #-}
104+fold :: Monad m => Unfold m a b -> Fold m b c -> a -> m c
105+fold (Unfold ustep inject) (Fold fstep initial extract) a =
106+ initial >>= \x -> inject a >>= go SPEC x
107+ where
108+ -- XXX !acc?
109+ {-# INLINE_LATE go #-}
110+ go !_ acc st = acc `seq` do
111+ r <- ustep st
112+ case r of
113+ Yield x s -> do
114+ acc' <- fstep acc x
115+ go SPEC acc' s
116+ Skip s -> go SPEC acc s
117+ Stop -> extract acc
118+
91-------------------------------------------------------------------------------
119-------------------------------------------------------------------------------
92-- Unfolds
120-- Unfolds
93-------------------------------------------------------------------------------
121-------------------------------------------------------------------------------