Skip to content

Commit 576aac4

Browse files
chadaustinchadaustin
committed
add filterMap and ifilterMap
1 parent c0308f1 commit 576aac4

File tree

8 files changed

+94
-6
lines changed

8 files changed

+94
-6
lines changed

‎Data/Vector.hs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,9 @@ module Data.Vector (
112112
-- * Working with predicates
113113

114114
-- ** Filtering
115-
filter, ifilter, filterM,
115+
filter, ifilter,
116+
filterMap, ifilterMap,
117+
filterM,
116118
takeWhile, dropWhile,
117119

118120
-- ** Partitioning
@@ -1079,6 +1081,16 @@ ifilter :: (Int -> a -> Bool) -> Vector a -> Vector a
10791081
{-# INLINE ifilter #-}
10801082
ifilter = G.ifilter
10811083

1084+
-- | /O(n)/ Drop elements when predicate returns Nothing
1085+
filterMap :: (a -> Maybe b) -> Vector a -> Vector b
1086+
{-# INLINE filterMap #-}
1087+
filterMap = G.filterMap
1088+
1089+
-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
1090+
ifilterMap :: (Int -> a -> Maybe b) -> Vector a -> Vector b
1091+
{-# INLINE ifilterMap #-}
1092+
ifilterMap = G.ifilterMap
1093+
10821094
-- | /O(n)/ Drop elements that do not satisfy the monadic predicate
10831095
filterM :: Monad m => (a -> m Bool) -> Vector a -> m (Vector a)
10841096
{-# INLINE filterM #-}

‎Data/Vector/Fusion/Stream/Monadic.hs

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ module Data.Vector.Fusion.Stream.Monadic (
4040
eq, cmp,
4141

4242
-- * Filtering
43-
filter, filterM, takeWhile, takeWhileM, dropWhile, dropWhileM,
43+
filter, filterM, filterMap, takeWhile, takeWhileM, dropWhile, dropWhileM,
4444

4545
-- * Searching
4646
elem, notElem, find, findM, findIndex, findIndexM,
@@ -676,6 +676,21 @@ filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
676676
{-# INLINE filter #-}
677677
filter f = filterM (return . f)
678678

679+
filterMap :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
680+
{-# INLINE_FUSED filterMap #-}
681+
filterMap f (Stream step t) = Stream step' t
682+
where
683+
{-# INLINE_INNER step' #-}
684+
step' s = do
685+
r <- step s
686+
case r of
687+
Yield x s' -> do
688+
return $ case f x of
689+
Nothing -> Skip s'
690+
Just b' -> Yield b' s'
691+
Skip s' -> return $ Skip s'
692+
Done -> return $ Done
693+
679694
-- | Drop elements which do not satisfy the monadic predicate
680695
filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
681696
{-# INLINE_FUSED filterM #-}

‎Data/Vector/Generic.hs

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,9 @@ module Data.Vector.Generic (
9595
-- * Working with predicates
9696

9797
-- ** Filtering
98-
filter, ifilter, filterM,
98+
filter, ifilter,
99+
filterMap, ifilterMap,
100+
filterM,
99101
takeWhile, dropWhile,
100102

101103
-- ** Partitioning
@@ -1275,6 +1277,19 @@ ifilter f = unstream
12751277
. inplace (S.map snd . S.filter (uncurry f) . S.indexed) toMax
12761278
. stream
12771279

1280+
-- | /O(n)/ Drop elements when predicate returns Nothing
1281+
filterMap :: (Vector v a, Vector v b) => (a -> Maybe b) -> v a -> v b
1282+
{-# INLINE filterMap #-}
1283+
filterMap f = unstream . inplace (S.filterMap f) toMax . stream
1284+
1285+
-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
1286+
ifilterMap :: (Vector v a, Vector v b) => (Int -> a -> Maybe b) -> v a -> v b
1287+
{-# INLINE ifilterMap #-}
1288+
ifilterMap f = unstream
1289+
. inplace (S.filterMap (uncurry f) . S.indexed) toMax
1290+
. stream
1291+
1292+
12781293
-- | /O(n)/ Drop elements that do not satisfy the monadic predicate
12791294
filterM :: (Monad m, Vector v a) => (a -> m Bool) -> v a -> m (v a)
12801295
{-# INLINE filterM #-}

‎Data/Vector/Primitive.hs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,9 @@ module Data.Vector.Primitive (
9191
-- * Working with predicates
9292

9393
-- ** Filtering
94-
filter, ifilter, filterM,
94+
filter, ifilter,
95+
filterMap, ifilterMap,
96+
filterM,
9597
takeWhile, dropWhile,
9698

9799
-- ** Partitioning
@@ -887,6 +889,16 @@ ifilter :: Prim a => (Int -> a -> Bool) -> Vector a -> Vector a
887889
{-# INLINE ifilter #-}
888890
ifilter = G.ifilter
889891

892+
-- | /O(n)/ Drop elements when predicate returns Nothing
893+
filterMap :: (Prim a, Prim b) => (a -> Maybe b) -> Vector a -> Vector b
894+
{-# INLINE filterMap #-}
895+
filterMap = G.filterMap
896+
897+
-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
898+
ifilterMap :: (Prim a, Prim b) => (Int -> a -> Maybe b) -> Vector a -> Vector b
899+
{-# INLINE ifilterMap #-}
900+
ifilterMap = G.ifilterMap
901+
890902
-- | /O(n)/ Drop elements that do not satisfy the monadic predicate
891903
filterM :: (Monad m, Prim a) => (a -> m Bool) -> Vector a -> m (Vector a)
892904
{-# INLINE filterM #-}

‎Data/Vector/Storable.hs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,9 @@ module Data.Vector.Storable (
8888
-- * Working with predicates
8989

9090
-- ** Filtering
91-
filter, ifilter, filterM,
91+
filter, ifilter,
92+
filterMap, ifilterMap,
93+
filterM,
9294
takeWhile, dropWhile,
9395

9496
-- ** Partitioning
@@ -897,6 +899,16 @@ ifilter :: Storable a => (Int -> a -> Bool) -> Vector a -> Vector a
897899
{-# INLINE ifilter #-}
898900
ifilter = G.ifilter
899901

902+
-- | /O(n)/ Drop elements when predicate returns Nothing
903+
filterMap :: (Storable a, Storable b) => (a -> Maybe b) -> Vector a -> Vector b
904+
{-# INLINE filterMap #-}
905+
filterMap = G.filterMap
906+
907+
-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
908+
ifilterMap :: (Storable a, Storable b) => (Int -> a -> Maybe b) -> Vector a -> Vector b
909+
{-# INLINE ifilterMap #-}
910+
ifilterMap = G.ifilterMap
911+
900912
-- | /O(n)/ Drop elements that do not satisfy the monadic predicate
901913
filterM :: (Monad m, Storable a) => (a -> m Bool) -> Vector a -> m (Vector a)
902914
{-# INLINE filterM #-}

‎Data/Vector/Unboxed.hs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,9 @@ module Data.Vector.Unboxed (
118118
-- * Working with predicates
119119

120120
-- ** Filtering
121-
filter, ifilter, filterM,
121+
filter, ifilter,
122+
filterMap, ifilterMap,
123+
filterM,
122124
takeWhile, dropWhile,
123125

124126
-- ** Partitioning
@@ -945,6 +947,16 @@ ifilter :: Unbox a => (Int -> a -> Bool) -> Vector a -> Vector a
945947
{-# INLINE ifilter #-}
946948
ifilter = G.ifilter
947949

950+
-- | /O(n)/ Drop elements when predicate returns Nothing
951+
filterMap :: (Unbox a, Unbox b) => (a -> Maybe b) -> Vector a -> Vector b
952+
{-# INLINE filterMap #-}
953+
filterMap = G.filterMap
954+
955+
-- | /O(n)/ Drop elements when predicate, applied to index and value, returns Nothing
956+
ifilterMap :: (Unbox a, Unbox b) => (Int -> a -> Maybe b) -> Vector a -> Vector b
957+
{-# INLINE ifilterMap #-}
958+
ifilterMap = G.ifilterMap
959+
948960
-- | /O(n)/ Drop elements that do not satisfy the monadic predicate
949961
filterM :: (Monad m, Unbox a) => (a -> m Bool) -> Vector a -> m (Vector a)
950962
{-# INLINE filterM #-}

‎tests/Tests/Vector.hs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,7 @@ testPolymorphicFunctions _ = $(testProperties [
159159

160160
-- Filtering
161161
'prop_filter, 'prop_ifilter, {- prop_filterM, -}
162+
'prop_filterMap, 'prop_ifilterMap,
162163
'prop_takeWhile, 'prop_dropWhile,
163164

164165
-- Paritioning
@@ -292,6 +293,8 @@ testPolymorphicFunctions _ = $(testProperties [
292293

293294
prop_filter :: P ((a -> Bool) -> v a -> v a) = V.filter `eq` filter
294295
prop_ifilter :: P ((Int -> a -> Bool) -> v a -> v a) = V.ifilter `eq` ifilter
296+
prop_filterMap :: P ((a -> Maybe a) -> v a -> v a) = V.filterMap `eq` filterMap
297+
prop_ifilterMap :: P ((Int -> a -> Maybe a) -> v a -> v a) = V.ifilterMap `eq` ifilterMap
295298
prop_takeWhile :: P ((a -> Bool) -> v a -> v a) = V.takeWhile `eq` takeWhile
296299
prop_dropWhile :: P ((a -> Bool) -> v a -> v a) = V.dropWhile `eq` dropWhile
297300
prop_partition :: P ((a -> Bool) -> v a -> (v a, v a))

‎tests/Utilities.hs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import Data.Function (on)
1616
import Data.Functor.Identity
1717
import Data.List ( sortBy )
1818
import Data.Monoid
19+
import Data.Maybe (catMaybes)
1920

2021
instance Show a => Show (S.Bundle v a) where
2122
show s = "Data.Vector.Fusion.Bundle.fromList " ++ show (S.toList s)
@@ -292,6 +293,12 @@ izipWith3 = withIndexFirst zipWith3
292293
ifilter :: (Int -> a -> Bool) -> [a] -> [a]
293294
ifilter f = map snd . withIndexFirst filter f
294295

296+
filterMap :: (a -> Maybe b) -> [a] -> [b]
297+
filterMap f = catMaybes . map f
298+
299+
ifilterMap :: (Int -> a -> Maybe b) -> [a] -> [b]
300+
ifilterMap f = catMaybes . withIndexFirst map f
301+
295302
indexedLeftFold fld f z = fld (uncurry . f) z . zip [0..]
296303

297304
ifoldl :: (a -> Int -> a -> a) -> a -> [a] -> a

0 commit comments

Comments
 (0)