3ce574a · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

implement covariant mapM/sequence for Fold

3ce574aHarendra Kumar7 years ago

src/Streamly/Fold.hs+70−7

compare with master

unifiedsplit

@@ -131,9 +131,13 @@ module Streamly.Fold

131 -- 'left' assuming left side is the input side, notice that in @Fold m a b@

132 -- the type variable @a@ is on the left side.

133

134+ -- ** Covariant Operations

135+ , sequence

136+ , mapM

137+

134 -- ** Mapping

138 -- ** Mapping

135 , lmap

139 , lmap

136− --, sequence

140+ --, lsequence

137 , lmapM

141 , lmapM

138

142

139 -- ** Scanning

143 -- ** Scanning

@@ -368,10 +372,11 @@ import Prelude

368 notElem, maximum, minimum, head, last, tail, length, null,

372 notElem, maximum, minimum, head, last, tail, length, null,

369 reverse, iterate, init, and, or, lookup, foldr1, (!!),

373 reverse, iterate, init, and, or, lookup, foldr1, (!!),

370 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,

374 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,

371− span, splitAt, break)

375+ span, splitAt, break, mapM)

372

376

373import qualified Data.Heap as H

377import qualified Data.Heap as H

374import qualified Data.Map.Strict as Map

378import qualified Data.Map.Strict as Map

379+import qualified Prelude

375

380

376import Streamly (MonadAsync, parallel)

381import Streamly (MonadAsync, parallel)

377import Streamly.Fold.Types (Fold(..))

382import Streamly.Fold.Types (Fold(..))

@@ -453,6 +458,30 @@ foldl' (Fold step begin done) = P.foldlMx' step begin done

453-- Transformations on fold inputs

458-- Transformations on fold inputs

454------------------------------------------------------------------------------

459------------------------------------------------------------------------------

455

460

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+

456-- | @(lmap f fold)@ maps the function @f@ on the input of the fold.

485-- | @(lmap f fold)@ maps the function @f@ on the input of the fold.

457--

486--

458-- >>> FL.foldl' (FL.lmap (\x -> x * x) FL.sum) (S.enumerateFromTo 1 100)

487-- >>> 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

737-- be very inefficient, consider using "Streamly.Array" instead.

766-- be very inefficient, consider using "Streamly.Array" instead.

738--

767--

739-- @since 0.7.0

768-- @since 0.7.0

740−{-# INLINABLE toStream #-}

769+{-# INLINE toStream #-}

741toStream :: Monad m => Fold m a (SerialT Identity a)

770toStream :: Monad m => Fold m a (SerialT Identity a)

742toStream = Fold (\f x -> return $ f . (x `K.cons`))

771toStream = Fold (\f x -> return $ f . (x `K.cons`))

743 (return id)

772 (return id)

@@ -1115,6 +1144,7 @@ chunksOf n f m = D.fromStreamD $ D.groupsOf n f (D.toStreamD m)

1115

1144

1116-- | Transform a fold from a pure input to a 'Maybe' input, consuming only

1145-- | Transform a fold from a pure input to a 'Maybe' input, consuming only

1117-- 'Just' values.

1146-- 'Just' values.

1147+{-# INLINE lcatMaybes #-}

1118lcatMaybes :: Monad m => Fold m a b -> Fold m (Maybe a) b

1148lcatMaybes :: Monad m => Fold m a b -> Fold m (Maybe a) b

1119lcatMaybes = lfilter isJust . lmap fromJust

1149lcatMaybes = lfilter isJust . lmap fromJust

1120

1150

@@ -1284,6 +1314,7 @@ _spanRollingBy cmp (Fold stepL initialL extractL) (Fold stepR initialR extractR)

1284-- > [[1,3,7],[0,2,5]]

1314-- > [[1,3,7],[0,2,5]]

1285--

1315--

1286-- @since 0.7.0

1316-- @since 0.7.0

1317+{-# INLINE groupsBy #-}

1287groupsBy

1318groupsBy

1288 :: (IsStream t, Monad m)

1319 :: (IsStream t, Monad m)

1289 => (a -> a -> Bool)

1320 => (a -> a -> Bool)

@@ -1304,6 +1335,7 @@ groupsBy cmp f m = D.fromStreamD $ D.groupsBy cmp f (D.toStreamD m)

1304-- > [[1,2,3],[7,8,9]]

1335-- > [[1,2,3],[7,8,9]]

1305--

1336--

1306-- @since 0.7.0

1337-- @since 0.7.0

1338+{-# INLINE groupsRollingBy #-}

1307groupsRollingBy

1339groupsRollingBy

1308 :: (IsStream t, Monad m)

1340 :: (IsStream t, Monad m)

1309 => (a -> a -> Bool)

1341 => (a -> a -> Bool)

@@ -1955,6 +1987,7 @@ partition = partitionBy id

1955-- @

1987-- @

1956--

1988--

1957-- @since 0.7.0

1989-- @since 0.7.0

1990+{-# INLINE demuxWith #-}

1958demuxWith :: (Monad m, Ord k)

1991demuxWith :: (Monad m, Ord k)

1959 => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)

1992 => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)

1960demuxWith f kv = Fold step initial extract

1993demuxWith f kv = Fold step initial extract

@@ -1974,7 +2007,7 @@ demuxWith f kv = Fold step initial extract

1974 Just (Fold step' acc extract') -> do

2007 Just (Fold step' acc extract') -> do

1975 !r <- acc >>= \x -> step' x a

2008 !r <- acc >>= \x -> step' x a

1976 return $ Map.insert k (Fold step' (return r) extract') mp

2009 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)

1978

2011

1979-- | Fold a stream of key value pairs using a map of specific folds for each

2012-- | Fold a stream of key value pairs using a map of specific folds for each

1980-- key into a map from keys to the results of fold outputs of the corresponding

2013-- 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

1989-- @

2022-- @

1990--

2023--

1991-- @since 0.7.0

2024-- @since 0.7.0

2025+{-# INLINE demux #-}

1992demux :: (Monad m, Ord k)

2026demux :: (Monad m, Ord k)

1993 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)

2027 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)

1994demux fs = demuxWith fst (Map.map (lmap snd) fs)

2028demux fs = demuxWith fst (Map.map (lmap snd) fs)

@@ -2012,13 +2046,16 @@ demux fs = demuxWith fst (Map.map (lmap snd) fs)

2012-- demuxWith_ can be slightly faster than demuxWith because we do not need to

2046-- demuxWith_ can be slightly faster than demuxWith because we do not need to

2013-- update the Map in this case. This may be significant only if the map is

2047-- update the Map in this case. This may be significant only if the map is

2014-- large.

2048-- large.

2049+{-# INLINE demuxWith_ #-}

2015demuxWith_ :: (Monad m, Ord k)

2050demuxWith_ :: (Monad m, Ord k)

2016 => (a -> k) -> Map k (Fold m a b) -> Fold m a ()

2051 => (a -> k) -> Map k (Fold m a b) -> Fold m a ()

2017demuxWith_ f kv = Fold step initial extract

2052demuxWith_ f kv = Fold step initial extract

2018

2053

2019 where

2054 where

2020

2055

2056+ initial = do

2057+ Prelude.mapM (\(Fold s i e) ->

2021− initial = mapM (\(Fold s i e) -> i >>= \r -> return (Fold s (return r) e)) kv

2058+ i >>= \r -> return (Fold s (return r) e)) kv

2022 step mp a =

2059 step mp a =

2023 -- XXX should we raise an exception in Nothing case?

2060 -- XXX should we raise an exception in Nothing case?

2024 -- Ideally we should enforce that it is a total map over k so that look

2061 -- 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

2028 Just (Fold step' acc _) -> do

2065 Just (Fold step' acc _) -> do

2029 _ <- acc >>= \x -> step' x a

2066 _ <- acc >>= \x -> step' x a

2030 return mp

2067 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 ()

2032

2069

2033-- | Given a stream of key value pairs and a map from keys to folds, fold the

2070-- | Given a stream of key value pairs and a map from keys to folds, fold the

2034-- values for each key using the corresponding folds, discarding the outputs.

2071-- values for each key using the corresponding folds, discarding the outputs.

@@ -2043,6 +2080,7 @@ demuxWith_ f kv = Fold step initial extract

2043-- @

2080-- @

2044--

2081--

2045-- @since 0.7.0

2082-- @since 0.7.0

2083+{-# INLINE demux_ #-}

2046demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()

2084demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()

2047demux_ fs = demuxWith_ fst (Map.map (lmap snd) fs)

2085demux_ fs = demuxWith_ fst (Map.map (lmap snd) fs)

2048

2086

@@ -2061,6 +2099,7 @@ demux_ fs = demuxWith_ fst (Map.map (lmap snd) fs)

2061-- @

2099-- @

2062--

2100--

2063-- @since 0.7.0

2101-- @since 0.7.0

2102+{-# INLINE classifyWith #-}

2064classifyWith :: (Monad m, Ord k) => (a -> k) -> Fold m a b -> Fold m a (Map k b)

2103classifyWith :: (Monad m, Ord k) => (a -> k) -> Fold m a b -> Fold m a (Map k b)

2065classifyWith f (Fold step initial extract) = Fold step' initial' extract'

2104classifyWith f (Fold step initial extract) = Fold step' initial' extract'

2066

2105

@@ -2077,7 +2116,7 @@ classifyWith f (Fold step initial extract) = Fold step' initial' extract'

2077 Just x -> do

2116 Just x -> do

2078 r <- step x a

2117 r <- step x a

2079 return $ Map.insert k r kv

2118 return $ Map.insert k r kv

2080− extract' = mapM extract

2119+ extract' = Prelude.mapM extract

2081

2120

2082-- | Given an input stream of key value pairs and a fold for values, fold all

2121-- | Given an input stream of key value pairs and a fold for values, fold all

2083-- the values belonging to each key. Useful for map/reduce, bucketizing the

2122-- 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'

2094-- > classify fld = classifyWith fst (lmap snd fld)

2133-- > classify fld = classifyWith fst (lmap snd fld)

2095--

2134--

2096-- @since 0.7.0

2135-- @since 0.7.0

2136+{-# INLINE classify #-}

2097classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)

2137classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)

2098classify fld = classifyWith fst (lmap snd fld)

2138classify fld = classifyWith fst (lmap snd fld)

2099

2139

@@ -2141,6 +2181,7 @@ unzipWith f = unzipWithM (return . f)

2141-- This is the consumer side dual of the producer side 'zip' operation.

2181-- This is the consumer side dual of the producer side 'zip' operation.

2142--

2182--

2143-- @since 0.7.0

2183-- @since 0.7.0

2184+{-# INLINE unzip #-}

2144unzip :: Monad m => Fold m a x -> Fold m b y -> Fold m (a,b) (x,y)

2185unzip :: Monad m => Fold m a x -> Fold m b y -> Fold m (a,b) (x,y)

2145unzip = unzipWith id

2186unzip = unzipWith id

2146

2187

@@ -2179,6 +2220,26 @@ lconcatMap s f1 f2 = undefined

2179-- All the grouping transformation that we apply to a stream can also be

2220-- All the grouping transformation that we apply to a stream can also be

2180-- applied to a fold input stream.

2221-- applied to a fold input stream.

2181

2222

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+

2182-- | Group the input stream into groups of @n@ elements each and then fold each

2243-- | Group the input stream into groups of @n@ elements each and then fold each

2183-- group using the provided fold function.

2244-- group using the provided fold function.

2184--

2245--

@@ -2188,6 +2249,7 @@ lconcatMap s f1 f2 = undefined

2188--

2249--

2189-- @

2250-- @

2190--

2251--

2252+{-# INLINE lchunksOf #-}

2191lchunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c

2253lchunksOf :: Monad m => Int -> Fold m a b -> Fold m b c -> Fold m a c

2192lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

2254lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

2193 Fold step' initial' extract'

2255 Fold step' initial' extract'

@@ -2219,6 +2281,7 @@ lchunksOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

2219-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

2281-- -----Fold m a b----|-Fold n a c-|-Fold n a c-|-...-|----Fold m a c

2220--

2282--

2221-- @

2283-- @

2284+{-# INLINE lsessionsOf #-}

2222lsessionsOf :: MonadAsync m => Double -> Fold m a b -> Fold m b c -> Fold m a c

2285lsessionsOf :: MonadAsync m => Double -> Fold m a b -> Fold m b c -> Fold m a c

2223lsessionsOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

2286lsessionsOf n (Fold step1 initial1 extract1) (Fold step2 initial2 extract2) =

2224 Fold step' initial' extract'

2287 Fold step' initial' extract'