diff --git a/containers-tests/tests/intmap-properties.hs b/containers-tests/tests/intmap-properties.hs index ec198d929..613ed1eb2 100644 --- a/containers-tests/tests/intmap-properties.hs +++ b/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 diff --git a/containers-tests/tests/intset-properties.hs b/containers-tests/tests/intset-properties.hs index 76b5813b4..d889976e1 100644 --- a/containers-tests/tests/intset-properties.hs +++ b/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 (), (<$)) -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 diff --git a/containers-tests/tests/set-properties.hs b/containers-tests/tests/set-properties.hs index aeb4e65c3..29bd94ca8 100644 --- a/containers-tests/tests/set-properties.hs +++ b/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)) diff --git a/containers/changelog.md b/containers/changelog.md index 1e3632985..8f53bbbce 100644 --- a/containers/changelog.md +++ b/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 diff --git a/containers/src/Data/IntSet.hs b/containers/src/Data/IntSet.hs index c78c3a0bb..f6f50d33c 100644 --- a/containers/src/Data/IntSet.hs +++ b/containers/src/Data/IntSet.hs @@ -145,6 +145,8 @@ module Data.IntSet ( , dropWhileAntitone , spanAntitone + , mapMaybe + , split , splitMember , splitRoot diff --git a/containers/src/Data/IntSet/Internal.hs b/containers/src/Data/IntSet/Internal.hs index 9008244e3..bfa2b7083 100644 --- a/containers/src/Data/IntSet/Internal.hs +++ b/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 diff --git a/containers/src/Data/Sequence.hs b/containers/src/Data/Sequence.hs index 9c9c52973..0adaeb2e0 100644 --- a/containers/src/Data/Sequence.hs +++ b/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 diff --git a/containers/src/Data/Sequence/Internal.hs b/containers/src/Data/Sequence/Internal.hs index ac9796ab3..59e13e84e 100644 --- a/containers/src/Data/Sequence/Internal.hs +++ b/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, diff --git a/containers/src/Data/Set.hs b/containers/src/Data/Set.hs index 7eaf80622..096420907 100644 --- a/containers/src/Data/Set.hs +++ b/containers/src/Data/Set.hs @@ -137,6 +137,7 @@ module Data.Set ( , takeWhileAntitone , dropWhileAntitone , spanAntitone + , mapMaybe , partition , split , splitMember diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs index 4682abc8c..0af7eb150 100644 --- a/containers/src/Data/Set/Internal.hs +++ b/containers/src/Data/Set/Internal.hs @@ -162,6 +162,7 @@ module Data.Set.Internal ( , takeWhileAntitone , dropWhileAntitone , spanAntitone + , mapMaybe , partition , split , splitMember @@ -993,6 +994,23 @@ partition p0 t0 = toPair $ go p0 t0 then t else link x l2 r2) +{-------------------------------------------------------------------- + Maybes +--------------------------------------------------------------------} + +-- | \(O(n \log n)\). Map elements and collect the 'Just' results. +-- +-- If the function is monotonically non-decreasing, this function takes \(O(n)\) +-- time. +-- +-- @since FIXME +mapMaybe :: Ord b => (a -> Maybe b) -> Set a -> Set b +mapMaybe f t = finishB (foldl' go emptyB t) + where go b x = case f x of + Nothing -> b + Just x' -> insertB x' b +{-# INLINABLE mapMaybe #-} + {---------------------------------------------------------------------- Map ----------------------------------------------------------------------}