@@ -67,6 +67,7 @@ module Streamly.Streams.StreamD
6767 -- ** General Folds
6868 , foldr
6969 , foldrM
70+ , foldr1
7071 , foldl'
7172 , foldlM'
7273
@@ -82,6 +83,10 @@ module Streamly.Streams.StreamD
8283 , any
8384 , maximum
8485 , minimum
86+ , findIndices
87+ , lookup
88+ , find
89+ , (!!)
8590
8691 -- ** Map and Fold
8792 , mapM_
@@ -110,6 +115,10 @@ module Streamly.Streams.StreamD
110115 -- * Mapping
111116 , map
112117 , mapM
118+ , sequence
119+
120+ -- * Inserting
121+ , insertBy
113122
114123 -- ** Map and Filter
115124 , mapMaybe
@@ -122,6 +131,12 @@ module Streamly.Streams.StreamD
122131 -- * Comparisions
123132 , eqBy
124133 , cmpBy
134+
135+ -- * Merging
136+ , mergeBy
137+
138+ -- * Transformation comprehensions
139+ , the
125140 )
126141where
127142
@@ -130,7 +145,8 @@ import GHC.Types ( SPEC(..) )
130145import Prelude
131146 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
132147 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
133− notElem, null, head, tail, zipWith)
148+ notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
149+ (!!))
134150
135151import Streamly.SVar (MonadAsync, State(..), defState, rstState)
136152
@@ -306,6 +322,14 @@ foldrM f z (Stream step state) = go SPEC state
306322foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
307323foldr f = foldrM (\a b -> return (f a b))
308324
325+{-# INLINE_NORMAL foldr1 #-}
326+foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
327+foldr1 f m = do
328+ r <- uncons m
329+ case r of
330+ Nothing -> return Nothing
331+ Just (h, t) -> fmap Just (foldr f h t)
332+
309333{-# INLINE_NORMAL foldlM' #-}
310334foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
311335foldlM' fstep begin (Stream step state) = go SPEC begin state
@@ -460,6 +484,53 @@ minimum (Stream step state) = go Nothing state
460484 Skip s -> go (Just acc) s
461485 Stop -> return (Just acc)
462486
487+{-# INLINE_NORMAL (!!) #-}
488+(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
489+(Stream step state) !! i = go i state
490+ where
491+ go n st = do
492+ r <- step defState st
493+ case r of
494+ Yield x s | n < 0 -> return Nothing
495+ | n == 0 -> return $ Just x
496+ | otherwise -> go (n - 1) s
497+ Skip s -> go n s
498+ Stop -> return Nothing
499+
500+{-# INLINE_NORMAL lookup #-}
501+lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
502+lookup e (Stream step state) = go state
503+ where
504+ go st = do
505+ r <- step defState st
506+ case r of
507+ Yield (a, b) s -> if e == a then return (Just b) else go s
508+ Skip s -> go s
509+ Stop -> return Nothing
510+
511+{-# INLINE_NORMAL find #-}
512+find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
513+find p (Stream step state) = go state
514+ where
515+ go st = do
516+ r <- step defState st
517+ case r of
518+ Yield x s -> if p x then return (Just x) else go s
519+ Skip s -> go s
520+ Stop -> return Nothing
521+
522+{-# INLINE_NORMAL findIndices #-}
523+findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
524+findIndices p (Stream step state) = Stream step' (state, 0)
525+ where
526+ {-# INLINE_LATE step' #-}
527+ step' gst (st, i) = do
528+ r <- step (rstState gst) st
529+ return $ case r of
530+ Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
531+ Skip s -> Skip (s, i+1)
532+ Stop -> Stop
533+
463534------------------------------------------------------------------------------
464535-- Map and Fold
465536------------------------------------------------------------------------------
@@ -660,6 +731,41 @@ mapM f (Stream step state) = Stream step' state
660731map :: Monad m => (a -> b) -> Stream m a -> Stream m b
661732map f = mapM (return . f)
662733
734+{-# INLINE_NORMAL sequence #-}
735+sequence :: Monad m => Stream m (m a) -> Stream m a
736+sequence (Stream step state) = Stream step' state
737+ where
738+ {-# INLINE_LATE step' #-}
739+ step' gst st = do
740+ r <- step (rstState gst) st
741+ case r of
742+ Yield x s -> x >>= \a -> return (Yield a s)
743+ Skip s -> return $ Skip s
744+ Stop -> return Stop
745+
746+------------------------------------------------------------------------------
747+-- Inserting
748+------------------------------------------------------------------------------
749+
750+{-# INLINE_NORMAL insertBy #-}
751+insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
752+insertBy cmp a (Stream step state) = Stream step' (state, False)
753+ where
754+ step' gst (st, False) = do
755+ r <- step (rstState gst) st
756+ case r of
757+ Yield x s -> case cmp a x of
758+ GT -> return $ Yield x (s, False)
759+ _ -> return $ Yield a (st, True)
760+ Skip s -> return $ Skip (s, False)
761+ Stop -> return $ Yield a (st, True)
762+ step' gst (st, True) = do
763+ r <- step (rstState gst) st
764+ case r of
765+ Yield x s -> return $ Yield x (s, True)
766+ Skip s -> return $ Skip (s, True)
767+ Stop -> return Stop
768+
663769------------------------------------------------------------------------------
664770-- Transformation by Map and Filter
665771------------------------------------------------------------------------------
@@ -674,7 +780,7 @@ mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
674780mapMaybeM f = fmap fromJust . filter isJust . mapM f
675781
676782------------------------------------------------------------------------------
677−-- Instances
783+-- Zipping
678784------------------------------------------------------------------------------
679785
680786{-# INLINE_NORMAL zipWithM #-}
@@ -768,6 +874,92 @@ cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
768874 Skip s2' -> cmp_null s2'
769875 Stop -> return EQ
770876
877+------------------------------------------------------------------------------
878+-- Merging
879+------------------------------------------------------------------------------
880+
881+data MergeWithState l r o
882+ = MergeBoth l r
883+ | MergeWithLeft l r o
884+ | MergeWithRight l r o
885+ | MergeNoMore (Either l r)
886+
887+{-# INLINE_NORMAL mergeBy #-}
888+mergeBy :: (Monad m) => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
889+mergeBy f (Stream stepa ta) (Stream stepb tb) =
890+ Stream step (MergeBoth ta tb)
891+ {-# INLINE_LATE step #-}
892+ step gst (MergeBoth sa sb) = do
893+ ra <- stepa (rstState gst) sa
894+ rb <- stepb (rstState gst) sb
895+ case (ra, rb) of
896+ (Yield xa sa', Yield xb sb') -> do
897+ case f xa xb of
898+ LT ->
899+ return $ Yield xa (MergeWithRight sa' sb' xb)
900+ _ ->
901+ return $ Yield xb (MergeWithLeft sa' sb' xa)
902+ (Yield xa sa', Stop) -> return $ Yield xa (MergeNoMore (Left sa'))
903+ (Stop, Yield xb sb') -> return $ Yield xb (MergeNoMore (Right sb'))
904+ (Yield xa sa', Skip sb') -> return $ Skip (MergeWithLeft sa' sb' xa)
905+ (Skip sa', Yield xb sb') -> return $ Skip (MergeWithRight sa' sb' xb)
906+ (Skip sa', Skip sb') -> return $ Skip (MergeBoth sa' sb')
907+ (Skip sa', Stop) -> return $ Skip (MergeNoMore (Left sa'))
908+ (Stop, Skip sb') -> return $ Skip (MergeNoMore (Right sb'))
909+ (Stop, Stop) -> return Stop
910+ step gst (MergeWithLeft sa sb x) = do
911+ rb <- stepb (rstState gst) sb
912+ case rb of
913+ Yield xb sb' -> do
914+ case f x xb of
915+ LT -> return $ Yield x (MergeWithRight sa sb' xb)
916+ _ -> return $ Yield xb (MergeWithLeft sa sb' x)
917+ Skip sb' -> return $ Skip (MergeWithLeft sa sb' x)
918+ Stop -> return $ Yield x (MergeNoMore (Left sa))
919+ step gst (MergeWithRight sa sb x) = do
920+ ra <- stepa (rstState gst) sa
921+ case ra of
922+ Yield xa sa' -> do
923+ case f xa x of
924+ LT -> return $ Yield xa (MergeWithRight sa' sb x)
925+ _ -> return $ Yield x (MergeWithLeft sa' sb xa)
926+ Skip sa' -> return $ Skip (MergeWithRight sa' sb x)
927+ Stop -> return $ Yield x (MergeNoMore (Right sb))
928+ step gst (MergeNoMore (Left sa)) = do
929+ ra <- stepa (rstState gst) sa
930+ case ra of
931+ Yield xa sa' -> return $ Yield xa (MergeNoMore (Left sa'))
932+ Skip sa' -> return $ Skip (MergeNoMore (Left sa'))
933+ Stop -> return Stop
934+ step gst (MergeNoMore (Right sb)) = do
935+ rb <- stepb (rstState gst) sb
936+ case rb of
937+ Yield xb sb' -> return $ Yield xb (MergeNoMore (Right sb'))
938+ Skip sb' -> return $ Skip (MergeNoMore (Right sb'))
939+ Stop -> return Stop
940+
941+------------------------------------------------------------------------------
942+-- Transformation comprehensions
943+------------------------------------------------------------------------------
944+
945+{-# INLINE_NORMAL the #-}
946+the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
947+the (Stream step state) = go state
948+ where
949+ go st = do
950+ r <- step defState st
951+ case r of
952+ Yield x s -> go' x s
953+ Skip s -> go s
954+ Stop -> return Nothing
955+ go' n st = do
956+ r <- step defState st
957+ case r of
958+ Yield x s | x == n -> go' n s
959+ | otherwise -> return Nothing
960+ Skip s -> go' n s
961+ Stop -> return (Just n)
962+
771963------------------------------------------------------------------------------
772964-- Instances
773965------------------------------------------------------------------------------