Add forM

containers-tests/benchmarks/Set.hs

@@ -1,12 +1,18 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE BangPatterns #-}
 
 module Main where
 
 import Control.DeepSeq (rnf)
 import Control.Exception (evaluate)
-import Gauge (bench, defaultMain, whnf)
+import Gauge (bench, defaultMain, whnf, nf)
 import Data.List (foldl')
 import qualified Data.Set as S
+import Data.Functor.Identity
+#if __GLASGOW_HASKELL__ >= 710
+import Control.Monad.ST
+import Data.STRef
+#endif
 
 main = do
     let s = S.fromAscList elems :: S.Set Int
@@ -21,6 +27,22 @@ main = do
         , bench "filter" $ whnf (S.filter ((== 0) . (`mod` 2))) s
         , bench "partition" $ whnf (S.partition ((== 0) . (`mod` 2))) s
         , bench "fold" $ whnf (S.fold (:) []) s
+#if __GLASGOW_HASKELL__ >= 710
+        , bench "forM (Identity)" $ nf (\s -> runIdentity $ S.forM s (\e -> pure $ e + 1)) s
+        , bench "for (Identity)" $ nf (\s -> runIdentity $ S.for s (\e -> pure $ e + 1)) s
+        , bench "forM (ST)" $ nf (\s -> runST $ do
+              ref <- newSTRef 0
+              S.forM s (\e -> do
+                              modifySTRef ref (+ 1)
+                              x <- readSTRef ref
+                              pure $ x + e * 36 `mod` 11)) (S.map (`mod` 51) s)
+        , bench "for (ST)" $ nf (\s -> runST $ do
+              ref <- newSTRef 0
+              S.for s (\e -> do
+                             modifySTRef ref (+ 1)
+                             x <- readSTRef ref
+                             pure $ x + e * 36 `mod` 11)) (S.map (`mod` 51) s)
+#endif
         , bench "delete" $ whnf (del elems) s
         , bench "findMin" $ whnf S.findMin s
         , bench "findMax" $ whnf S.findMax s

containers-tests/tests/set-properties.hs

@@ -5,7 +5,7 @@ import qualified Data.List as List
 import Data.Monoid (mempty)
 import Data.Maybe
 import Data.Set
-import Prelude hiding (lookup, null, map, filter, foldr, foldl, all, take, drop, splitAt)
+import Prelude hiding (lookup, null, map, filter, foldr, foldl, forM, all, take, drop, splitAt)
 import Test.Framework
 import Test.Framework.Providers.HUnit
 import Test.Framework.Providers.QuickCheck2
@@ -88,6 +88,10 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT
                    , testProperty "prop_map" prop_map
                    , testProperty "prop_map2" prop_map2
                    , testProperty "prop_mapMonotonic" prop_mapMonotonic
+#if __GLASGOW_HASKELL__ >= 710
+                   , testProperty "prop_forM" prop_forM
+                   , testProperty "prop_for" prop_for
+#endif
                    , testProperty "prop_maxView" prop_maxView
                    , testProperty "prop_minView" prop_minView
                    , testProperty "prop_split" prop_split
@@ -591,6 +595,14 @@ prop_map2 f g s = map (apply f) (map (apply g) s) === map (apply f . apply g) s
 prop_mapMonotonic :: Set Int -> Property
 prop_mapMonotonic s = mapMonotonic id s === s
 
+#if __GLASGOW_HASKELL__ >= 710
+prop_forM :: Set Int -> Bool
+prop_forM s = runIdentity (forM s (pure . id)) == map id s
+
+prop_for :: Set Int -> Bool
+prop_for s = runIdentity (for s (pure . id)) == map id s
+#endif
+
 prop_maxView :: Set Int -> Bool
 prop_maxView s = case maxView s of
     Nothing -> null s

containers/src/Data/Set.hs

@@ -136,6 +136,12 @@ module Data.Set (
             -- * Map
             , S.map
             , mapMonotonic
+#if __GLASGOW_HASKELL__ >= 710
+            , S.forM
+            , S.for
+            , S.mapM
+            , S.traverse
+#endif
 
             -- * Folds
             , S.foldr

containers/src/Data/Set/Internal.hs

@@ -181,6 +181,12 @@ module Data.Set.Internal (
             -- * Map
             , map
             , mapMonotonic
+#if __GLASGOW_HASKELL__ >= 710
+            , mapM
+            , forM
+            , traverse
+            , for
+#endif
 
             -- * Folds
             , foldr
@@ -230,7 +236,8 @@ module Data.Set.Internal (
             , merge
             ) where
 
-import Prelude hiding (filter,foldl,foldr,null,map,take,drop,splitAt)
+import Prelude hiding (filter,foldl,foldr,null,map,take,drop,splitAt,mapM,traverse)
+import qualified Prelude
 import Control.Applicative (Const(..))
 import qualified Data.List as List
 import Data.Bits (shiftL, shiftR)
@@ -953,6 +960,33 @@ mapMonotonic :: (a->b) -> Set a -> Set b
 mapMonotonic _ Tip = Tip
 mapMonotonic f (Bin sz x l r) = Bin sz (f x) (mapMonotonic f l) (mapMonotonic f r)
 
+
+#if __GLASGOW_HASKELL__ >= 710
+-- | /O(n*log n)/.
+--
+-- Like 'traverse' from Data.Traversable. This is less generic, since 'Set'
+-- does not have a Traversable instance.
+traverse :: (Ord b, Applicative m) => (a -> m b) -> Set a -> m (Set b)
+-- This implementation is significantly faster. We don't know why.
+traverse f s0 = fmap (GHCExts.fromListN (size s0)) . Prelude.traverse f . toList $ s0
+
+-- | 'for' is 'traverse' with its arguments flipped.
+for :: (Ord b, Applicative m) => Set a -> (a -> m b) -> m (Set b)
+for = flip traverse
+
+-- | /O(n*log n)/.
+--
+-- Like 'mapM' from Data.Traversable. This is less generic, since 'Set'
+-- does not have a Traversable instance.
+mapM :: (Ord b, Monad m) => (a -> m b) -> Set a -> m (Set b)
+mapM = traverse
+
+-- | 'forM' is 'mapM' with its arguments flipped.
+forM :: (Ord b, Monad m) => Set a -> (a -> m b) -> m (Set b)
+forM = flip mapM
+#endif
+
+
 {--------------------------------------------------------------------
   Fold
 --------------------------------------------------------------------}