From e46d341819e05367af7083ae5a9ff0d3c67259a3 Mon Sep 17 00:00:00 2001
From: Xiaokui Shu <subbyte@gmail.com>
Date: Sat, 26 May 2018 15:58:04 -0400
Subject: [PATCH] add mapMaybe and mapMaybeM

---
 Changelog.md            |  1 +
 src/Streamly/Prelude.hs | 29 +++++++++++++++++++++++++++++
 test/Prop.hs            |  4 ++++
 3 files changed, 34 insertions(+)

diff --git a/Changelog.md b/Changelog.md
index d1548a4f..88cd2ee3 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -26,6 +26,7 @@
   concurrently when used at appropriate stream types.
 * Added concurrent function application operators to run stages of a
   stream processing function application pipeline concurrently.
+* Added two functions `mapMaybe` and `mapMaybeM`.
 
 ## 0.2.0
 
diff --git a/src/Streamly/Prelude.hs b/src/Streamly/Prelude.hs
index 61fb76f5..cb6246a4 100644
--- a/src/Streamly/Prelude.hs
+++ b/src/Streamly/Prelude.hs
@@ -100,12 +100,14 @@ module Streamly.Prelude
     , takeWhile
     , drop
     , dropWhile
+    , mapMaybe
 
     -- * Reordering
     , reverse
 
     -- * Mapping
     , mapM
+    , mapMaybeM
     , sequence
 
     -- * Zipping
@@ -129,6 +131,7 @@ where
 import           Control.Monad (void)
 import           Control.Monad.IO.Class      (MonadIO(..))
 import           Data.Semigroup              (Semigroup(..))
+import           Data.Maybe                  (isJust, fromJust)
 import           Prelude hiding              (filter, drop, dropWhile, take,
                                               takeWhile, zipWith, foldr, foldl,
                                               mapM, mapM_, sequence, all, any,
@@ -570,6 +573,24 @@ dropWhile p m = fromStream $ go (toStream m)
                       | otherwise = yld a r
          in (S.runStream m1) ctx stp single yield
 
+-- | Map + filter: compute a stream of b from a stream of a and filter out
+-- elements that are mapped to Nothing. It is similar to 'mapMaybe' in
+-- 'Data.Maybe' but applies to streams instead of list.
+--
+-- @since 0.2.1
+{-# INLINE mapMaybe #-}
+mapMaybe :: (IsStream t) => (a -> Maybe b) -> t m a -> t m b
+mapMaybe f m = go (toStream m)
+  where
+    go m1 = fromStream $ Stream $ \_ stp sng yld ->
+        let single a = case f a of
+                Just b  -> sng b
+                Nothing -> stp
+            yield a r = case f a of
+                Just b  -> yld b (toStream $ go r)
+                Nothing -> (S.runStream r) Nothing stp single yield
+        in (S.runStream m1) Nothing stp single yield
+
 -- | Determine whether all elements of a stream satisfy a predicate.
 --
 -- @since 0.1.0
@@ -750,6 +771,14 @@ mapM f m = go (toStream m)
             yield a r = S.runStream (toStream (f a |: (go r))) svr stp sng yld
          in (S.runStream m1) Nothing stp single yield
 
+-- | Monadic version mapMaybe.
+--
+-- @since 0.2.1
+{-# INLINE mapMaybeM #-}
+mapMaybeM :: (IsStream t, MonadAsync m, Functor (t m))
+          => (a -> m (Maybe b)) -> t m a -> t m b
+mapMaybeM f = fmap fromJust . filter isJust . mapM f
+
 -- XXX this can utilize parallel mapping if we implement it as runStream . mapM
 -- | Apply a monadic action to each element of the stream and discard the
 -- output of the action.
diff --git a/test/Prop.hs b/test/Prop.hs
index df03d29f..a4955767 100644
--- a/test/Prop.hs
+++ b/test/Prop.hs
@@ -9,6 +9,7 @@ import Control.Applicative (ZipList(..))
 import Control.Concurrent (MVar, takeMVar, putMVar, newEmptyMVar)
 import Control.Monad (replicateM, replicateM_)
 import Data.List (sort, foldl', scanl')
+import Data.Maybe (mapMaybe)
 import GHC.Word (Word8)
 
 import Test.Hspec.QuickCheck (prop)
@@ -295,6 +296,9 @@ transformOps constr desc t eq = do
     prop (desc ++ " takeWhile > 0") $
         transform (takeWhile (> 0)) $ t . (A.takeWhile (> 0))
 
+    let f x = if odd x then Just (x + 100) else Nothing
+    prop (desc ++ " mapMaybe") $ transform (mapMaybe f) $ t . (A.mapMaybe f)
+
     prop (desc ++ " drop maxBound") $
         transform (drop maxBound) $ t . (A.drop maxBound)
     prop (desc ++ " drop 0") $ transform (drop 0) $ t . (A.drop 0)