diff --git a/containers-tests/benchmarks/IntMap.hs b/containers-tests/benchmarks/IntMap.hs index d9f40cda3..c3f5522f0 100644 --- a/containers-tests/benchmarks/IntMap.hs +++ b/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 diff --git a/containers-tests/benchmarks/IntSet.hs b/containers-tests/benchmarks/IntSet.hs index 9ef9ff89d..e32841d54 100644 --- a/containers-tests/benchmarks/IntSet.hs +++ b/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 diff --git a/containers/src/Data/IntMap/Internal.hs b/containers/src/Data/IntMap/Internal.hs index 6fc058c97..54e6f986c 100644 --- a/containers/src/Data/IntMap/Internal.hs +++ b/containers/src/Data/IntMap/Internal.hs @@ -2642,6 +2642,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. @@ -2651,11 +2654,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@. @@ -2666,8 +2672,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@ @@ -2689,8 +2695,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 @@ -3413,7 +3419,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 @@ -3481,6 +3488,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 diff --git a/containers/src/Data/IntMap/Strict/Internal.hs b/containers/src/Data/IntMap/Strict/Internal.hs index fac2953bc..623d61ed0 100644 --- a/containers/src/Data/IntMap/Strict/Internal.hs +++ b/containers/src/Data/IntMap/Strict/Internal.hs @@ -949,6 +949,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@. @@ -959,7 +962,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 diff --git a/containers/src/Data/IntSet/Internal.hs b/containers/src/Data/IntSet/Internal.hs index 94459d16b..55570fc39 100644 --- a/containers/src/Data/IntSet/Internal.hs +++ b/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 @@ -1441,17 +1441,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. @@ -1475,6 +1465,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