on master help sign in
implement covariant mapM/sequence for Fold

Also add some missing INLINE

commit3ce574aa533e5d428287d91f6cfebaf683e63bf3 authorHarendra Kumar <harendra.kumar@gmail.com> dateSun 2 Jun 2019, 05:17 (7 years ago) parentf767cf5
1 file changed+70−7
src/Streamly/Fold.hs+70−7
src/Streamly/Fold.hs+70−7
@@ -131,9 +131,13 @@ module Streamly.Fold
131131 -- 'left' assuming left side is the input side, notice that in @Fold m a b@
132132 -- the type variable @a@ is on the left side.
133133
134+ -- ** Covariant Operations
135+ , sequence
136+ , mapM
137+
134138 -- ** Mapping
135139 , lmap
136− --, sequence
140+ --, lsequence
137141 , lmapM
138142
139143 -- ** Scanning
@@ -368,10 +372,11 @@ import Prelude
368372 notElem, maximum, minimum, head, last, tail, length, null,
369373 reverse, iterate, init, and, or, lookup, foldr1, (!!),
370374 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,
371− span, splitAt, break)
375+ span, splitAt, break, mapM)
372376
373377import qualified Data.Heap as H
374378import qualified Data.Map.Strict as Map
379+import qualified Prelude
375380
376381import Streamly (MonadAsync, parallel)
377382import Streamly.Fold.Types (Fold(..))
@@ -453,6 +458,30 @@ foldl' (Fold step begin done) = P.foldlMx' step begin done
453458-- Transformations on fold inputs
454459------------------------------------------------------------------------------
455460
461+-- | Convert a fold that generates a monadic action to a fold that generates
462+-- the output of that monadic action.
463+--
464+-- @since 0.7.0
465+{-# INLINE sequence #-}
466+sequence :: Monad m => Fold m a (m b) -> Fold m a b
467+sequence (Fold step initial extract) = Fold step initial extract'
468+ where
469+ extract' x = do
470+ act <- extract x
471+ act >>= return
472+
473+-- | Convert a fold that generates a monadic action to a fold that generates
474+-- the output of that monadic action.
475+--
476+-- @since 0.7.0
477+{-# INLINE mapM #-}
478+mapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
479+mapM f = sequence . fmap f
480+
481+------------------------------------------------------------------------------
482+-- Transformations on fold inputs
483+------------------------------------------------------------------------------
484+
456485-- | @(lmap f fold)@ maps the function @f@ on the input of the fold.
457486--
458487-- >>> FL.foldl' (FL.lmap (\x -> x * x) FL.sum) (S.enumerateFromTo 1 100)
@@ -737,7 +766,7 @@ toListRev = Fold (\xs x -> return $ x:xs) (return []) return
737766-- be very inefficient, consider using "Streamly.Array" instead.
738767--
739768-- @since 0.7.0
740−{-# INLINABLE toStream #-}
769+{-# INLINE toStream #-}
741770toStream :: Monad m => Fold m a (SerialT Identity a)
742771toStream = Fold (\f x -> return $ f . (x `K.cons`))
743772 (return id)
@@ -1115,6 +1144,7 @@ chunksOf n f m = D.fromStreamD $ D.groupsOf n f (D.toStreamD m)
11151144
11161145-- | Transform a fold from a pure input to a 'Maybe' input, consuming only
11171146-- 'Just' values.
1147+{-# INLINE lcatMaybes #-}
11181148lcatMaybes :: Monad m => Fold m a b -> Fold m (Maybe a) b
11191149lcatMaybes = lfilter isJust . lmap fromJust
11201150
@@ -1284,6 +1314,7 @@ _spanRollingBy cmp (Fold stepL initialL extractL) (Fold stepR initialR extractR)
12841314-- > [[1,3,7],[0,2,5]]
12851315--
12861316-- @since 0.7.0
1317+{-# INLINE groupsBy #-}
12871318groupsBy
12881319 :: (IsStream t, Monad m)
12891320 => (a -> a -> Bool)
@@ -1304,6 +1335,7 @@ groupsBy cmp f m = D.fromStreamD $ D.groupsBy cmp f (D.toStreamD m)
13041335-- > [[1,2,3],[7,8,9]]
13051336--
13061337-- @since 0.7.0
1338+{-# INLINE groupsRollingBy #-}
13071339groupsRollingBy
13081340 :: (IsStream t, Monad m)
13091341 => (a -> a -> Bool)
@@ -1955,6 +1987,7 @@ partition = partitionBy id
19551987-- @
19561988--
19571989-- @since 0.7.0
1990+{-# INLINE demuxWith #-}
19581991demuxWith :: (Monad m, Ord k)
19591992 => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)
19601993demuxWith f kv = Fold step initial extract
@@ -1974,7 +2007,7 @@ demuxWith f kv = Fold step initial extract
19742007 Just (Fold step' acc extract') -> do
19752008 !r <- acc >>= \x -> step' x a
19762009 return $ Map.insert k (Fold step' (return r) extract') mp
1977− extract = mapM (\(Fold _ acc e) -> acc >>= e)
2010+ extract = Prelude.mapM (\(Fold _ acc e) -> acc >>= e)
19782011
19792012-- | Fold a stream of key value pairs using a map of specific folds for each
19802013-- key into a map from keys to the results of fold outputs of the corresponding
@@ -1989,6 +2022,7 @@ demuxWith f kv = Fold step initial extract
19892022-- @
19902023--
19912024-- @since 0.7.0
2025+{-# INLINE demux #-}
19922026demux :: (Monad m, Ord k)
19932027 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
19942028demux fs = demuxWith fst (Map.map (lmap snd) fs)
@@ -2012,13 +2046,16 @@ demux fs = demuxWith fst (Map.map (lmap snd) fs)
20122046-- demuxWith_ can be slightly faster than demuxWith because we do not need to
20132047-- update the Map in this case. This may be significant only if the map is
20142048-- large.
2049+{-# INLINE demuxWith_ #-}
20152050demuxWith_ :: (Monad m, Ord k)
20162051 => (a -> k) -> Map k (Fold m a b) -> Fold m a ()
20172052demuxWith_ f kv = Fold step initial extract
20182053
20192054 where
20202055
2021− initial = mapM (\(Fold s i e) -> i >>= \r -> return (Fold s (return r) e)) kv
2056+ initial = do
2057+ Prelude.mapM (\(Fold s i e) ->
2058+ i >>= \r -> return (Fold s (return r) e)) kv
20222059 step mp a =
20232060 -- XXX should we raise an exception in Nothing case?
20242061 -- Ideally we should enforce that it is a total map over k so that look
@@ -2028,7 +2065,7 @@ demuxWith_ f kv = Fold step initial extract
20282065 Just (Fold step' acc _) -> do
20292066 _ <- acc >>= \x -> step' x a
20302067 return mp
2031− extract mp = mapM (\(Fold _ acc e) -> acc >>= e) mp >> return ()
2068+ extract mp = Prelude.mapM (\(Fold _ acc e) -> acc >>= e) mp >> return ()
20322069
20332070-- | Given a stream of key value pairs and a map from keys to folds, fold the
20342071-- values for each key using the corresponding folds, discarding the outputs.
@@ -2043,6 +2080,7 @@ demuxWith_ f kv = Fold step initial extract
20432080-- @
20442081--
20452082-- @since 0.7.0
2083+{-# INLINE demux_ #-}
20462084demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
20472085demux_ fs = demuxWith_ fst (Map.map (lmap snd) fs)
20482086
@@ -2061,6 +2099,7 @@ demux_ fs = demuxWith_ fst (Map.map (lmap snd) fs)
20612099-- @
20622100--
20632101-- @since 0.7.0
2102+{-# INLINE classifyWith #-}
20642103classifyWith :: (Monad m, Ord k) => (a -> k) -> Fold m a b -> Fold m a (Map k b)
20652104classifyWith f (Fold step initial extract) = Fold step' initial' extract'
20662105
@@ -2077,7 +2116,7 @@ classifyWith f (Fold step initial extract) = Fold step' initial' extract'
20772116 Just x -> do
20782117 r <- step x a
20792118 return $ Map.insert k r kv
2080− extract' = mapM extract
2119+ extract' = Prelude.mapM extract
20812120
20822121-- | Given an input stream of key value pairs and a fold for values, fold all
20832122-- the values belonging to each key. Useful for map/reduce, bucketizing the
@@ -2094,6 +2133,7 @@ classifyWith f (Fold step initial extract) = Fold step' initial' extract'
20942133-- > classify fld = classifyWith fst (lmap snd fld)
20952134--
20962135-- @since 0.7.0
2136+{-# INLINE classify #-}
20972137classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
20982138classify fld = classifyWith fst (lmap snd fld)
20992139
@@ -2141,6 +2181,7 @@ unzipWith f = unzipWithM (return . f)
21412181-- This is the consumer side dual of the producer side 'zip' operation.
21422182--
21432183-- @since 0.7.0
2184+{-# INLINE unzip #-}
21442185unzip :: Monad m => Fold m a x -> Fold m b y -> Fold m (a,b) (x,y)
21452186unzip = unzipWith id
21462187
@@ -2179,6 +2220,26 @@ lconcatMap s f1 f2 = undefined
21792220-- All the grouping transformation that we apply to a stream can also be
21802221-- applied to a fold input stream.
21812222
2223+{-
2224+-- | Group the input stream into groups of elements between @low@ and @high@.
2225+-- Collection starts in chunks of @low@ and then keeps doubling until we reach
2226+-- @high@. Each chunk is folded using the provided fold function.
2227+--
2228+-- This could be useful, for example, when we are folding a stream of unknown
2229+-- size to a stream of arrays and we want to minimize the number of
2230+-- allocations.
2231+--
2232+-- @
2233+--
2234+-- XXX we should be able to implement it with parsers/terminating folds.
2235+--
2236+{-# INLINE lchunksInRange #-}
2237+lchunksInRange :: Monad m
2238+ => Int -> Int -> Fold m a b -> Fold m b c -> Fold m a c
2239+lchunksInRange low high (Fold step1 initial1 extract1)
2240+ (Fold step2 initial2 extract2) = undefined
2241+-}
2242+
21822243-- | Group the input stream into groups of @n@ elements each and then fold each
21832244-- group using the provided fold function.
21842245--
@@ -2188,6 +2249,7 @@ lconcatMap s f1 f2 = undefined
21882249--
21892250-- @
21902251--
2252+{-# INLINE lchunksOf #-}
21912253lchunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c
21922254lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =
21932255 Fold step' initial' extract'
@@ -2219,6 +2281,7 @@ lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =
22192281-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c
22202282--
22212283-- @
2284+{-# INLINE lsessionsOf #-}
22222285lsessionsOf :: MonadAsync m => Double -> Fold m a b -> Fold m b c -> Fold m a c
22232286lsessionsOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =
22242287 Fold step' initial' extract'