Add pop for maps and sets

containers-tests/tests/intmap-properties.hs

@@ -222,6 +222,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "isSubmapOfBy"         prop_isSubmapOfBy
              , testProperty "insert"               prop_insert
              , testProperty "delete"               prop_delete
+             , testProperty "pop"                  prop_pop
              , testProperty "insertWith"           prop_insertWith
              , testProperty "insertWithKey"        prop_insertWithKey
              , testProperty "insertLookupWithKey"  prop_insertLookupWithKey
@@ -1763,6 +1764,16 @@ prop_delete k m = valid m' .&&. toList m' === kxs'
     kxs = toList m
     kxs' = maybe kxs (\v -> kxs List.\\ [(k,v)]) (List.lookup k kxs)
 
+prop_pop :: Int -> IntMap A -> Property
+prop_pop k m = case pop k m of
+  Nothing -> property $ notMember k m
+  Just (x, m') ->
+    valid m' .&&.
+    Just x === List.lookup k kxs .&&.
+    toList m' === kxs List.\\ [(k,x)]
+    where
+      kxs = toList m
+
 prop_insertWith :: Fun (A, A) A -> Int -> A -> IntMap A -> Property
 prop_insertWith f k x m = valid m' .&&. toList m' === kxs'
   where

containers-tests/tests/intset-properties.hs

@@ -88,6 +88,7 @@ main = defaultMain $ testGroup "intset-properties"
                    , testProperty "intersections_lazy" prop_intersections_lazy
                    , testProperty "insert" prop_insert
                    , testProperty "delete" prop_delete
+                   , testProperty "pop" prop_pop
                    , testProperty "deleteMin" prop_deleteMin
                    , testProperty "deleteMax" prop_deleteMax
                    , testProperty "fromAscList" prop_fromAscList
@@ -521,6 +522,11 @@ prop_delete x s = valid s' .&&. toList s' === toList s List.\\ [x]
   where
     s' = delete x s
 
+prop_pop :: Int -> IntSet -> Property
+prop_pop x s = case pop x s of
+  Nothing -> property $ notMember x s
+  Just s' -> valid s' .&&. toList s' === toList s List.\\ [x]
+
 prop_deleteMin :: IntSet -> Property
 prop_deleteMin s = toList (deleteMin s) === if null s then [] else tail (toList s)
 

containers-tests/tests/map-properties.hs

@@ -269,6 +269,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "compare"              prop_compare
          , testProperty "insert"               prop_insert
          , testProperty "delete"               prop_delete
+         , testProperty "pop"                  prop_pop
          , testProperty "insertWith"           prop_insertWith
          , testProperty "insertWithKey"        prop_insertWithKey
          , testProperty "insertLookupWithKey"  prop_insertLookupWithKey
@@ -1747,6 +1748,16 @@ prop_delete k m = valid m' .&&. toList m' === kxs'
     kxs = toList m
     kxs' = maybe kxs (\v -> kxs List.\\ [(k,v)]) (List.lookup k kxs)
 
+prop_pop :: Int -> Map Int A -> Property
+prop_pop k m = case pop k m of
+  Nothing -> property $ notMember k m
+  Just (x, m') ->
+    valid m' .&&.
+    Just x === List.lookup k kxs .&&.
+    toList m' === kxs List.\\ [(k,x)]
+    where
+      kxs = toList m
+
 prop_insertWith :: Fun (A, A) A -> Int -> A -> Map Int A -> Property
 prop_insertWith f k x m = valid m' .&&. toList m' === kxs'
   where

containers-tests/tests/set-properties.hs

@@ -113,6 +113,7 @@ main = defaultMain $ testGroup "set-properties"
                    , testProperty "intersections_lazy" prop_intersections_lazy
                    , testProperty "insert" prop_insert
                    , testProperty "delete" prop_delete
+                   , testProperty "pop" prop_pop
                    , testProperty "deleteMin" prop_deleteMin
                    , testProperty "deleteMax" prop_deleteMax
                    , testProperty "findIndex" prop_findIndex
@@ -715,6 +716,11 @@ prop_delete x s = valid s' .&&. toList s' === toList s List.\\ [x]
   where
     s' = delete x s
 
+prop_pop :: Int -> Set Int -> Property
+prop_pop x s = case pop x s of
+  Nothing -> property $ notMember x s
+  Just s' -> valid s' .&&. toList s' === toList s List.\\ [x]
+
 prop_deleteMin :: Set Int -> Property
 prop_deleteMin s = toList (deleteMin s) === if null s then [] else tail (toList s)
 

containers/src/Data/IntMap/Internal.hs

@@ -118,6 +118,7 @@ module Data.IntMap.Internal (
 
     -- ** Delete\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
@@ -965,6 +966,43 @@ delete k t@(Tip ky _)
   | otherwise     = t
 delete _k Nil = Nil
 
+-- | \(O(\min(n,W))\). Pop an entry from the map.
+--
+-- Returns @Nothing@ if the key is not in the map. Otherwise returns @Just@ the
+-- value at the key and a map with the entry removed.
+--
+-- @
+-- pop 1 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Nothing
+-- pop 2 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Just ("b",fromList [(0,"a"),(4,"c")])
+-- @
+--
+-- @since FIXME
+pop :: Key -> IntMap a -> Maybe (a, IntMap a)
+pop k0 t0 = case go k0 t0 of
+  Popped (Just y) t -> Just (y, t)
+  _ -> Nothing
+  where
+    go !k (Bin p l r)
+      | nomatch k p = Popped Nothing Nil
+      | left k p = case go k l of
+          Popped y@(Just _) l' -> Popped y (binCheckL p l' r)
+          q -> q
+      | otherwise = case go k r of
+          Popped y@(Just _) r' -> Popped y (binCheckR p l r')
+          q -> q
+    go !k (Tip ky y)
+      | k == ky = Popped (Just y) Nil
+      | otherwise = Popped Nothing Nil
+    go !_ Nil = Popped Nothing Nil
+
+-- See Note [Popped impl] in Data.Map.Internal
+data Popped k a = Popped
+#if __GLASGOW_HASKELL__ >= 906
+  {-# UNPACK #-}
+#endif
+  !(Maybe a)
+  !(IntMap a)
+
 -- | \(O(\min(n,W))\). Adjust a value at a specific key. When the key is not
 -- a member of the map, the original map is returned.
 --
@@ -1056,6 +1094,8 @@ upsert f !k Nil = Tip k (f Nothing)
 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")])
 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing,  fromList [(3, "b"), (5, "a")])
 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
+--
+-- See also: 'pop'
 
 updateLookupWithKey ::  (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a)
 updateLookupWithKey f !k (Bin p l r)

containers/src/Data/IntMap/Lazy.hs

@@ -123,6 +123,7 @@ module Data.IntMap.Lazy (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update

containers/src/Data/IntMap/Strict.hs

@@ -141,6 +141,7 @@ module Data.IntMap.Strict (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update

containers/src/Data/IntMap/Strict/Internal.hs

@@ -84,6 +84,7 @@ module Data.IntMap.Strict.Internal (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
@@ -277,6 +278,7 @@ import Data.IntMap.Internal
   , mergeWithKey'
   , compose
   , delete
+  , pop
   , deleteMin
   , deleteMax
   , deleteFindMax

containers/src/Data/IntSet.hs

@@ -109,6 +109,7 @@ module Data.IntSet (
 
             -- * Deletion
             , delete
+            , pop
 
             -- * Generalized insertion/deletion
             , alterF

containers/src/Data/IntSet/Internal.hs

@@ -115,6 +115,7 @@ module Data.IntSet.Internal (
     , fromRange
     , insert
     , delete
+    , pop
     , alterF
 
     -- * Combine
@@ -558,6 +559,38 @@ deleteBM kx bm t@(Tip kx' bm')
   | otherwise = t
 deleteBM _ _ Nil = Nil
 
+-- | \(O(\min(n,W))\). Pop an element from the set.
+--
+-- Returns @Nothing@ if the element is not a member of the set. Otherwise
+-- returns @Just@ the set with the element removed.
+--
+-- @
+-- pop 1 (fromList [0,2,4]) == Nothing
+-- pop 2 (fromList [0,2,4]) == Just (fromList [0,4])
+-- @
+--
+-- @since FIXME
+pop :: Key -> IntSet -> Maybe IntSet
+pop x0 t0 = case go x0 t0 of
+  True :*: t -> Just t
+  _ -> Nothing
+  where
+    -- See Note [Popped impl] in Data.Map.Internal
+    go !x (Bin p l r)
+      | nomatch x p = False :*: Nil
+      | left x p = case go x l of
+          True :*: l' -> True :*: binCheckL p l' r
+          q -> q
+      | otherwise = case go x r of
+          True :*: r' -> True :*: binCheckR p l r'
+          q -> q
+    go !x (Tip ky bmy)
+      | prefixOf x == ky && bmx .&. bmy /= 0 = True :*: tip ky (bmx `xor` bmy)
+      | otherwise = False :*: Nil
+      where
+        bmx = bitmapOf x
+    go !_ Nil = False :*: 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@.
 --

containers/src/Data/Map/Internal.hs

@@ -159,6 +159,7 @@ module Data.Map.Internal (
 
     -- ** Delete\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
@@ -1035,6 +1036,53 @@ delete = go
 {-# INLINE delete #-}
 #endif
 
+-- | \(O(\log n)\). Pop an entry from the map.
+--
+-- Returns @Nothing@ if the key is not in the map. Otherwise returns @Just@ the
+-- value at the key and a map with the entry removed.
+--
+-- @
+-- pop 1 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Nothing
+-- pop 2 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Just ("b",fromList [(0,"a"),(4,"c")])
+-- @
+--
+-- @since FIXME
+pop :: Ord k => k -> Map k a -> Maybe (a, Map k a)
+pop k0 t0 = case go k0 t0 of
+  Popped (Just y) t -> Just (y, t)
+  _ -> Nothing
+  where
+    go !k (Bin _ kx x l r) = case compare k kx of
+      LT -> case go k l of
+        Popped y@(Just _) l' -> Popped y (balanceR kx x l' r)
+        q -> q
+      EQ -> Popped (Just x) (glue l r)
+      GT -> case go k r of
+        Popped y@(Just _) r' -> Popped y (balanceL kx x l r')
+        q -> q
+    go !_ Tip = Popped Nothing Tip
+
+-- Note [Popped impl]
+-- ~~~~~~~~~~~~~~~~~~
+-- Popped is implemented as a pair, though a sum makes more sense:
+--   data Popped k a = NotPopped | Popped a !(Map k a)
+-- This is because GHC optimizes a return value of `Popped k a` to
+-- `(# Maybe a, Map k a #)`, avoiding all Popped allocations in `pop`.
+-- GHC cannot do this with a sum type yet, see GHC #14259. Manually using
+-- unboxed sums avoids the allocations but GHC loses strictness information,
+-- see #25988.
+--
+-- On GHC>=9.6 we unbox the Maybe and avoid that allocation too, so `pop`'s `go`
+-- returns `(# (# (# #) | a #), Map k a #)`.
+
+data Popped k a = Popped
+#if __GLASGOW_HASKELL__ >= 906
+  {-# UNPACK #-}
+#endif
+  !(Maybe a)
+  !(Map k a)
+{-# INLINABLE pop #-}
+
 -- | \(O(\log n)\). Update a value at a specific key with the result of the provided function.
 -- When the key is not
 -- a member of the map, the original map is returned.
@@ -1148,6 +1196,8 @@ upsert f !k Tip = singleton k (f Nothing)
 -- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")])
 -- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing,  fromList [(3, "b"), (5, "a")])
 -- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a")
+--
+-- See also: 'pop'
 
 -- See Note: Type of local 'go' function
 updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a)

containers/src/Data/Map/Lazy.hs

@@ -133,6 +133,7 @@ module Data.Map.Lazy (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update

containers/src/Data/Map/Strict.hs

@@ -147,6 +147,7 @@ module Data.Map.Strict
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update

containers/src/Data/Map/Strict/Internal.hs

@@ -97,6 +97,7 @@ module Data.Map.Strict.Internal
 
     -- ** Delete\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
@@ -337,6 +338,7 @@ import Data.Map.Internal
   , elems
   , empty
   , delete
+  , pop
   , deleteAt
   , deleteFindMax
   , deleteFindMin

containers/src/Data/Set.hs

@@ -102,6 +102,7 @@ module Data.Set (
 
             -- * Deletion
             , delete
+            , pop
 
             -- * Generalized insertion/deletion
 

containers/src/Data/Set/Internal.hs

@@ -142,6 +142,7 @@ module Data.Set.Internal (
             , singleton
             , insert
             , delete
+            , pop
             , alterF
             , powerSet
 
@@ -587,6 +588,34 @@ delete = go
 {-# INLINE delete #-}
 #endif
 
+-- | \(O(\log n)\). Pop an element from the set.
+--
+-- Returns @Nothing@ if the element is not a member of the set. Otherwise
+-- returns @Just@ the set with the element removed.
+--
+-- @
+-- pop 1 (fromList [0,2,4]) == Nothing
+-- pop 2 (fromList [0,2,4]) == Just (fromList [0,4])
+-- @
+--
+-- @since FIXME
+pop :: Ord a => a -> Set a -> Maybe (Set a)
+pop x0 t0 = case go x0 t0 of
+  True :*: t -> Just t
+  _ -> Nothing
+  where
+    -- See Note [Popped impl] in Data.Map.Internal
+    go !x (Bin _ y l r) = case compare x y of
+      LT -> case go x l of
+        True :*: l' -> True :*: balanceR y l' r
+        q -> q
+      EQ -> True :*: glue l r
+      GT -> case go x r of
+        True :*: r' -> True :*: balanceL y l r'
+        q -> q
+    go !_ Tip = False :*: Tip
+{-# INLINABLE pop #-}
+
 -- | \(O(\log n)\) @('alterF' f x s)@ can delete or insert @x@ in @s@ depending on
 -- whether an equal element is found in @s@.
 --