Improve mapping over keys for IntMap and IntSet (#1148)

* Implement the functions in terms of the builders instead of converting
  to and from a list. This reduces time and allocations.
* Document that mapping takes linear time for monotonic functions.

Author: meooow25

containers-tests/benchmarks/IntMap.hs

@@ -10,7 +10,7 @@ import qualified Data.IntMap.Strict as MS
 import qualified Data.IntSet as S
 import Data.Maybe (fromMaybe)
 import Data.Word (Word8)
-import System.Random (StdGen, mkStdGen, randoms)
+import System.Random (StdGen, mkStdGen, random, randoms)
 import Prelude hiding (lookup)
 
 import Utils.Fold (foldBenchmarks, foldWithKeyBenchmarks)
@@ -55,6 +55,10 @@ main = do
         , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m
         , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m
         , bench "fromList:asc" $ whnf M.fromList elems_asc
+        , bench "mapKeys:asc" $ whnf (M.mapKeys (+1)) m
+        , bench "mapKeys:random" $ whnf (M.mapKeys (fst . random . mkStdGen)) m
+        , bench "mapKeysWith:asc:dups" $ whnf (M.mapKeysWith (+) (`div` 2)) m
+        , bench "mapKeysMonotonic" $ whnf (M.mapKeysMonotonic (+1)) m
         , bench "fromList:asc:fusion" $
             whnf (\n -> M.fromList (unitValues [1..n])) bound
         , bench "fromList:random" $ whnf M.fromList elems_random

containers-tests/benchmarks/IntSet.hs

@@ -15,7 +15,7 @@ import qualified Data.Set as S
 import qualified Data.IntMap as IM
 import qualified Data.Map.Strict as M
 import Data.Word (Word8)
-import System.Random (StdGen, mkStdGen, randoms, randomRs)
+import System.Random (StdGen, mkStdGen, randoms, random, randomRs)
 
 import Utils.Fold (foldBenchmarks)
 
@@ -36,7 +36,9 @@ main = do
     defaultMain
         [ bench "member" $ whnf (member elems) s
         , bench "insert" $ whnf (ins elems) IS.empty
-        , bench "map" $ whnf (IS.map (+ 1)) s
+        , bench "map:asc" $ whnf (IS.map (+ 1)) s
+        , bench "map:random" $ whnf (IS.map (fst . random . mkStdGen)) s
+        , bench "mapMonotonic" $ whnf (IS.mapMonotonic (+1)) s
         , bench "filter" $ whnf (IS.filter ((== 0) . (`mod` 2))) s
         , bench "partition" $ whnf (IS.partition ((== 0) . (`mod` 2))) s
         , bench "delete" $ whnf (del elems) s

containers/src/Data/IntMap/Internal.hs

@@ -2653,6 +2653,9 @@ mapAccumRWithKey f a t
 -- | \(O(n \min(n,W))\).
 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
 --
+-- If `f` is monotonically non-decreasing or monotonically non-increasing, this
+-- function takes \(O(n)\) time.
+--
 -- The size of the result may be smaller if @f@ maps two or more distinct
 -- keys to the same new key.  In this case the value at the greatest of the
 -- original keys is retained.
@@ -2662,11 +2665,14 @@ mapAccumRWithKey f a t
 -- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c"
 
 mapKeys :: (Key->Key) -> IntMap a -> IntMap a
-mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
+mapKeys f t = finishB (foldlWithKey' (\b kx x -> insertB (f kx) x b) emptyB t)
 
 -- | \(O(n \min(n,W))\).
 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
 --
+-- If `f` is monotonically non-decreasing or monotonically non-increasing, this
+-- function takes \(O(n)\) time.
+--
 -- The size of the result may be smaller if @f@ maps two or more distinct
 -- keys to the same new key.  In this case the associated values will be
 -- combined using @c@.
@@ -2677,8 +2683,8 @@ mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) []
 -- Also see the performance note on 'fromListWith'.
 
 mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a
-mapKeysWith c f
-  = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
+mapKeysWith c f t =
+  finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB t)
 
 -- | \(O(n)\).
 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
@@ -2700,8 +2706,8 @@ mapKeysWith c f
 -- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")]
 
 mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a
-mapKeysMonotonic f
-  = fromDistinctAscList . foldrWithKey (\k x xs -> (f k, x) : xs) []
+mapKeysMonotonic f t =
+  ascLinkAll (foldlWithKey' (\s kx x -> ascInsert s (f kx) x) MSNada t)
 
 {--------------------------------------------------------------------
   Filter
@@ -3487,7 +3493,8 @@ fromListWithKey f xs =
 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
 
 fromAscList :: [(Key,a)] -> IntMap a
-fromAscList xs = fromAscListWithKey (\_ x _ -> x) xs
+fromAscList xs =
+  ascLinkAll (Foldable.foldl' (\s (ky, y) -> ascInsert s ky y) MSNada xs)
 {-# INLINE fromAscList #-} -- Inline for list fusion
 
 -- | \(O(n)\). Build a map from a list of key\/value pairs where
@@ -3555,6 +3562,17 @@ data MonoState a
   = MSNada
   | MSPush {-# UNPACK #-} !Key a !(Stack a)
 
+-- Insert an entry. The key must be >= the last inserted key. If it is equal
+-- to the previous key, the previous value is replaced.
+ascInsert :: MonoState a -> Int -> a -> MonoState a
+ascInsert s !ky y = case s of
+  MSNada -> MSPush ky y Nada
+  MSPush kx x stk
+    | kx == ky -> MSPush ky y stk
+    | otherwise -> let m = branchMask kx ky
+                   in MSPush ky y (ascLinkTop stk kx (Tip kx x) m)
+{-# INLINE ascInsert #-}
+
 ascLinkTop :: Stack a -> Int -> IntMap a -> Int -> Stack a
 ascLinkTop stk !rk r !rm = case stk of
   Nada -> Push rm r stk

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

@@ -970,6 +970,9 @@ mapAccumRWithKey f0 a0 t0 = toPair $ go f0 a0 t0
 -- | \(O(n \min(n,W))\).
 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
 --
+-- If `f` is monotonically non-decreasing or monotonically non-increasing, this
+-- function takes \(O(n)\) time.
+--
 -- The size of the result may be smaller if @f@ maps two or more distinct
 -- keys to the same new key.  In this case the associated values will be
 -- combined using @c@.
@@ -980,7 +983,8 @@ mapAccumRWithKey f0 a0 t0 = toPair $ go f0 a0 t0
 -- Also see the performance note on 'fromListWith'.
 
 mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a
-mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) []
+mapKeysWith c f t =
+  finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB t)
 
 {--------------------------------------------------------------------
   Filter

containers/src/Data/IntSet/Internal.hs

@@ -1183,11 +1183,14 @@ deleteMax = maybe Nil snd . maxView
 -- | \(O(n \min(n,W))\).
 -- @'map' f s@ is the set obtained by applying @f@ to each element of @s@.
 --
+-- If `f` is monotonically non-decreasing or monotonically non-increasing, this
+-- function takes \(O(n)\) time.
+--
 -- It's worth noting that the size of the result may be smaller if,
 -- for some @(x,y)@, @x \/= y && f x == f y@
 
 map :: (Key -> Key) -> IntSet -> IntSet
-map f = fromList . List.map f . toList
+map f t = finishB (foldl' (\b x -> insertB (f x) b) emptyB t)
 
 -- | \(O(n)\). The
 --
@@ -1203,11 +1206,8 @@ map f = fromList . List.map f . toList
 -- precondition may not hold.
 --
 -- @since 0.6.3.1
-
--- Note that for now the test is insufficient to support any fancier implementation.
 mapMonotonic :: (Key -> Key) -> IntSet -> IntSet
-mapMonotonic f = fromDistinctAscList . List.map f . toAscList
-
+mapMonotonic f t = ascLinkAll (foldl' (\s x -> ascInsert s (f x)) MSNada t)
 
 {--------------------------------------------------------------------
   Fold
@@ -1464,17 +1464,7 @@ fromRange (lx,rx)
 
 -- See Note [fromAscList implementation] in Data.IntMap.Internal.
 fromAscList :: [Key] -> IntSet
-fromAscList xs = ascLinkAll (Foldable.foldl' next MSNada xs)
-  where
-    next s !ky = case s of
-      MSNada -> MSPush py bmy Nada
-      MSPush px bmx stk
-        | px == py -> MSPush py (bmx .|. bmy) stk
-        | otherwise -> let m = branchMask px py
-                       in MSPush py bmy (ascLinkTop stk px (Tip px bmx) m)
-      where
-        py = prefixOf ky
-        bmy = bitmapOf ky
+fromAscList xs = ascLinkAll (Foldable.foldl' ascInsert MSNada xs)
 {-# INLINE fromAscList #-} -- Inline for list fusion
 
 -- | \(O(n)\). Build a set from an ascending list of distinct elements.
@@ -1498,6 +1488,19 @@ data MonoState
   = MSNada
   | MSPush {-# UNPACK #-} !Int {-# UNPACK #-} !BitMap !Stack
 
+-- Insert an element. The element must be >= the last inserted element.
+ascInsert :: MonoState -> Int -> MonoState
+ascInsert s !ky = case s of
+  MSNada -> MSPush py bmy Nada
+  MSPush px bmx stk
+    | px == py -> MSPush py (bmx .|. bmy) stk
+    | otherwise -> let m = branchMask px py
+                   in MSPush py bmy (ascLinkTop stk px (Tip px bmx) m)
+  where
+    py = prefixOf ky
+    bmy = bitmapOf ky
+{-# INLINE ascInsert #-}
+
 ascLinkTop :: Stack -> Int -> IntSet -> Int -> Stack
 ascLinkTop stk !rk r !rm = case stk of
   Nada -> Push rm r stk