Add Data.IntSet.alterF (#720)

* Add Data.IntSet.alterF

Fixes #719.

* Add Const rule and specialize for Identity

* Add tests

* Docs

* Work around missing Show instance for Const in GHC 7.8

Author: sjakobi

containers-tests/tests/intset-properties.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+import Control.Applicative (Const(..))
 import Data.Bits ((.&.), popCount)
 import Data.Word (Word)
 import Data.IntSet
@@ -71,6 +72,8 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT
                    , testProperty "prop_partition" prop_partition
                    , testProperty "prop_filter" prop_filter
                    , testProperty "prop_bitcount" prop_bitcount
+                   , testProperty "prop_alterF_list" prop_alterF_list
+                   , testProperty "prop_alterF_const" prop_alterF_const
                    ]
 
 ----------------------------------------------------------------
@@ -425,3 +428,21 @@ prop_bitcount a w = bitcount_orig a w == bitcount_new a w
       where go a 0 = a
             go a x = go (a + 1) (x .&. (x-1))
     bitcount_new a x = a + popCount x
+
+prop_alterF_list
+    :: Fun Bool [Bool]
+    -> Int
+    -> IntSet
+    -> Property
+prop_alterF_list f k s =
+        fmap toSet (alterF     (applyFun f) k s)
+    ===             Set.alterF (applyFun f) k (toSet s)
+
+prop_alterF_const
+    :: Fun Bool Bool
+    -> Int
+    -> IntSet
+    -> Property
+prop_alterF_const f k s =
+        getConst (alterF     (Const . applyFun f) k s        )
+    === getConst (Set.alterF (Const . applyFun f) k (toSet s))

containers/src/Data/IntSet.hs

@@ -86,6 +86,9 @@ module Data.IntSet (
             -- * Deletion
             , delete
 
+            -- * Generalized insertion/deletion
+            , alterF
+
             -- * Query
             , member
             , notMember

containers/src/Data/IntSet/Internal.hs

@@ -126,6 +126,7 @@ module Data.IntSet.Internal (
     , singleton
     , insert
     , delete
+    , alterF
 
     -- * Combine
     , union
@@ -188,6 +189,7 @@ module Data.IntSet.Internal (
     , zero
     ) where
 
+import Control.Applicative (Const(..))
 import Control.DeepSeq (NFData(rnf))
 import Data.Bits
 import qualified Data.List as List
@@ -222,7 +224,9 @@ import qualified GHC.Exts
 #endif
 
 import qualified Data.Foldable as Foldable
-#if !MIN_VERSION_base(4,8,0)
+#if MIN_VERSION_base(4,8,0)
+import Data.Functor.Identity (Identity(..))
+#else
 import Data.Foldable (Foldable())
 #endif
 
@@ -502,6 +506,40 @@ deleteBM kx bm t@(Tip kx' bm')
   | otherwise = t
 deleteBM _ _ Nil = Nil
 
+-- | /O(min(n,W))/. @('alterF' f x s)@ can delete or insert @x@ in @s@ depending
+-- on whether it is already present in @s@.
+--
+-- In short:
+--
+-- @
+-- 'member' x \<$\> 'alterF' f x s = f ('member' x s)
+-- @
+--
+-- Note: 'alterF' is a variant of the @at@ combinator from "Control.Lens.At".
+alterF :: Functor f => (Bool -> f Bool) -> Key -> IntSet -> f IntSet
+alterF f k s = fmap choose (f member_)
+  where
+    member_ = member k s
+
+    (inserted, deleted)
+      | member_   = (s         , delete k s)
+      | otherwise = (insert k s, s         )
+
+    choose True  = inserted
+    choose False = deleted
+#ifndef __GLASGOW_HASKELL__
+{-# INLINE alterF #-}
+#else
+{-# INLINABLE [2] alterF #-}
+
+{-# RULES
+"alterF/Const" forall k (f :: Bool -> Const a Bool) . alterF f k = \s -> Const . getConst . f $ member k s
+ #-}
+#endif
+
+#if MIN_VERSION_base(4,8,0)
+{-# SPECIALIZE alterF :: (Bool -> Identity Bool) -> Key -> IntSet -> Identity IntSet #-}
+#endif
 
 {--------------------------------------------------------------------
   Union