Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit aa64020

Browse files
Improve IntMap's restrictKeys and withoutKeys (#1131)
* Properly implement the early return optimization * No need to have a separate withoutBM * Remove the unnecessary closure in withoutKeys * Add some benchmarks The dense case benchmarks improve by 98% since they benefit maximally from the early return. For the random case, withoutKeys improves by 14% and restrictKeys improves by 20%.
1 parent 2658197 commit aa64020

File tree

4 files changed

+90
-90
lines changed

4 files changed

+90
-90
lines changed

‎containers-tests/benchmarks/IntMap.hs

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import Test.Tasty.Bench (bench, bgroup, defaultMain, whnf)
77
import Data.List (foldl')
88
import qualified Data.IntMap as M
99
import qualified Data.IntMap.Strict as MS
10+
import qualified Data.IntSet as S
1011
import Data.Maybe (fromMaybe)
1112
import System.Random (StdGen, mkStdGen, randoms, randomRs)
1213
import Prelude hiding (lookup)
@@ -19,8 +20,13 @@ main = do
1920
let m'' = M.fromAscList elems_most :: M.IntMap Int
2021
let m''' = M.fromAscList elems_misses :: M.IntMap Int
2122
let m'''' = M.fromAscList elems_mixed :: M.IntMap Int
23+
m_random = M.fromList elems_random
24+
s = S.fromList keys
25+
s_random2 = S.fromList keys_random2
2226
evaluate $ rnf [elems_asc, elems_random, elems_randomDups]
2327
evaluate $ rnf [m, m', m'', m''', m'''']
28+
evaluate $ rnf m_random
29+
evaluate $ rnf [s, s_random2]
2430
defaultMain
2531
[ bench "lookup_hits" $ whnf (lookup keys) m
2632
, bench "lookup_half" $ whnf (lookup keys) m'
@@ -65,6 +71,14 @@ main = do
6571
, bench "splitLookup" $ whnf (M.splitLookup key_mid) m
6672
, bench "eq" $ whnf (\m' -> m' == m') m -- worst case, compares everything
6773
, bench "compare" $ whnf (\m' -> compare m' m') m -- worst case, compares everything
74+
, bench "withoutKeys" $ -- dense, equal keys
75+
whnf (uncurry M.withoutKeys) (m, s)
76+
, bench "restrictKeys" $ -- dense, equal keys
77+
whnf (uncurry M.restrictKeys) (m, s)
78+
, bench "withoutKeys:random" $ -- large keys, no overlap
79+
whnf (uncurry M.withoutKeys) (m_random, s_random2)
80+
, bench "restrictKeys:random" $ -- large keys, no overlap
81+
whnf (uncurry M.restrictKeys) (m_random, s_random2)
6882
, bgroup "folds" $ foldBenchmarks M.foldr M.foldl M.foldr' M.foldl' foldMap m
6983
, bgroup "folds with key" $
7084
foldWithKeyBenchmarks M.foldrWithKey M.foldlWithKey M.foldrWithKey' M.foldlWithKey' M.foldMapWithKey m
@@ -89,6 +103,7 @@ main = do
89103
mixedKeys = interleave keys keys'
90104
values = [1..bound]
91105
key_mid = bound `div` 2
106+
keys_random2 = take bound (randoms gen2)
92107
--------------------------------------------------------
93108
sum k v1 v2 = k + v1 + v2
94109
consPair k v xs = (k, v) : xs
@@ -149,5 +164,6 @@ unitValues :: [Int] -> [(Int, ())]
149164
unitValues = map (flip (,) ())
150165
{-# INLINE unitValues #-}
151166

152-
gen :: StdGen
167+
gen, gen2 :: StdGen
153168
gen = mkStdGen 42
169+
gen2 = mkStdGen 90

‎containers/changelog.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,13 @@
11
# Changelog for [`containers` package](http://github.com/haskell/containers)
22

3+
## Next release
4+
5+
### Performance improvements
6+
7+
* Improved performance for `Data.IntMap.restrictKeys` and
8+
`Data.IntMap.withoutKeys`. (Soumik Sarkar)
9+
([#1131](https://github.com/haskell/containers/pull/1131))
10+
311
## 0.8 *March 2025*
412

513
### Breaking changes

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

Lines changed: 63 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,7 @@ import Utils.Containers.Internal.Prelude hiding
320320
(lookup, map, filter, foldr, foldl, foldl', null)
321321
import Prelude ()
322322

323+
import Data.IntSet.Internal (IntSet)
323324
import qualified Data.IntSet.Internal as IntSet
324325
import Data.IntSet.Internal.IntTreeCommons
325326
( Key
@@ -335,7 +336,7 @@ import Data.IntSet.Internal.IntTreeCommons
335336
, i2w
336337
, Order(..)
337338
)
338-
import Utils.Containers.Internal.BitUtil (shiftLL, shiftRL, iShiftRL)
339+
import Utils.Containers.Internal.BitUtil (shiftLL, shiftRL, iShiftRL, wordSize)
339340
import Utils.Containers.Internal.StrictPair
340341

341342
#ifdef __GLASGOW_HASKELL__
@@ -402,20 +403,11 @@ data IntMap a = Bin {-# UNPACK #-} !Prefix
402403
-- See Note [Okasaki-Gill] for how the implementation here relates to the one in
403404
-- Okasaki and Gill's paper.
404405

405-
-- Some stuff from "Data.IntSet.Internal", for 'restrictKeys' and
406-
-- 'withoutKeys' to use.
407-
type IntSetPrefix = Int
408-
type IntSetBitMap = Word
409-
410406
#ifdef __GLASGOW_HASKELL__
411407
-- | @since 0.6.6
412408
deriving instance Lift a => Lift (IntMap a)
413409
#endif
414410

415-
bitmapOf :: Int -> IntSetBitMap
416-
bitmapOf x = shiftLL 1 (x .&. IntSet.suffixBitMask)
417-
{-# INLINE bitmapOf #-}
418-
419411
{--------------------------------------------------------------------
420412
Operators
421413
--------------------------------------------------------------------}
@@ -1196,7 +1188,7 @@ differenceWithKey f m1 m2
11961188
-- @
11971189
--
11981190
-- @since 0.5.8
1199-
withoutKeys :: IntMap a -> IntSet.IntSet -> IntMap a
1191+
withoutKeys :: IntMap a -> IntSet -> IntMap a
12001192
withoutKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1 p2 of
12011193
ABL -> binCheckLeft p1 (withoutKeys l1 t2) r1
12021194
ABR -> binCheckRight p1 l1 (withoutKeys r1 t2)
@@ -1205,49 +1197,26 @@ withoutKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1
12051197
EQL -> bin p1 (withoutKeys l1 l2) (withoutKeys r1 r2)
12061198
NOM -> t1
12071199
where
1208-
withoutKeys t1@(Bin p1 _ _) (IntSet.Tip p2 bm2) =
1209-
let px1 = unPrefix p1
1210-
minbit = bitmapOf (px1 .&. (px1-1))
1211-
lt_minbit = minbit - 1
1212-
maxbit = bitmapOf (px1 .|. (px1-1))
1213-
gt_maxbit = (-maxbit) `xor` maxbit
1214-
-- TODO(wrengr): should we manually inline/unroll 'updatePrefix'
1215-
-- and 'withoutBM' here, in order to avoid redundant case analyses?
1216-
in updatePrefix p2 t1 $ withoutBM (bm2 .|. lt_minbit .|. gt_maxbit)
1200+
withoutKeys t1@(Bin _ _ _) (IntSet.Tip p2 bm2) = withoutKeysTip t1 p2 bm2
12171201
withoutKeys t1@(Bin _ _ _) IntSet.Nil = t1
12181202
withoutKeys t1@(Tip k1 _) t2
12191203
| k1 `IntSet.member` t2 = Nil
12201204
| otherwise = t1
12211205
withoutKeys Nil _ = Nil
12221206

1223-
1224-
updatePrefix
1225-
:: IntSetPrefix -> IntMap a -> (IntMap a -> IntMap a) -> IntMap a
1226-
updatePrefix !kp t@(Bin p l r) f
1227-
| unPrefix p .&. IntSet.suffixBitMask /= 0 =
1228-
if unPrefix p .&. IntSet.prefixBitMask == kp then f t else t
1229-
| nomatch kp p = t
1230-
| left kp p = binCheckLeft p (updatePrefix kp l f) r
1231-
| otherwise = binCheckRight p l (updatePrefix kp r f)
1232-
updatePrefix kp t@(Tip kx _) f
1233-
| kx .&. IntSet.prefixBitMask == kp = f t
1234-
| otherwise = t
1235-
updatePrefix _ Nil _ = Nil
1236-
1237-
1238-
withoutBM :: IntSetBitMap -> IntMap a -> IntMap a
1239-
withoutBM 0 t = t
1240-
withoutBM bm (Bin p l r) =
1241-
let leftBits = bitmapOf (unPrefix p) - 1
1242-
bmL = bm .&. leftBits
1243-
bmR = bm `xor` bmL -- = (bm .&. complement leftBits)
1244-
in bin p (withoutBM bmL l) (withoutBM bmR r)
1245-
withoutBM bm t@(Tip k _)
1246-
-- TODO(wrengr): need we manually inline 'IntSet.Member' here?
1247-
| k `IntSet.member` IntSet.Tip (k .&. IntSet.prefixBitMask) bm = Nil
1248-
| otherwise = t
1249-
withoutBM _ Nil = Nil
1250-
1207+
withoutKeysTip :: IntMap a -> Int -> IntSet.BitMap -> IntMap a
1208+
withoutKeysTip t@(Bin p l r) !p2 !bm2
1209+
| IntSet.suffixOf (unPrefix p) /= 0 =
1210+
if IntSet.prefixOf (unPrefix p) == p2
1211+
then restrictBM t (complement bm2)
1212+
else t
1213+
| nomatch p2 p = t
1214+
| left p2 p = binCheckLeft p (withoutKeysTip l p2 bm2) r
1215+
| otherwise = binCheckRight p l (withoutKeysTip r p2 bm2)
1216+
withoutKeysTip t@(Tip kx _) !p2 !bm2
1217+
| IntSet.prefixOf kx == p2 && IntSet.bitmapOf kx .&. bm2 /= 0 = Nil
1218+
| otherwise = t
1219+
withoutKeysTip Nil !_ !_ = Nil
12511220

12521221
{--------------------------------------------------------------------
12531222
Intersection
@@ -1270,58 +1239,63 @@ intersection m1 m2
12701239
-- @
12711240
--
12721241
-- @since 0.5.8
1273-
restrictKeys :: IntMap a -> IntSet.IntSet -> IntMap a
1242+
restrictKeys :: IntMap a -> IntSet -> IntMap a
12741243
restrictKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1 p2 of
12751244
ABL -> restrictKeys l1 t2
12761245
ABR -> restrictKeys r1 t2
12771246
BAL -> restrictKeys t1 l2
12781247
BAR -> restrictKeys t1 r2
12791248
EQL -> bin p1 (restrictKeys l1 l2) (restrictKeys r1 r2)
12801249
NOM -> Nil
1281-
restrictKeys t1@(Bin p1 _ _) (IntSet.Tip p2 bm2) =
1282-
let px1 = unPrefix p1
1283-
minbit = bitmapOf (px1 .&. (px1-1))
1284-
ge_minbit = complement (minbit - 1)
1285-
maxbit = bitmapOf (px1 .|. (px1-1))
1286-
le_maxbit = maxbit .|. (maxbit - 1)
1287-
-- TODO(wrengr): should we manually inline/unroll 'lookupPrefix'
1288-
-- and 'restrictBM' here, in order to avoid redundant case analyses?
1289-
in restrictBM (bm2 .&. ge_minbit .&. le_maxbit) (lookupPrefix p2 t1)
1250+
restrictKeys t1@(Bin _ _ _) (IntSet.Tip p2 bm2) = restrictKeysTip t1 p2 bm2
12901251
restrictKeys (Bin _ _ _) IntSet.Nil = Nil
12911252
restrictKeys t1@(Tip k1 _) t2
12921253
| k1 `IntSet.member` t2 = t1
12931254
| otherwise = Nil
12941255
restrictKeys Nil _ = Nil
12951256

1296-
1297-
-- | \(O(\min(n,W))\). Restrict to the sub-map with all keys matching
1298-
-- a key prefix.
1299-
lookupPrefix :: IntSetPrefix -> IntMap a -> IntMap a
1300-
lookupPrefix !kp t@(Bin p l r)
1301-
| unPrefix p .&. IntSet.suffixBitMask /= 0 =
1302-
if unPrefix p .&. IntSet.prefixBitMask == kp then t else Nil
1303-
| nomatch kp p = Nil
1304-
| left kp p = lookupPrefix kp l
1305-
| otherwise = lookupPrefix kp r
1306-
lookupPrefix kp t@(Tip kx _)
1307-
| (kx .&. IntSet.prefixBitMask) == kp = t
1308-
| otherwise = Nil
1309-
lookupPrefix _ Nil = Nil
1310-
1311-
1312-
restrictBM :: IntSetBitMap -> IntMap a -> IntMap a
1313-
restrictBM 0 _ = Nil
1314-
restrictBM bm (Bin p l r) =
1315-
let leftBits = bitmapOf (unPrefix p) - 1
1316-
bmL = bm .&. leftBits
1317-
bmR = bm `xor` bmL -- = (bm .&. complement leftBits)
1318-
in bin p (restrictBM bmL l) (restrictBM bmR r)
1319-
restrictBM bm t@(Tip k _)
1320-
-- TODO(wrengr): need we manually inline 'IntSet.Member' here?
1321-
| k `IntSet.member` IntSet.Tip (k .&. IntSet.prefixBitMask) bm = t
1322-
| otherwise = Nil
1323-
restrictBM _ Nil = Nil
1324-
1257+
restrictKeysTip :: IntMap a -> Int -> IntSet.BitMap -> IntMap a
1258+
restrictKeysTip t@(Bin p l r) !p2 !bm2
1259+
| IntSet.suffixOf (unPrefix p) /= 0 =
1260+
if IntSet.prefixOf (unPrefix p) == p2
1261+
then restrictBM t bm2
1262+
else Nil
1263+
| nomatch p2 p = Nil
1264+
| left p2 p = restrictKeysTip l p2 bm2
1265+
| otherwise = restrictKeysTip r p2 bm2
1266+
restrictKeysTip t@(Tip kx _) !p2 !bm2
1267+
| IntSet.prefixOf kx == p2 && IntSet.bitmapOf kx .&. bm2 /= 0 = t
1268+
| otherwise = Nil
1269+
restrictKeysTip Nil !_ !_ = Nil
1270+
1271+
-- Must be called on an IntMap whose keys fit in the given IntSet BitMap's Tip.
1272+
-- Keeps keys that match the BitMap.
1273+
-- Returns early as an optimization, i.e. if the tree can be entirely kept or
1274+
-- discarded there is no need to recursively visit the children.
1275+
restrictBM :: IntMap a -> IntSet.BitMap -> IntMap a
1276+
restrictBM t@(Bin p l r) !bm
1277+
| bm' == 0 = Nil
1278+
| bm' == -1 = t
1279+
| otherwise = bin p (restrictBM l bm) (restrictBM r bm)
1280+
where
1281+
-- Here we care about the "submask" of bm corresponding the current Bin's
1282+
-- range. So we create bm', where this submask is at the lowest position and
1283+
-- and all other bits are set to the highest bit of the submask (using an
1284+
-- arithmetric shiftR). Now bm' is 0 when the submask is empty and -1 when
1285+
-- the submask is full.
1286+
px = IntSet.suffixOf (unPrefix p)
1287+
px1 = px - 1
1288+
min_ = px .&. px1
1289+
max_ = px .|. px1
1290+
sh = (wordSize - 1) - max_
1291+
bm' = (w2i bm `unsafeShiftL` sh) `unsafeShiftR` (sh + min_)
1292+
restrictBM t@(Tip k _) !bm
1293+
| IntSet.bitmapOf k .&. bm /= 0 = t
1294+
| otherwise = Nil
1295+
restrictBM Nil !_ = Nil
1296+
1297+
w2i :: Word -> Int
1298+
w2i = fromIntegral
13251299

13261300
-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\).
13271301
-- The intersection with a combining function.
@@ -3194,7 +3168,7 @@ assocs = toAscList
31943168
-- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5]
31953169
-- > keysSet empty == Data.IntSet.empty
31963170

3197-
keysSet :: IntMap a -> IntSet.IntSet
3171+
keysSet :: IntMap a -> IntSet
31983172
keysSet Nil = IntSet.Nil
31993173
keysSet (Tip kx _) = IntSet.singleton kx
32003174
keysSet (Bin p l r)
@@ -3212,7 +3186,7 @@ keysSet (Bin p l r)
32123186
-- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")]
32133187
-- > fromSet undefined Data.IntSet.empty == empty
32143188

3215-
fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a
3189+
fromSet :: (Key -> a) -> IntSet -> IntMap a
32163190
fromSet _ IntSet.Nil = Nil
32173191
fromSet f (IntSet.Bin p l r) = Bin p (fromSet f l) (fromSet f r)
32183192
fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1)

‎containers/src/Data/IntSet/Internal.hs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,8 @@ module Data.IntSet.Internal (
183183
-- * Internals
184184
, suffixBitMask
185185
, prefixBitMask
186+
, prefixOf
187+
, suffixOf
186188
, bitmapOf
187189
) where
188190

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /