Add mapMaybe to sets and Seq (#1159)

Author: ChickenProp

containers-tests/tests/intmap-properties.hs

@@ -189,6 +189,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "filterKeys"           prop_filterKeys
              , testProperty "filterWithKey"        prop_filterWithKey
              , testProperty "partition"            prop_partition
+             , testProperty "mapMaybe"             prop_mapMaybe
              , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
              , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
              , testProperty "spanAntitone"         prop_spanAntitone
@@ -1552,6 +1553,13 @@ prop_filterWithKey fun m =
   where
     m' = filterWithKey (applyFun2 fun) m
 
+prop_mapMaybe :: Fun Int (Maybe Bool) -> IMap -> Property
+prop_mapMaybe f m =
+  valid m' .&&.
+  toList m' === Maybe.mapMaybe (\(k,x) -> (,) k <$> applyFun f x) (toList m)
+  where
+    m' = mapMaybe (applyFun f) m
+
 prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property
 prop_partition p ys = length ys > 0 ==>
   let xs = List.nubBy ((==) `on` fst) ys

containers-tests/tests/intset-properties.hs

@@ -6,6 +6,7 @@ import Data.IntSet
 import Data.List (nub,sort)
 import qualified Data.List as List
 import Data.Maybe (listToMaybe)
+import qualified Data.Maybe as Maybe
 import Data.Monoid (mempty)
 import Data.List.NonEmpty (NonEmpty(..))
 import qualified Data.List.NonEmpty as NE
@@ -79,6 +80,7 @@ main = defaultMain $ testGroup "intset-properties"
                    , testProperty "prop_splitRoot" prop_splitRoot
                    , testProperty "prop_partition" prop_partition
                    , testProperty "prop_filter" prop_filter
+                   , testProperty "prop_mapMaybe" prop_mapMaybe
                    , testProperty "takeWhileAntitone" prop_takeWhileAntitone
                    , testProperty "dropWhileAntitone" prop_dropWhileAntitone
                    , testProperty "spanAntitone" prop_spanAntitone
@@ -457,6 +459,12 @@ prop_filter s i =
      valid evens .&&.
      parts === (odds, evens)
 
+prop_mapMaybe :: Fun Int (Maybe Int) -> IntSet -> Property
+prop_mapMaybe f s =
+  let mapped = mapMaybe (applyFun f) s
+  in valid mapped .&&.
+     mapped === fromList (Maybe.mapMaybe (applyFun f) $ toList s)
+
 prop_takeWhileAntitone :: Int -> [Int] -> Property
 prop_takeWhileAntitone x ys =
   let l = takeWhileAntitone (<x) (fromList ys)

containers-tests/tests/seq-properties.hs

@@ -24,7 +24,8 @@ import Data.Array (listArray)
 import Data.Coerce (coerce)
 import Data.Foldable (Foldable(foldl, foldl1, foldr, foldr1, foldMap, fold), toList, all, sum, foldl', foldr')
 import Data.Functor ((<$>), (<$))
-import Data.Maybe
+import Data.Maybe (listToMaybe)
+import qualified Data.Maybe as Maybe
 import Data.Function (on)
 import Data.Monoid (Monoid(..), All(..), Endo(..), Dual(..))
 import Data.Semigroup (stimes, stimesMonoid)
@@ -100,6 +101,7 @@ main = defaultMain $ testGroup "seq-properties"
        , testProperty "breakr" prop_breakr
        , testProperty "partition" prop_partition
        , testProperty "filter" prop_filter
+       , testProperty "mapMaybe" prop_mapMaybe
        , testProperty "sort" prop_sort
        , testProperty "sortStable" prop_sortStable
        , testProperty "sortBy" prop_sortBy
@@ -553,6 +555,10 @@ prop_filter (Positive n) xs =
     toList' (filter p xs) ~= Prelude.filter p (toList xs)
   where p x = x `mod` n == 0
 
+prop_mapMaybe :: Fun Int (Maybe Int) -> Seq Int -> Bool
+prop_mapMaybe f xs =
+  toList' (mapMaybe (applyFun f) xs) ~= Maybe.mapMaybe (applyFun f) (toList xs)
+
 -- * Sorting
 
 prop_sort :: Seq OrdA -> Bool

containers-tests/tests/set-properties.hs

@@ -2,7 +2,8 @@
 import qualified Data.IntSet as IntSet
 import Data.List (nub, sort, sortBy)
 import qualified Data.List as List
-import Data.Maybe
+import Data.Maybe (isJust, fromJust)
+import qualified Data.Maybe as Maybe
 import Data.Set
 import Data.Set.Internal (link, merge)
 import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl', all, take, drop, splitAt)
@@ -98,6 +99,7 @@ main = defaultMain $ testGroup "set-properties"
                    , testProperty "prop_splitRoot" prop_splitRoot
                    , testProperty "prop_partition" prop_partition
                    , testProperty "prop_filter" prop_filter
+                   , testProperty "prop_mapMaybe" prop_mapMaybe
                    , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
                    , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
                    , testProperty "spanAntitone"         prop_spanAntitone
@@ -618,6 +620,12 @@ prop_partition s i = case partition odd s of
 prop_filter :: Set Int -> Int -> Bool
 prop_filter s i = partition odd s == (filter odd s, filter even s)
 
+prop_mapMaybe :: Fun Int (Maybe Int) -> Set Int -> Property
+prop_mapMaybe f s =
+  let mapped = mapMaybe (applyFun f) s
+  in valid mapped .&&.
+     mapped === fromList (Maybe.mapMaybe (applyFun f) $ toList s)
+
 prop_take :: Int -> Set Int -> Property
 prop_take n xs = valid taken .&&.
                  taken === fromDistinctAscList (List.take n (toList xs))

containers/changelog.md

@@ -7,6 +7,9 @@
 * Add `compareSize` for `IntSet` and `IntMap`. (Soumik Sarkar)
   ([#1135](https://github.com/haskell/containers/pull/1135))
 
+* Add `mapMaybe` for `Seq`, `Set` and `IntSet`. (Phil Hazelden)
+  ([#1159](https://github.com/haskell/containers/pull/1159)
+
 ### Performance improvements
 
 * Improved performance for `Data.IntMap.restrictKeys` and

containers/src/Data/IntSet.hs

@@ -145,6 +145,8 @@ module Data.IntSet (
             , dropWhileAntitone
             , spanAntitone
 
+            , mapMaybe
+
             , split
             , splitMember
             , splitRoot

containers/src/Data/IntSet/Internal.hs

@@ -134,6 +134,8 @@ module Data.IntSet.Internal (
     , dropWhileAntitone
     , spanAntitone
 
+    , mapMaybe
+
     , split
     , splitMember
     , splitRoot
@@ -880,6 +882,18 @@ filter predicate t
                          | otherwise           = bm
         {-# INLINE bitPred #-}
 
+-- | \(O(n \min(n,W))\). Map elements and collect the 'Just' results.
+--
+-- If the function is monotonically non-decreasing or monotonically
+-- non-increasing, 'mapMaybe' takes \(O(n)\) time.
+--
+-- @since FIXME
+mapMaybe :: (Key -> Maybe Key) -> IntSet -> IntSet
+mapMaybe f t = finishB (foldl' go emptyB t)
+  where go b x = case f x of
+          Nothing -> b
+          Just x' -> insertB x' b
+
 -- | \(O(n)\). partition the set according to some predicate.
 partition :: (Key -> Bool) -> IntSet -> (IntSet,IntSet)
 partition predicate0 t0 = toPair $ go predicate0 t0

containers/src/Data/Sequence.hs

@@ -192,6 +192,7 @@ module Data.Sequence (
     breakr,         -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a)
     partition,      -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a)
     filter,         -- :: (a -> Bool) -> Seq a -> Seq a
+    mapMaybe,       -- :: (a -> Maybe b) -> Seq a -> Seq b
     -- * Sorting
     sort,           -- :: Ord a => Seq a -> Seq a
     sortBy,         -- :: (a -> a -> Ordering) -> Seq a -> Seq a

containers/src/Data/Sequence/Internal.hs

@@ -128,6 +128,7 @@ module Data.Sequence.Internal (
     breakr,         -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a)
     partition,      -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a)
     filter,         -- :: (a -> Bool) -> Seq a -> Seq a
+    mapMaybe,       -- :: (a -> Maybe b) -> Seq a -> Seq b
     -- * Indexing
     lookup,         -- :: Int -> Seq a -> Maybe a
     (!?),           -- :: Seq a -> Int -> Maybe a
@@ -4199,6 +4200,15 @@ partition p = toPair . foldl' part (empty :*: empty)
 filter :: (a -> Bool) -> Seq a -> Seq a
 filter p = foldl' (\ xs x -> if p x then xs `snoc'` x else xs) empty
 
+-- | \( O(n) \). Map elements and collect the 'Just' results.
+--
+-- @since FIXME
+mapMaybe :: (a -> Maybe b) -> Seq a -> Seq b
+mapMaybe f = foldl' go empty
+  where go xs x = case f x of
+          Nothing -> xs
+          Just x' -> xs `snoc'` x'
+
 -- Indexing sequences
 
 -- | 'elemIndexL' finds the leftmost index of the specified element,

containers/src/Data/Set.hs

@@ -137,6 +137,7 @@ module Data.Set (
             , takeWhileAntitone
             , dropWhileAntitone
             , spanAntitone
+            , mapMaybe
             , partition
             , split
             , splitMember