5bb0a49 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add concatMap, concatMapM, generateM, generate, findM to Prelude.

5bb0a49Pranay Sashank8 years ago

src/Streamly/Prelude.hs+33−1

compare with master

unifiedsplit

@@ -91,6 +91,8 @@ module Streamly.Prelude

91 -- ** Generate From

92 -- | Convert an input structure, container or source into a stream. All of

93 -- these can be expressed in terms of primitives.

94+ , generate

95+ , generateM

94 , P.fromList

96 , P.fromList

95 , fromListM

97 , fromListM

96 , K.fromFoldable

98 , K.fromFoldable

@@ -137,6 +139,7 @@ module Streamly.Prelude

137 , elemIndex

139 , elemIndex

138 , notElem

140 , notElem

139 , lookup

141 , lookup

142+ , findM

140 , find

143 , find

141 , findIndex

144 , findIndex

142 , all

145 , all

@@ -146,6 +149,8 @@ module Streamly.Prelude

146 , eqBy

149 , eqBy

147 , cmpBy

150 , cmpBy

148 , (!!)

151 , (!!)

152+ , concatMapM

153+ , concatMap

149

154

150 -- Full folds - need to go through all elements

155 -- Full folds - need to go through all elements

151 , length

156 , length

@@ -258,7 +263,7 @@ import Prelude

258 foldl, mapM, mapM_, sequence, all, any, sum, product, elem,

263 foldl, mapM, mapM_, sequence, all, any, sum, product, elem,

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

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

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

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

261− splitAt, scanl, scanl1, replicate)

266+ splitAt, scanl, scanl1, replicate, concatMap)

262import qualified Prelude

267import qualified Prelude

263import qualified System.IO as IO

268import qualified System.IO as IO

264

269

@@ -403,6 +408,19 @@ yield = K.yield

403yieldM :: (Monad m, IsStream t) => m a -> t m a

408yieldM :: (Monad m, IsStream t) => m a -> t m a

404yieldM = K.yieldM

409yieldM = K.yieldM

405

410

411+-- | Generate a stream of given length, from a given pure

412+-- function of its indices.

413+{-# INLINE generate #-}

414+generate :: (IsStream t, Monad m) => Int -> (Int -> a) -> t m a

415+generate n = fromStreamD . D.generate n

416+

417+-- | Generate a stream of given length, from a given monadic

418+-- function of its indices.

419+--

420+{-# INLINE generateM #-}

421+generateM :: (IsStream t, Monad m) => Int -> (Int -> m a) -> t m a

422+generateM n = fromStreamD . D.generateM n

423+

406-- | Generate a stream by performing a monadic action @n@ times. Can be

424-- | Generate a stream by performing a monadic action @n@ times. Can be

407-- expressed as @stimes n (yieldM m)@.

425-- expressed as @stimes n (yieldM m)@.

408--

426--

@@ -769,6 +787,12 @@ lookup a m = S.lookup a (toStreamS m)

769find :: Monad m => (a -> Bool) -> SerialT m a -> m (Maybe a)

787find :: Monad m => (a -> Bool) -> SerialT m a -> m (Maybe a)

770find p m = S.find p (toStreamS m)

788find p m = S.find p (toStreamS m)

771

789

790+-- | Like 'find' but with a monadic predicate.

791+--

792+{-# INLINE findM #-}

793+findM :: Monad m => (a -> m Bool) -> SerialT m a -> m (Maybe a)

794+findM p m = S.findM p (toStreamS m)

795+

772-- | Find all the indices where the element in the stream satisfies the given

796-- | Find all the indices where the element in the stream satisfies the given

773-- predicate.

797-- predicate.

774--

798--

@@ -800,6 +824,14 @@ elemIndices a = findIndices (==a)

800elemIndex :: (Monad m, Eq a) => a -> SerialT m a -> m (Maybe Int)

824elemIndex :: (Monad m, Eq a) => a -> SerialT m a -> m (Maybe Int)

801elemIndex a = findIndex (==a)

825elemIndex a = findIndex (==a)

802

826

827+{-# INLINE concatMap #-}

828+concatMap :: Monad m => (a -> SerialT m b) -> SerialT m a -> SerialT m b

829+concatMap f m = fromStreamD $ D.concatMap (toStreamD . f) (toStreamD m)

830+

831+{-# INLINE concatMapM #-}

832+concatMapM :: Monad m => (a -> m (SerialT m b)) -> SerialT m a -> SerialT m b

833+concatMapM f m = fromStreamD $ D.concatMapM (fmap toStreamD . f) (toStreamD m)

834+

803------------------------------------------------------------------------------

835------------------------------------------------------------------------------

804-- Substreams

836-- Substreams

805------------------------------------------------------------------------------

837------------------------------------------------------------------------------

src/Streamly/Streams/StreamD.hs+58−7

@@ -58,6 +58,8 @@ module Streamly.Streams.StreamD

58 -- | Direct style stream does not support @fromFoldable@.

59 , yield

60 , yieldM

61+ , generate

62+ , generateM

61 , fromList

63 , fromList

62 , fromListM

64 , fromListM

63 , fromStreamK

65 , fromStreamK

@@ -87,8 +89,11 @@ module Streamly.Streams.StreamD

87 , minimumBy

89 , minimumBy

88 , findIndices

90 , findIndices

89 , lookup

91 , lookup

92+ , findM

90 , find

93 , find

91 , (!!)

94 , (!!)

95+ , concatMapM

96+ , concatMap

93 -- ** Substreams

98 -- ** Substreams

94 , isPrefixOf

99 , isPrefixOf

@@ -162,7 +167,7 @@ import Prelude

162 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,

167 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,

163 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,

168 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,

164 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,

169 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,

165− (!!), scanl, scanl1)

170+ (!!), scanl, scanl1, concatMap)

166

171

167import Streamly.SVar (MonadAsync, State(..), defState, rstState)

172import Streamly.SVar (MonadAsync, State(..), defState, rstState)

168

173

@@ -285,6 +290,20 @@ yieldM m = Stream step True

285 step _ True = m >>= \x -> return $ Yield x False

290 step _ True = m >>= \x -> return $ Yield x False

286 step _ False = return Stop

291 step _ False = return Stop

287

292

293+{-# INLINE generate #-}

294+generate :: Monad m => Int -> (Int -> a) -> Stream m a

295+generate n gen = generateM n (return . gen)

296+

297+{-# INLINE_NORMAL generateM #-}

298+generateM :: Monad m => Int -> (Int -> m a) -> Stream m a

299+generateM n gen = n `seq` Stream step 0

300+ where

301+ {-# INLINE_LATE step #-}

302+ step _ i | i < n = do

303+ x <- gen i

304+ return $ Yield x (i+1)

305+ | otherwise = return Stop

306+

288-- XXX we need the MonadAsync constraint because of a rewrite rule.

307-- XXX we need the MonadAsync constraint because of a rewrite rule.

289-- | Convert a list of monadic actions to a 'Stream'

308-- | Convert a list of monadic actions to a 'Stream'

290{-# INLINE_LATE fromListM #-}

309{-# INLINE_LATE fromListM #-}

@@ -562,17 +581,23 @@ lookup e (Stream step state) = go state

562 Skip s -> go s

581 Skip s -> go s

563 Stop -> return Nothing

582 Stop -> return Nothing

564

583

565−{-# INLINE_NORMAL find #-}

584+{-# INLINE_NORMAL findM #-}

566−find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)

585+findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)

567−find p (Stream step state) = go state

586+findM p (Stream step state) = go SPEC state

568 where

587 where

569− go st = do

588+ go !_ st = do

570 r <- step defState st

589 r <- step defState st

571 case r of

590 case r of

591+ Yield x s -> do

592+ b <- p x

572− Yield x s -> if p x then return (Just x) else go s

593+ if b then return (Just x) else go SPEC s

573− Skip s -> go s

594+ Skip s -> go SPEC s

574 Stop -> return Nothing

595 Stop -> return Nothing

575

596

597+{-# INLINE find #-}

598+find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)

599+find p = findM (return . p)

600+

576{-# INLINE_NORMAL findIndices #-}

601{-# INLINE_NORMAL findIndices #-}

577findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int

602findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int

578findIndices p (Stream step state) = Stream step' (state, 0)

603findIndices p (Stream step state) = Stream step' (state, 0)

@@ -585,6 +610,32 @@ findIndices p (Stream step state) = Stream step' (state, 0)

585 Skip s -> Skip (s, i+1)

610 Skip s -> Skip (s, i+1)

586 Stop -> Stop

611 Stop -> Stop

587

612

613+{-# INLINE concatMap #-}

614+concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b

615+concatMap f = concatMapM (return . f)

616+

617+{-# INLINE_NORMAL concatMapM #-}

618+concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b

619+concatMapM f (Stream step state) = Stream step' (Left state)

620+ where

621+ step' gst (Left st) = do

622+ r <- step (rstState gst) st

623+ case r of

624+ Yield a s -> do

625+ b_stream <- f a

626+ return $ Skip (Right (b_stream, s))

627+ Skip s -> return $ Skip (Left s)

628+ Stop -> return Stop

629+

630+ step' _ (Right (Stream inner_step inner_st, st)) = do

631+ r <- inner_step defState inner_st

632+ case r of

633+ Yield b inner_s ->

634+ return $ Yield b (Right (Stream inner_step inner_s, st))

635+ Skip inner_s ->

636+ return $ Skip (Right (Stream inner_step inner_s, st))

637+ Stop -> return $ Skip (Left st)

638+

588------------------------------------------------------------------------------

639------------------------------------------------------------------------------

589-- Substreams

640-- Substreams

590------------------------------------------------------------------------------

641------------------------------------------------------------------------------

src/Streamly/Streams/StreamK.hs+14−7

@@ -95,6 +95,7 @@ module Streamly.Streams.StreamK

95 , maximumBy

96 , findIndices

97 , lookup

98+ , findM

98 , find

99 , find

99 , (!!)

100 , (!!)

100

101

@@ -802,17 +803,23 @@ lookup e m = go (toStream m)

802 | otherwise = go x

803 | otherwise = go x

803 in unStream m1 defState (return Nothing) single yieldk

804 in unStream m1 defState (return Nothing) single yieldk

804

805

805−{-# INLINE find #-}

806+{-# INLINE findM #-}

806−find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)

807+findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)

807−find p m = go (toStream m)

808+findM p m = go (toStream m)

808 where

809 where

809 go m1 =

810 go m1 =

810− let single a | p a = return $ Just a

811+ let single a = do

812+ b <- p a

811− | otherwise = return Nothing

813+ if b then return $ Just a else return Nothing

812− yieldk a x | p a = return $ Just a

814+ yieldk a x = do

813− | otherwise = go x

815+ b <- p a

816+ if b then return $ Just a else go x

814 in unStream m1 defState (return Nothing) single yieldk

817 in unStream m1 defState (return Nothing) single yieldk

815

818

819+{-# INLINE find #-}

820+find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)

821+find p = findM (return . p)

822+

816{-# INLINE findIndices #-}

823{-# INLINE findIndices #-}

817findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int

824findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int

818findIndices p = fromStream . go 0 . toStream

825findIndices p = fromStream . go 0 . toStream

test/Prop.hs+3

@@ -404,6 +404,9 @@ transformOps constr desc eq t = do

404 prop (desc <> " deleteBy (==) 4") $

405 transform (deleteBy (==) 4) $ t . S.deleteBy (==) 4

406

407+ prop (desc <> " concatMap") $ transform (concatMap (const [1..10])) $

408+ t . (\g -> adapt (S.concatMap (const (S.fromList [1..10])) (adapt g)))

409+

407concurrentOps

410concurrentOps

408 :: IsStream t

411 :: IsStream t

409 => ([Word8] -> t IO Word8)

412 => ([Word8] -> t IO Word8)