Add Data.Set.alterF
authorSimon Jakobi <simon.jakobi@gmail.com>
Mon, 7 Oct 2019 17:44:11 +0000 (19:44 +0200)
committerDavid Feuer <David.Feuer@gmail.com>
Sat, 21 Dec 2019 23:46:53 +0000 (18:46 -0500)
Fixes #629.

containers-tests/benchmarks/Set.hs
containers-tests/tests/set-properties.hs
containers/src/Data/Set.hs
containers/src/Data/Set/Internal.hs

index f90524e..eca66ab 100644 (file)
@@ -12,6 +12,7 @@ main = do
     let s = S.fromAscList elems :: S.Set Int
         s_even = S.fromAscList elems_even :: S.Set Int
         s_odd = S.fromAscList elems_odd :: S.Set Int
+        strings_s = S.fromList strings
     evaluate $ rnf [s, s_even, s_odd]
     defaultMain
         [ bench "member" $ whnf (member elems) s
@@ -37,11 +38,19 @@ main = do
         , bench "disjoint:true" $ whnf (S.disjoint s_odd) s_even
         , bench "null.intersection:false" $ whnf (S.null. S.intersection s) s_even
         , bench "null.intersection:true" $ whnf (S.null. S.intersection s_odd) s_even
+        , bench "alterF:member" $ whnf (alterF_member elems) s
+        , bench "alterF:insert" $ whnf (alterF_ins elems) S.empty
+        , bench "alterF:delete" $ whnf (alterF_del elems) s
+        , bench "alterF:four" $ whnf (alterF_four elems) s
+        , bench "alterF:four:strings" $ whnf (alterF_four strings) strings_s
+        , bench "alterF_naive:four" $ whnf (alterF_naive_four elems) s
+        , bench "alterF_naive:four:strings" $ whnf (alterF_naive_four strings) strings_s
         ]
   where
     elems = [1..2^12]
     elems_even = [2,4..2^12]
     elems_odd = [1,3..2^12]
+    strings = map show elems
 
 member :: [Int] -> S.Set Int -> Int
 member xs s = foldl' (\n x -> if S.member x s then n + 1 else n) 0 xs
@@ -51,3 +60,47 @@ ins xs s0 = foldl' (\s a -> S.insert a s) s0 xs
 
 del :: [Int] -> S.Set Int -> S.Set Int
 del xs s0 = foldl' (\s k -> S.delete k s) s0 xs
+
+alterF_member :: [Int] -> S.Set Int -> Int
+alterF_member xs s = foldl' (\n x -> if member' x s then n + 1 else n) 0 xs
+  where
+    member' k s = getConsty (S.alterF (\b -> Consty b) k s)
+
+alterF_ins :: [Int] -> S.Set Int -> S.Set Int
+alterF_ins xs s0 = foldl' (\s a -> insert' a s) s0 xs
+  where
+    insert' k s = runIdent (S.alterF (const (Ident True)) k s)
+
+alterF_del :: [Int] -> S.Set Int -> S.Set Int
+alterF_del xs s0 = foldl' (\s k -> delete' k s) s0 xs
+  where
+    delete' k s = runIdent (S.alterF (const (Ident False)) k s)
+
+alterF_four :: Ord a => [a] -> S.Set a -> S.Set a
+alterF_four xs s0 = foldl' (\s k -> S.alterF four k s `seq` s) s0 xs
+
+alterF_naive_four :: Ord a => [a] -> S.Set a -> S.Set a
+alterF_naive_four xs s0 = foldl' (\s k -> alterF_naive four k s `seq` s) s0 xs
+
+alterF_naive :: (Ord a, Functor f) => (Bool -> f Bool) -> a -> S.Set a -> f (S.Set a)
+alterF_naive f k s = fmap g (f (k `S.member` s))
+  where
+    g True  = S.insert k s
+    g False = S.delete k s
+
+four :: Bool -> Four Bool
+               -- insert  delete  reinsert  toggle
+four True  = Four True    False   True      False
+four False = Four True    False   False     True
+
+newtype Consty a b = Consty { getConsty :: a}
+instance Functor (Consty a) where
+  fmap _ (Consty a) = Consty a
+
+newtype Ident a = Ident { runIdent :: a }
+instance Functor Ident where
+  fmap f (Ident a) = Ident (f a)
+
+data Four a = Four !a !a !a !a
+instance Functor Four where
+  fmap f (Four a b c d) = Four (f a) (f b) (f c) (f d)
index 14ba921..f8e7d26 100644 (file)
@@ -44,6 +44,12 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT
                    , testProperty "prop_InsertDelete" prop_InsertDelete
                    , testProperty "prop_InsertBiased" prop_InsertBiased
                    , testProperty "prop_DeleteValid" prop_DeleteValid
+                   , testProperty "alterF" prop_alterF
+                   , testProperty "alterF/delete" prop_alterF_delete
+                   , testProperty "alterF/insert" prop_alterF_insert
+                   , testProperty "alterF/member" prop_alterF_member
+                   , testProperty "alterF/four" prop_alterF_four
+                   , testProperty "alterF/valid" prop_alterF_valid
                    , testProperty "prop_Link" prop_Link
                    , testProperty "prop_Merge" prop_Merge
                    , testProperty "prop_UnionValid" prop_UnionValid
@@ -369,6 +375,52 @@ prop_DeleteValid :: Int -> Property
 prop_DeleteValid k = forValidUnitTree $ \t -> valid (delete k (insert k t))
 
 {--------------------------------------------------------------------
+  alterF
+--------------------------------------------------------------------}
+
+newtype Ident a = Ident { runIdent :: a }
+instance Functor Ident where
+  fmap f (Ident a) = Ident (f a)
+
+newtype Consty a b = Consty { getConsty :: a}
+instance Functor (Consty a) where
+  fmap _ (Consty a) = Consty a
+
+data Four a = Four a a a a
+  deriving (Eq, Show)
+instance Functor Four where
+  fmap f (Four a b c d) = Four (f a) (f b) (f c) (f d)
+
+four :: Bool -> Four Bool
+               -- insert  delete  id     toggle
+four True  = Four True    False   True   False
+four False = Four True    False   False  True
+
+toggle :: Ord a => a -> Set a -> Set a
+toggle k s =
+    if member k s
+        then delete k s
+        else insert k s
+
+prop_alterF :: Fun Bool [Bool] -> Int -> Set Int -> Property
+prop_alterF f k s = fmap (member k) (alterF (apply f) k s) === apply f (member k s)
+
+prop_alterF_insert :: Int -> Set Int -> Property
+prop_alterF_insert k s = runIdent (alterF (const (Ident True)) k s) === insert k s
+
+prop_alterF_delete :: Int -> Set Int -> Property
+prop_alterF_delete k s = runIdent (alterF (const (Ident False)) k s) === delete k s
+
+prop_alterF_member :: Int -> Set Int -> Property
+prop_alterF_member k s = getConsty (alterF (\b -> Consty b) k s) === member k s
+
+prop_alterF_four :: Int -> Set Int -> Property
+prop_alterF_four k s = alterF four k s === Four (insert k s) (delete k s) s (toggle k s)
+
+prop_alterF_valid :: Int -> Set Int -> Property
+prop_alterF_valid k s = fmap valid (alterF four k s) === Four True True True True
+
+{--------------------------------------------------------------------
   Balance
 --------------------------------------------------------------------}
 prop_Link :: Int -> Property
index 2dea90b..e3f7281 100644 (file)
@@ -88,6 +88,10 @@ module Data.Set (
             -- * Deletion
             , delete
 
+            -- * Generalized insertion/deletion
+
+            , alterF
+
             -- * Query
             , member
             , notMember
index c6ce3f2..97b11cf 100644 (file)
@@ -148,6 +148,7 @@ module Data.Set.Internal (
             , singleton
             , insert
             , delete
+            , alterF
             , powerSet
 
             -- * Combine
@@ -230,6 +231,7 @@ module Data.Set.Internal (
             ) where
 
 import Prelude hiding (filter,foldl,foldr,null,map,take,drop,splitAt)
+import Control.Applicative (Const(..))
 import qualified Data.List as List
 import Data.Bits (shiftL, shiftR)
 #if !MIN_VERSION_base(4,8,0)
@@ -245,6 +247,9 @@ import Data.Semigroup (Semigroup((<>)))
 import Data.Semigroup (stimesIdempotentMonoid)
 import Data.Functor.Classes
 #endif
+#if MIN_VERSION_base(4,8,0)
+import Data.Functor.Identity (Identity)
+#endif
 import qualified Data.Foldable as Foldable
 #if !MIN_VERSION_base(4,8,0)
 import Data.Foldable (Foldable (foldMap))
@@ -594,6 +599,70 @@ delete = go
 {-# INLINE delete #-}
 #endif
 
+-- | /O(log n)/ @('alterF' f x s)@ can delete or insert @x@ in @s@ depending on
+-- whether an equal element is found in @s@.
+--
+-- In short:
+--
+-- @
+-- 'member' x \<$\> 'alterF' f x s = f ('member' x s)
+-- @
+--
+-- Note that unlike 'insert', 'alterF' will /not/ replace an element equal to
+-- the given value.
+--
+-- Note: 'alterF' is a variant of the @at@ combinator from "Control.Lens.At".
+alterF :: (Ord a, Functor f) => (Bool -> f Bool) -> a -> Set a -> f (Set a)
+alterF f k s = fmap choose (f member_)
+  where
+    (member_, inserted, deleted) = case alteredSet k s of
+        Deleted d           -> (True , s, d)
+        Inserted i          -> (False, i, 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 :: Ord a => (Bool -> Identity Bool) -> a -> Set a -> Identity (Set a) #-}
+#endif
+
+data AlteredSet a
+      -- | The needle is present in the original set.
+      -- We return the set where the needle is deleted.
+    = Deleted !(Set a)
+
+      -- | The needle is not present in the original set.
+      -- We return the set with the needle inserted.
+    | Inserted !(Set a)
+
+alteredSet :: Ord a => a -> Set a -> AlteredSet a
+alteredSet x0 s0 = go x0 s0
+  where
+    go :: Ord a => a -> Set a -> AlteredSet a
+    go x Tip           = Inserted (singleton x)
+    go x (Bin _ y l r) = case compare x y of
+        LT -> case go x l of
+            Deleted d           -> Deleted (balanceR y d r)
+            Inserted i          -> Inserted (balanceL y i r)
+        GT -> case go x r of
+            Deleted d           -> Deleted (balanceL y l d)
+            Inserted i          -> Inserted (balanceR y l i)
+        EQ -> Deleted (glue l r)
+#if __GLASGOW_HASKELL__
+{-# INLINABLE alteredSet #-}
+#else
+{-# INLINE alteredSet #-}
+#endif
+
 {--------------------------------------------------------------------
   Subset
 --------------------------------------------------------------------}