Skip to content

Commit 56c99f3

Browse files
meooow25meooow25
authored
Add antitone functions for IntSet and IntMap (#874)
* Add IntMap antitone funcs * Add IntSet antitone funcs * Add antitone property tests * Add antitone benchmarks
1 parent aebfe5e commit 56c99f3

File tree

10 files changed

+296
-1
lines changed

10 files changed

+296
-1
lines changed

‎containers-tests/benchmarks/IntMap.hs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ main = do
5050
, bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems
5151
, bench "minView" $ whnf (maybe 0 (\((k,v), m) -> k+v+M.size m) . M.minViewWithKey)
5252
(M.fromList $ zip [1..10] [1..10])
53+
, bench "spanAntitone" $ whnf (M.spanAntitone (<key_mid)) m
5354
]
5455
where
5556
elems = elems_hits
@@ -64,6 +65,7 @@ main = do
6465
keys'' = fmap (* 2) [1..2^12]
6566
mixedKeys = interleave keys keys'
6667
values = [1..2^12]
68+
key_mid = 2^11
6769
--------------------------------------------------------
6870
sum k v1 v2 = k + v1 + v2
6971
consPair k v xs = (k, v) : xs

‎containers-tests/benchmarks/IntSet.hs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ main = do
1919
let s = IS.fromAscList elems :: IS.IntSet
2020
s_even = IS.fromAscList elems_even :: IS.IntSet
2121
s_odd = IS.fromAscList elems_odd :: IS.IntSet
22-
evaluate $ rnf [s, s_even, s_odd]
22+
s_sparse = IS.fromAscList elems_sparse :: IS.IntSet
23+
evaluate $ rnf [s, s_even, s_odd, s_sparse]
2324
defaultMain
2425
[ bench "member" $ whnf (member elems) s
2526
, bench "insert" $ whnf (ins elems) IS.empty
@@ -47,11 +48,16 @@ main = do
4748
$ whnf (num_transitions . det 2 0) $ hard_nfa 1 16
4849
, bench "instanceOrd:sparse" -- many Bin, each Tip is singleton
4950
$ whnf (num_transitions . det 2 0) $ hard_nfa 1111 16
51+
, bench "spanAntitone:dense" $ whnf (IS.spanAntitone (<elem_mid)) s
52+
, bench "spanAntitone:sparse" $ whnf (IS.spanAntitone (<elem_sparse_mid)) s_sparse
5053
]
5154
where
5255
elems = [1..2^12]
5356
elems_even = [2,4..2^12]
5457
elems_odd = [1,3..2^12]
58+
elem_mid = 2^11
59+
elems_sparse = map (*64) elems -- when built into a map, each Tip is a singleton
60+
elem_sparse_mid = 64 * elem_mid
5561

5662
member :: [Int] -> IS.IntSet -> Int
5763
member xs s = foldl' (\n x -> if IS.member x s then n + 1 else n) 0 xs

‎containers-tests/tests/intmap-properties.hs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,9 @@ main = defaultMain $ testGroup "intmap-properties"
180180
, testProperty "deleteMax" prop_deleteMaxModel
181181
, testProperty "filter" prop_filter
182182
, testProperty "partition" prop_partition
183+
, testProperty "takeWhileAntitone" prop_takeWhileAntitone
184+
, testProperty "dropWhileAntitone" prop_dropWhileAntitone
185+
, testProperty "spanAntitone" prop_spanAntitone
183186
, testProperty "map" prop_map
184187
, testProperty "fmap" prop_fmap
185188
, testProperty "mapkeys" prop_mapkeys
@@ -1469,6 +1472,26 @@ prop_partition p ys = length ys > 0 ==>
14691472
m === let (a,b) = (List.partition (apply p . snd) xs)
14701473
in (fromList a, fromList b)
14711474

1475+
prop_takeWhileAntitone :: Int -> [(Int, Int)] -> Property
1476+
prop_takeWhileAntitone x ys =
1477+
let l = takeWhileAntitone (<x) (fromList ys)
1478+
in valid l .&&.
1479+
l === fromList (List.filter ((<x) . fst) ys)
1480+
1481+
prop_dropWhileAntitone :: Int -> [(Int, Int)] -> Property
1482+
prop_dropWhileAntitone x ys =
1483+
let r = dropWhileAntitone (<x) (fromList ys)
1484+
in valid r .&&.
1485+
r === fromList (List.filter ((>=x) . fst) ys)
1486+
1487+
prop_spanAntitone :: Int -> [(Int, Int)] -> Property
1488+
prop_spanAntitone x ys =
1489+
let (l, r) = spanAntitone (<x) (fromList ys)
1490+
in valid l .&&.
1491+
valid r .&&.
1492+
l === fromList (List.filter ((<x) . fst) ys) .&&.
1493+
r === fromList (List.filter ((>=x) . fst) ys)
1494+
14721495
prop_map :: Fun Int Int -> [(Int, Int)] -> Property
14731496
prop_map f ys = length ys > 0 ==>
14741497
let xs = List.nubBy ((==) `on` fst) ys

‎containers-tests/tests/intset-properties.hs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ main = defaultMain $ testGroup "intset-properties"
7070
, testProperty "prop_splitRoot" prop_splitRoot
7171
, testProperty "prop_partition" prop_partition
7272
, testProperty "prop_filter" prop_filter
73+
, testProperty "takeWhileAntitone" prop_takeWhileAntitone
74+
, testProperty "dropWhileAntitone" prop_dropWhileAntitone
75+
, testProperty "spanAntitone" prop_spanAntitone
7376
, testProperty "prop_bitcount" prop_bitcount
7477
, testProperty "prop_alterF_list" prop_alterF_list
7578
, testProperty "prop_alterF_const" prop_alterF_const
@@ -420,6 +423,26 @@ prop_filter s i =
420423
valid evens .&&.
421424
parts === (odds, evens)
422425

426+
prop_takeWhileAntitone :: Int -> [Int] -> Property
427+
prop_takeWhileAntitone x ys =
428+
let l = takeWhileAntitone (<x) (fromList ys)
429+
in valid l .&&.
430+
l === fromList (List.filter (<x) ys)
431+
432+
prop_dropWhileAntitone :: Int -> [Int] -> Property
433+
prop_dropWhileAntitone x ys =
434+
let r = dropWhileAntitone (<x) (fromList ys)
435+
in valid r .&&.
436+
r === fromList (List.filter (>=x) ys)
437+
438+
prop_spanAntitone :: Int -> [Int] -> Property
439+
prop_spanAntitone x ys =
440+
let (l, r) = spanAntitone (<x) (fromList ys)
441+
in valid l .&&.
442+
valid r .&&.
443+
l === fromList (List.filter (<x) ys) .&&.
444+
r === fromList (List.filter (>=x) ys)
445+
423446
prop_bitcount :: Int -> Word -> Bool
424447
prop_bitcount a w = bitcount_orig a w == bitcount_new a w
425448
where

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

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,10 @@ module Data.IntMap.Internal (
228228
, partition
229229
, partitionWithKey
230230

231+
, takeWhileAntitone
232+
, dropWhileAntitone
233+
, spanAntitone
234+
231235
, mapMaybe
232236
, mapMaybeWithKey
233237
, mapEither
@@ -2614,6 +2618,101 @@ partitionWithKey predicate0 t0 = toPair $ go predicate0 t0
26142618
| otherwise -> (Nil :*: t)
26152619
Nil -> (Nil :*: Nil)
26162620

2621+
-- | \(O(\min(n,W))\). Take while a predicate on the keys holds.
2622+
-- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@.
2623+
-- See note at 'spanAntitone'.
2624+
--
2625+
-- @
2626+
-- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' (p . fst) . 'toList'
2627+
-- takeWhileAntitone p = 'filterWithKey' (\\k _ -> p k)
2628+
-- @
2629+
--
2630+
-- @since FIXME
2631+
takeWhileAntitone :: (Key -> Bool) -> IntMap a -> IntMap a
2632+
takeWhileAntitone predicate t =
2633+
case t of
2634+
Bin p m l r
2635+
| m < 0 ->
2636+
if predicate 0 -- handle negative numbers.
2637+
then bin p m (go predicate l) r
2638+
else go predicate r
2639+
_ -> go predicate t
2640+
where
2641+
go predicate' (Bin p m l r)
2642+
| predicate' $! p+m = bin p m l (go predicate' r)
2643+
| otherwise = go predicate' l
2644+
go predicate' t'@(Tip ky _)
2645+
| predicate' ky = t'
2646+
| otherwise = Nil
2647+
go _ Nil = Nil
2648+
2649+
-- | \(O(\min(n,W))\). Drop while a predicate on the keys holds.
2650+
-- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@.
2651+
-- See note at 'spanAntitone'.
2652+
--
2653+
-- @
2654+
-- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' (p . fst) . 'toList'
2655+
-- dropWhileAntitone p = 'filterWithKey' (\\k _ -> not (p k))
2656+
-- @
2657+
--
2658+
-- @since FIXME
2659+
dropWhileAntitone :: (Key -> Bool) -> IntMap a -> IntMap a
2660+
dropWhileAntitone predicate t =
2661+
case t of
2662+
Bin p m l r
2663+
| m < 0 ->
2664+
if predicate 0 -- handle negative numbers.
2665+
then go predicate l
2666+
else bin p m l (go predicate r)
2667+
_ -> go predicate t
2668+
where
2669+
go predicate' (Bin p m l r)
2670+
| predicate' $! p+m = go predicate' r
2671+
| otherwise = bin p m (go predicate' l) r
2672+
go predicate' t'@(Tip ky _)
2673+
| predicate' ky = Nil
2674+
| otherwise = t'
2675+
go _ Nil = Nil
2676+
2677+
-- | \(O(\min(n,W))\). Divide a map at the point where a predicate on the keys stops holding.
2678+
-- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@.
2679+
--
2680+
-- @
2681+
-- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs)
2682+
-- spanAntitone p xs = 'partitionWithKey' (\\k _ -> p k) xs
2683+
-- @
2684+
--
2685+
-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the map
2686+
-- at some /unspecified/ point.
2687+
--
2688+
-- @since FIXME
2689+
spanAntitone :: (Key -> Bool) -> IntMap a -> (IntMap a, IntMap a)
2690+
spanAntitone predicate t =
2691+
case t of
2692+
Bin p m l r
2693+
| m < 0 ->
2694+
if predicate 0 -- handle negative numbers.
2695+
then
2696+
case go predicate l of
2697+
(lt :*: gt) ->
2698+
let !lt' = bin p m lt r
2699+
in (lt', gt)
2700+
else
2701+
case go predicate r of
2702+
(lt :*: gt) ->
2703+
let !gt' = bin p m l gt
2704+
in (lt, gt')
2705+
_ -> case go predicate t of
2706+
(lt :*: gt) -> (lt, gt)
2707+
where
2708+
go predicate' (Bin p m l r)
2709+
| predicate' $! p+m = case go predicate' r of (lt :*: gt) -> bin p m l lt :*: gt
2710+
| otherwise = case go predicate' l of (lt :*: gt) -> lt :*: bin p m gt r
2711+
go predicate' t'@(Tip ky _)
2712+
| predicate' ky = (t' :*: Nil)
2713+
| otherwise = (Nil :*: t')
2714+
go _ Nil = (Nil :*: Nil)
2715+
26172716
-- | \(O(n)\). Map values and collect the 'Just' results.
26182717
--
26192718
-- > let f x = if x == "a" then Just "new a" else Nothing

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,10 @@ module Data.IntMap.Lazy (
198198
, partition
199199
, partitionWithKey
200200

201+
, takeWhileAntitone
202+
, dropWhileAntitone
203+
, spanAntitone
204+
201205
, mapMaybe
202206
, mapMaybeWithKey
203207
, mapEither

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,10 @@ module Data.IntMap.Strict (
217217
, partition
218218
, partitionWithKey
219219

220+
, takeWhileAntitone
221+
, dropWhileAntitone
222+
, spanAntitone
223+
220224
, mapMaybe
221225
, mapMaybeWithKey
222226
, mapEither

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,10 @@ module Data.IntMap.Strict.Internal (
217217
, partition
218218
, partitionWithKey
219219

220+
, takeWhileAntitone
221+
, dropWhileAntitone
222+
, spanAntitone
223+
220224
, mapMaybe
221225
, mapMaybeWithKey
222226
, mapEither
@@ -327,6 +331,9 @@ import Data.IntMap.Internal
327331
, null
328332
, partition
329333
, partitionWithKey
334+
, takeWhileAntitone
335+
, dropWhileAntitone
336+
, spanAntitone
330337
, restrictKeys
331338
, size
332339
, split

‎containers/src/Data/IntSet.hs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,11 @@ module Data.IntSet (
111111
-- * Filter
112112
, IS.filter
113113
, partition
114+
115+
, takeWhileAntitone
116+
, dropWhileAntitone
117+
, spanAntitone
118+
114119
, split
115120
, splitMember
116121
, splitRoot

0 commit comments

Comments
 (0)