From cbd884c4dc4410b73837e7272f5fbfd6e4e0841d Mon Sep 17 00:00:00 2001
From: meooow25 <soumiksarkar.3120@gmail.com>
Date: Sat, 28 Jun 2025 21:18:35 +0530
Subject: [PATCH] Add pop for maps and sets

This operation has been requested since it comes up once in a while.
Currently the best way to perform it is to separately use lookup and
delete or use updateLookupWithKey. Both of these are less efficient and
harder to read.
---
 containers-tests/tests/intmap-properties.hs   | 11 ++++
 containers-tests/tests/intset-properties.hs   |  6 +++
 containers-tests/tests/map-properties.hs      | 11 ++++
 containers-tests/tests/set-properties.hs      |  6 +++
 containers/src/Data/IntMap/Internal.hs        | 40 +++++++++++++++
 containers/src/Data/IntMap/Lazy.hs            |  1 +
 containers/src/Data/IntMap/Strict.hs          |  1 +
 containers/src/Data/IntMap/Strict/Internal.hs |  2 +
 containers/src/Data/IntSet.hs                 |  1 +
 containers/src/Data/IntSet/Internal.hs        | 33 ++++++++++++
 containers/src/Data/Map/Internal.hs           | 50 +++++++++++++++++++
 containers/src/Data/Map/Lazy.hs               |  1 +
 containers/src/Data/Map/Strict.hs             |  1 +
 containers/src/Data/Map/Strict/Internal.hs    |  2 +
 containers/src/Data/Set.hs                    |  1 +
 containers/src/Data/Set/Internal.hs           | 29 +++++++++++
 16 files changed, 196 insertions(+)

diff --git a/containers-tests/tests/intmap-properties.hs b/containers-tests/tests/intmap-properties.hs
index ec198d929..17e4d0819 100644
--- a/containers-tests/tests/intmap-properties.hs
+++ b/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
diff --git a/containers-tests/tests/intset-properties.hs b/containers-tests/tests/intset-properties.hs
index 76b5813b4..f7770cd2b 100644
--- a/containers-tests/tests/intset-properties.hs
+++ b/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)
 
diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs
index a585391c4..920f0e0da 100644
--- a/containers-tests/tests/map-properties.hs
+++ b/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
diff --git a/containers-tests/tests/set-properties.hs b/containers-tests/tests/set-properties.hs
index aeb4e65c3..d77871dc6 100644
--- a/containers-tests/tests/set-properties.hs
+++ b/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)
 
diff --git a/containers/src/Data/IntMap/Internal.hs b/containers/src/Data/IntMap/Internal.hs
index ee65250b2..f36d77dd0 100644
--- a/containers/src/Data/IntMap/Internal.hs
+++ b/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)
diff --git a/containers/src/Data/IntMap/Lazy.hs b/containers/src/Data/IntMap/Lazy.hs
index e6be3148e..617bc4c8f 100644
--- a/containers/src/Data/IntMap/Lazy.hs
+++ b/containers/src/Data/IntMap/Lazy.hs
@@ -123,6 +123,7 @@ module Data.IntMap.Lazy (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
diff --git a/containers/src/Data/IntMap/Strict.hs b/containers/src/Data/IntMap/Strict.hs
index 91c083152..84660312e 100644
--- a/containers/src/Data/IntMap/Strict.hs
+++ b/containers/src/Data/IntMap/Strict.hs
@@ -141,6 +141,7 @@ module Data.IntMap.Strict (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
diff --git a/containers/src/Data/IntMap/Strict/Internal.hs b/containers/src/Data/IntMap/Strict/Internal.hs
index 85dd7d63f..01ad0bbce 100644
--- a/containers/src/Data/IntMap/Strict/Internal.hs
+++ b/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
diff --git a/containers/src/Data/IntSet.hs b/containers/src/Data/IntSet.hs
index c78c3a0bb..864f6fe45 100644
--- a/containers/src/Data/IntSet.hs
+++ b/containers/src/Data/IntSet.hs
@@ -109,6 +109,7 @@ module Data.IntSet (
 
             -- * Deletion
             , delete
+            , pop
 
             -- * Generalized insertion/deletion
             , alterF
diff --git a/containers/src/Data/IntSet/Internal.hs b/containers/src/Data/IntSet/Internal.hs
index 9008244e3..9a6b79993 100644
--- a/containers/src/Data/IntSet/Internal.hs
+++ b/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@.
 --
diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index e6491b2d2..e3963fee2 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/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)
diff --git a/containers/src/Data/Map/Lazy.hs b/containers/src/Data/Map/Lazy.hs
index 56ca3c536..dbabbc7e8 100644
--- a/containers/src/Data/Map/Lazy.hs
+++ b/containers/src/Data/Map/Lazy.hs
@@ -133,6 +133,7 @@ module Data.Map.Lazy (
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
diff --git a/containers/src/Data/Map/Strict.hs b/containers/src/Data/Map/Strict.hs
index de0fb18c4..ae362cf38 100644
--- a/containers/src/Data/Map/Strict.hs
+++ b/containers/src/Data/Map/Strict.hs
@@ -147,6 +147,7 @@ module Data.Map.Strict
 
     -- * Deletion\/Update
     , delete
+    , pop
     , adjust
     , adjustWithKey
     , update
diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs
index d70977e38..6790c988b 100644
--- a/containers/src/Data/Map/Strict/Internal.hs
+++ b/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
diff --git a/containers/src/Data/Set.hs b/containers/src/Data/Set.hs
index 7eaf80622..a4fe4a811 100644
--- a/containers/src/Data/Set.hs
+++ b/containers/src/Data/Set.hs
@@ -102,6 +102,7 @@ module Data.Set (
 
             -- * Deletion
             , delete
+            , pop
 
             -- * Generalized insertion/deletion
 
diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs
index 9d2021c8b..cae5de72e 100644
--- a/containers/src/Data/Set/Internal.hs
+++ b/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@.
 --