Merge pull request #74 from chadaustin/master

add mapMaybe and imapMaybe

Author: dolio

Data/Vector.hs

@@ -112,7 +112,9 @@ module Data.Vector (
   -- * Working with predicates
 
   -- ** Filtering
-  filter, ifilter, filterM,
+  filter, ifilter,
+  mapMaybe, imapMaybe,
+  filterM,
   takeWhile, dropWhile,
 
   -- ** Partitioning
@@ -1097,6 +1099,16 @@ ifilter :: (Int -> a -> Bool) -> Vector a -> Vector a
 {-# INLINE ifilter #-}
 ifilter = G.ifilter
 
+-- | /O(n)/ Drop elements when predicate returns Nothing
+mapMaybe :: (a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE mapMaybe #-}
+mapMaybe = G.mapMaybe
+
+-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
+imapMaybe :: (Int -> a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE imapMaybe #-}
+imapMaybe = G.imapMaybe
+
 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
 filterM :: Monad m => (a -> m Bool) -> Vector a -> m (Vector a)
 {-# INLINE filterM #-}

Data/Vector/Fusion/Stream/Monadic.hs

@@ -40,7 +40,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   eq, cmp,
 
   -- * Filtering
-  filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
+  filter, filterM, mapMaybe, takeWhile, takeWhileM, dropWhile, dropWhileM,
 
   -- * Searching
   elem, notElem, find, findM, findIndex, findIndexM,
@@ -688,6 +688,21 @@ filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
 {-# INLINE filter #-}
 filter f = filterM (return . f)
 
+mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
+{-# INLINE_FUSED mapMaybe #-}
+mapMaybe f (Stream step t) = Stream step' t
+  where
+    {-# INLINE_INNER step' #-}
+    step' s = do
+                r <- step s
+                case r of
+                  Yield x s' -> do
+                                  return $ case f x of
+                                    Nothing -> Skip s'
+                                    Just b' -> Yield b' s'
+                  Skip    s' -> return $ Skip s'
+                  Done       -> return $ Done
+
 -- | Drop elements which do not satisfy the monadic predicate
 filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
 {-# INLINE_FUSED filterM #-}

Data/Vector/Generic.hs

@@ -95,7 +95,9 @@ module Data.Vector.Generic (
   -- * Working with predicates
 
   -- ** Filtering
-  filter, ifilter, filterM,
+  filter, ifilter,
+  mapMaybe, imapMaybe,
+  filterM,
   takeWhile, dropWhile,
 
   -- ** Partitioning
@@ -1287,6 +1289,19 @@ ifilter f = unstream
           . inplace (S.map snd . S.filter (uncurry f) . S.indexed) toMax
           . stream
 
+-- | /O(n)/ Drop elements when predicate returns Nothing
+mapMaybe :: (Vector v a, Vector v b) => (a -> Maybe b) -> v a -> v b
+{-# INLINE mapMaybe #-}
+mapMaybe f = unstream . inplace (S.mapMaybe f) toMax . stream
+
+-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
+imapMaybe :: (Vector v a, Vector v b) => (Int -> a -> Maybe b) -> v a -> v b
+{-# INLINE imapMaybe #-}
+imapMaybe f = unstream
+          . inplace (S.mapMaybe (uncurry f) . S.indexed) toMax
+          . stream
+
+
 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
 filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
 {-# INLINE filterM #-}

Data/Vector/Primitive.hs

@@ -91,7 +91,9 @@ module Data.Vector.Primitive (
   -- * Working with predicates
 
   -- ** Filtering
-  filter, ifilter, filterM,
+  filter, ifilter,
+  mapMaybe, imapMaybe,
+  filterM,
   takeWhile, dropWhile,
 
   -- ** Partitioning
@@ -903,6 +905,16 @@ ifilter :: Prim a => (Int -> a -> Bool) -> Vector a -> Vector a
 {-# INLINE ifilter #-}
 ifilter = G.ifilter
 
+-- | /O(n)/ Drop elements when predicate returns Nothing
+mapMaybe :: (Prim a, Prim b) => (a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE mapMaybe #-}
+mapMaybe = G.mapMaybe
+
+-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
+imapMaybe :: (Prim a, Prim b) => (Int -> a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE imapMaybe #-}
+imapMaybe = G.imapMaybe
+
 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
 filterM :: (Monad m, Prim a) => (a -> m Bool) -> Vector a -> m (Vector a)
 {-# INLINE filterM #-}

Data/Vector/Storable.hs

@@ -88,7 +88,9 @@ module Data.Vector.Storable (
   -- * Working with predicates
 
   -- ** Filtering
-  filter, ifilter, filterM,
+  filter, ifilter,
+  mapMaybe, imapMaybe,
+  filterM,
   takeWhile, dropWhile,
 
   -- ** Partitioning
@@ -913,6 +915,16 @@ ifilter :: Storable a => (Int -> a -> Bool) -> Vector a -> Vector a
 {-# INLINE ifilter #-}
 ifilter = G.ifilter
 
+-- | /O(n)/ Drop elements when predicate returns Nothing
+mapMaybe :: (Storable a, Storable b) => (a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE mapMaybe #-}
+mapMaybe = G.mapMaybe
+
+-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
+imapMaybe :: (Storable a, Storable b) => (Int -> a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE imapMaybe #-}
+imapMaybe = G.imapMaybe
+
 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
 filterM :: (Monad m, Storable a) => (a -> m Bool) -> Vector a -> m (Vector a)
 {-# INLINE filterM #-}

Data/Vector/Unboxed.hs

@@ -118,7 +118,9 @@ module Data.Vector.Unboxed (
   -- * Working with predicates
 
   -- ** Filtering
-  filter, ifilter, filterM,
+  filter, ifilter,
+  mapMaybe, imapMaybe,
+  filterM,
   takeWhile, dropWhile,
 
   -- ** Partitioning
@@ -960,6 +962,16 @@ ifilter :: Unbox a => (Int -> a -> Bool) -> Vector a -> Vector a
 {-# INLINE ifilter #-}
 ifilter = G.ifilter
 
+-- | /O(n)/ Drop elements when predicate returns Nothing
+mapMaybe :: (Unbox a, Unbox b) => (a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE mapMaybe #-}
+mapMaybe = G.mapMaybe
+
+-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
+imapMaybe :: (Unbox a, Unbox b) => (Int -> a -> Maybe b) -> Vector a -> Vector b
+{-# INLINE imapMaybe #-}
+imapMaybe = G.imapMaybe
+
 -- | /O(n)/ Drop elements that do not satisfy the monadic predicate
 filterM :: (Monad m, Unbox a) => (a -> m Bool) -> Vector a -> m (Vector a)
 {-# INLINE filterM #-}

tests/Tests/Vector.hs

@@ -160,6 +160,7 @@ testPolymorphicFunctions _ = $(testProperties [
 
         -- Filtering
         'prop_filter, 'prop_ifilter, {- prop_filterM, -}
+        'prop_mapMaybe, 'prop_imapMaybe,
         'prop_takeWhile, 'prop_dropWhile,
 
         -- Paritioning
@@ -298,6 +299,8 @@ testPolymorphicFunctions _ = $(testProperties [
 
     prop_filter :: P ((a -> Bool) -> v a -> v a) = V.filter `eq` filter
     prop_ifilter :: P ((Int -> a -> Bool) -> v a -> v a) = V.ifilter `eq` ifilter
+    prop_mapMaybe :: P ((a -> Maybe a) -> v a -> v a) = V.mapMaybe `eq` mapMaybe
+    prop_imapMaybe :: P ((Int -> a -> Maybe a) -> v a -> v a) = V.imapMaybe `eq` imapMaybe
     prop_takeWhile :: P ((a -> Bool) -> v a -> v a) = V.takeWhile `eq` takeWhile
     prop_dropWhile :: P ((a -> Bool) -> v a -> v a) = V.dropWhile `eq` dropWhile
     prop_partition :: P ((a -> Bool) -> v a -> (v a, v a))

tests/Utilities.hs

@@ -16,6 +16,7 @@ import Data.Function (on)
 import Data.Functor.Identity
 import Data.List ( sortBy )
 import Data.Monoid
+import Data.Maybe (catMaybes)
 
 instance Show a => Show (S.Bundle v a) where
     show s = "Data.Vector.Fusion.Bundle.fromList " ++ show (S.toList s)
@@ -292,6 +293,12 @@ izipWith3 = withIndexFirst zipWith3
 ifilter :: (Int -> a -> Bool) -> [a] -> [a]
 ifilter f = map snd . withIndexFirst filter f
 
+mapMaybe :: (a -> Maybe b) -> [a] -> [b]
+mapMaybe f = catMaybes . map f
+
+imapMaybe :: (Int -> a -> Maybe b) -> [a] -> [b]
+imapMaybe f = catMaybes . withIndexFirst map f
+
 indexedLeftFold fld f z = fld (uncurry . f) z . zip [0..]
 
 ifoldl :: (a -> Int -> a -> a) -> a -> [a] -> a