Add foldMap for IntSet (#1048)

Author: meooow25

containers-tests/benchmarks/IntSet.hs

@@ -7,7 +7,7 @@ import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
 import Test.Tasty.Bench (bench, defaultMain, whnf)
 import Data.List (foldl')
-import Data.Monoid (Sum(..))
+import Data.Monoid (Sum(..), All(..))
 import qualified Data.IntSet as IS
 -- benchmarks for "instance Ord IntSet"
 -- uses IntSet as keys of maps, and elements of sets
@@ -56,6 +56,8 @@ main = do
         , bench "splitMember:dense" $ whnf (IS.splitMember elem_mid) s
         , bench "splitMember:sparse" $ whnf (IS.splitMember elem_sparse_mid) s_sparse
         , bench "eq" $ whnf (\s' -> s' == s') s -- worst case, compares everything
+        , bench "foldMap:dense" $ whnf (IS.foldMap (All . (>0))) s
+        , bench "foldMap:sparse" $ whnf (IS.foldMap (All . (>0))) s_sparse
         ]
   where
     bound = 2^12

containers-tests/tests/intset-properties.hs

@@ -14,7 +14,7 @@ import qualified Data.Foldable1 as Foldable1
 #endif
 import qualified Data.Set as Set
 import IntSetValidity (valid)
-import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl')
+import Prelude hiding (lookup, null, map, filter, foldr, foldl, foldl', foldMap)
 import Test.Tasty
 import Test.Tasty.HUnit
 import Test.Tasty.QuickCheck hiding ((.&.))
@@ -71,6 +71,7 @@ main = defaultMain $ testGroup "intset-properties"
                    , testProperty "prop_foldR'" prop_foldR'
                    , testProperty "prop_foldL" prop_foldL
                    , testProperty "prop_foldL'" prop_foldL'
+                   , testProperty "prop_foldMap" prop_foldMap
                    , testProperty "prop_map" prop_map
                    , testProperty "prop_mapMonotonicId" prop_mapMonotonicId
                    , testProperty "prop_mapMonotonicLinear" prop_mapMonotonicLinear
@@ -386,6 +387,9 @@ prop_foldL s = foldl (flip (:)) [] s == List.foldl (flip (:)) [] (toList s)
 prop_foldL' :: IntSet -> Bool
 prop_foldL' s = foldl' (flip (:)) [] s == List.foldl' (flip (:)) [] (toList s)
 
+prop_foldMap :: IntSet -> Property
+prop_foldMap s = foldMap (:[]) s === toList s
+
 prop_map :: IntSet -> Bool
 prop_map s = map id s == s
 

containers/changelog.md

@@ -51,6 +51,8 @@
 * Add `Intersection` and `intersections` for `Data.Set` and `Data.IntSet`.
   (Reed Mullanix, Soumik Sarkar)
 
+* Add `foldMap` for `Data.IntSet`. (Soumik Sarkar)
+
 ## Unreleased with `@since` annotation for 0.7.1:
 
 ### Additions

containers/src/Data/IntSet.hs

@@ -134,6 +134,7 @@ module Data.IntSet (
             -- * Folds
             , IS.foldr
             , IS.foldl
+            , IS.foldMap
             -- ** Strict folds
             , IS.foldr'
             , IS.foldl'

containers/src/Data/IntSet/Internal.hs

@@ -150,6 +150,7 @@ module Data.IntSet.Internal (
     -- * Folds
     , foldr
     , foldl
+    , foldMap
     -- ** Strict folds
     , foldr'
     , foldl'
@@ -206,7 +207,7 @@ import qualified Data.Foldable1 as Foldable1
 import Data.List.NonEmpty (NonEmpty(..))
 #endif
 import Utils.Containers.Internal.Prelude hiding
-  (filter, foldr, foldl, foldl', null, map)
+  (filter, foldr, foldl, foldl', foldMap, null, map)
 import Prelude ()
 
 import Utils.Containers.Internal.BitUtil
@@ -1252,6 +1253,29 @@ foldl' f z = \t ->      -- Use lambda t to be inlinable with two arguments only.
     go z' (Bin _ l r) = go (go z' l) r
 {-# INLINE foldl' #-}
 
+-- | \(O(n))\). Map the elements in the set to a monoid and combine with @(<>)@.
+foldMap :: Monoid a => (Key -> a) -> IntSet -> a
+foldMap f = \t ->  -- Use lambda t to be inlinable with one argument only.
+  case t of
+    Bin p l r
+#if MIN_VERSION_base(4,11,0)
+      | signBranch p -> go r <> go l  -- handle negative numbers
+      | otherwise -> go l <> go r
+#else
+      | signBranch p -> go r `mappend` go l  -- handle negative numbers
+      | otherwise -> go l `mappend` go r
+#endif
+    _ -> go t
+  where
+#if MIN_VERSION_base(4,11,0)
+    go (Bin _ l r) = go l <> go r
+#else
+    go (Bin _ l r) = go l `mappend` go r
+#endif
+    go (Tip kx bm) = foldMapBits kx f bm
+    go Nil = mempty
+{-# INLINE foldMap #-}
+
 {--------------------------------------------------------------------
   List variations
 --------------------------------------------------------------------}
@@ -1675,6 +1699,11 @@ foldlBits :: Int -> (a -> Int -> a) -> a -> Nat -> a
 foldl'Bits :: Int -> (a -> Int -> a) -> a -> Nat -> a
 foldrBits :: Int -> (Int -> a -> a) -> a -> Nat -> a
 foldr'Bits :: Int -> (Int -> a -> a) -> a -> Nat -> a
+#if MIN_VERSION_base(4,11,0)
+foldMapBits :: Semigroup a => Int -> (Int -> a) -> Nat -> a
+#else
+foldMapBits :: Monoid a => Int -> (Int -> a) -> Nat -> a
+#endif
 takeWhileAntitoneBits :: Int -> (Int -> Bool) -> Nat -> Nat
 
 {-# INLINE lowestBitSet #-}
@@ -1683,6 +1712,7 @@ takeWhileAntitoneBits :: Int -> (Int -> Bool) -> Nat -> Nat
 {-# INLINE foldl'Bits #-}
 {-# INLINE foldrBits #-}
 {-# INLINE foldr'Bits #-}
+{-# INLINE foldMapBits #-}
 {-# INLINE takeWhileAntitoneBits #-}
 
 lowestBitMask :: Nat -> Nat
@@ -1738,6 +1768,20 @@ foldr'Bits prefix f z bitmap = go (revNat bitmap) z
           where !bitmask = lowestBitMask bm
                 !bi = countTrailingZeros bitmask
 
+foldMapBits prefix f bitmap = go (prefix + bi0) (bitmap `xor` bitmask0)
+  where
+    bitmask0 = lowestBitMask bitmap
+    bi0 = countTrailingZeros bitmask0
+    go !x 0 = f x
+#if MIN_VERSION_base(4,11,0)
+    go !x bm = f x <> go (prefix + bi) (bm `xor` bitmask)
+#else
+    go !x bm = f x `mappend` go (prefix + bi) (bm `xor` bitmask)
+#endif
+      where
+        bitmask = lowestBitMask bm
+        bi = countTrailingZeros bitmask
+
 takeWhileAntitoneBits prefix predicate bitmap =
   -- Binary search for the first index where the predicate returns false, but skip a predicate
   -- call if the high half of the current range is empty. This ensures
@@ -1810,6 +1854,19 @@ foldr'Bits prefix f z bm = let lb = lowestBitSet bm
         go bi n | n `testBit` 0 = f bi $! go (bi + 1) (n `shiftRL` 1)
                 | otherwise     =         go (bi + 1) (n `shiftRL` 1)
 
+foldMapBits prefix f bm = go x0 (x0 + 1) ((bm `shiftRL` lb) `shiftRL` 1)
+  where
+    lb = lowestBitSet bm
+    x0 = prefix + lb
+    go !x !_ 0 = f x
+    go !x !bi n
+#if MIN_VERSION_base(4,11,0)
+      | n `testBit` 0 = f x <> go bi (bi + 1) (n `shiftRL` 1)
+#else
+      | n `testBit` 0 = f x `mappend` go bi (bi + 1) (n `shiftRL` 1)
+#endif
+      | otherwise = go x (bi + 1) (n `shiftRL` 1)
+
 takeWhileAntitoneBits prefix predicate = foldl'Bits prefix f 0 -- Does not use antitone property
   where
     f acc bi | predicate bi = acc .|. bitmapOf bi