From dc1903c8b7fcf9787fbce514ebb07ac6a329aa16 Mon Sep 17 00:00:00 2001 From: Reed Mullanix Date: Fri, 1 Oct 2021 20:16:00 -0700 Subject: [PATCH 1/6] Add 'between' for Data.Map --- containers/src/Data/Map/Internal.hs | 13 +++++++++++++ containers/src/Data/Map/Lazy.hs | 1 + containers/src/Data/Map/Strict.hs | 1 + containers/src/Data/Map/Strict/Internal.hs | 2 ++ 4 files changed, 17 insertions(+) diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index a2edf6914..bbcb23146 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -299,6 +299,7 @@ module Data.Map.Internal ( , withoutKeys , partition , partitionWithKey + , between , mapMaybe , mapMaybeWithKey @@ -2983,6 +2984,18 @@ spanAntitone p0 m = toPair (go p0 m) | p kx = let u :*: v = go p r in link kx x l u :*: v | otherwise = let u :*: v = go p l in u :*: link kx x v r +-- | /O(log n)/. Filter a map such that its keys lie between two values (inclusive). +-- +-- > between 2 4 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == fromList [(3, "a"), (4, "c"), (2, "e")] +-- > between 4 2 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == empty +-- > between 2 2 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == singleton 2 "e" +between :: (Ord k) => k -> k -> Map k a -> Map k a +between _ _ Tip = Tip +between lo hi (Bin _ kx x l r) + | kx < lo = between lo hi r + | kx > hi = between lo hi l + | otherwise = link kx x (between lo hi l) (between lo hi r) + -- | /O(n)/. Partition the map according to a predicate. The first -- map contains all elements that satisfy the predicate, the second all -- elements that fail the predicate. See also 'split'. diff --git a/containers/src/Data/Map/Lazy.hs b/containers/src/Data/Map/Lazy.hs index 6b05603a1..d85161d93 100644 --- a/containers/src/Data/Map/Lazy.hs +++ b/containers/src/Data/Map/Lazy.hs @@ -222,6 +222,7 @@ module Data.Map.Lazy ( , withoutKeys , partition , partitionWithKey + , between , takeWhileAntitone , dropWhileAntitone , spanAntitone diff --git a/containers/src/Data/Map/Strict.hs b/containers/src/Data/Map/Strict.hs index 8eea8c329..cf5045999 100644 --- a/containers/src/Data/Map/Strict.hs +++ b/containers/src/Data/Map/Strict.hs @@ -238,6 +238,7 @@ module Data.Map.Strict , withoutKeys , partition , partitionWithKey + , between , takeWhileAntitone , dropWhileAntitone diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs index d8b5325ba..4817bf6a4 100644 --- a/containers/src/Data/Map/Strict/Internal.hs +++ b/containers/src/Data/Map/Strict/Internal.hs @@ -256,6 +256,7 @@ module Data.Map.Strict.Internal , takeWhileAntitone , dropWhileAntitone , spanAntitone + , between , mapMaybe , mapMaybeWithKey @@ -393,6 +394,7 @@ import Data.Map.Internal , null , partition , partitionWithKey + , between , restrictKeys , size , spanAntitone From c1e0d6621d2119149046a0347507e82b1578c8e0 Mon Sep 17 00:00:00 2001 From: Reed Mullanix Date: Fri, 1 Oct 2021 20:26:54 -0700 Subject: [PATCH 2/6] Add property test for 'Map.between' --- containers-tests/tests/map-properties.hs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs index fc0f39909..45379f664 100644 --- a/containers-tests/tests/map-properties.hs +++ b/containers-tests/tests/map-properties.hs @@ -214,6 +214,7 @@ main = defaultMain $ testGroup "map-properties" , testProperty "deleteMax" prop_deleteMaxModel , testProperty "filter" prop_filter , testProperty "partition" prop_partition + , testProperty "between" prop_between , testProperty "map" prop_map , testProperty "fmap" prop_fmap , testProperty "mapkeys" prop_mapkeys @@ -1397,6 +1398,12 @@ prop_splitAt n xs = valid taken .&&. where (taken, dropped) = splitAt n xs +prop_between :: Int -> Int -> [(Int, Int)] -> Property +prop_between lo hi xs' = valid tw .&&. tw === (filterWithKey (\k _ -> lo <= k && k <= hi) xs) + where + xs = fromList xs' + tw = between lo hi xs + prop_takeWhileAntitone :: [(Either Int Int, Int)] -> Property prop_takeWhileAntitone xs' = valid tw .&&. (tw === filterWithKey (\k _ -> isLeft k) xs) where From e7688c6580fe82dd39b32195a15828c7dc743ac5 Mon Sep 17 00:00:00 2001 From: Reed Mullanix Date: Fri, 1 Oct 2021 20:42:25 -0700 Subject: [PATCH 3/6] Add between for 'Data.Set' --- containers/src/Data/Set.hs | 1 + containers/src/Data/Set/Internal.hs | 12 ++++++++++++ 2 files changed, 13 insertions(+) diff --git a/containers/src/Data/Set.hs b/containers/src/Data/Set.hs index e3f7281b6..bf144aab6 100644 --- a/containers/src/Data/Set.hs +++ b/containers/src/Data/Set.hs @@ -123,6 +123,7 @@ module Data.Set ( , split , splitMember , splitRoot + , between -- * Indexed , lookupIndex diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs index ffcca01ed..067b147e2 100644 --- a/containers/src/Data/Set/Internal.hs +++ b/containers/src/Data/Set/Internal.hs @@ -168,6 +168,7 @@ module Data.Set.Internal ( , split , splitMember , splitRoot + , between -- * Indexed , lookupIndex @@ -1535,6 +1536,17 @@ spanAntitone p0 m = toPair (go p0 m) | p x = let u :*: v = go p r in link x l u :*: v | otherwise = let u :*: v = go p l in u :*: link x v r +-- | /O(log n)/. Filter a map set that its values lie between a lower and upper bound (inclusive). +-- +-- > between 2 4 (fromList [3, 5, 4, 1 2]) == fromList [3, 4, 2] +-- > between 4 2 (fromList [3, 5, 4, 1 2]) == empty +-- > between 2 2 (fromList [3, 5, 4, 1 2]) == singleton 2 +between :: (Ord a) => a -> a -> Set a -> Set a +between _ _ Tip = Tip +between lo hi (Bin _ x l r) + | x < lo = between lo hi r + | x > hi = between lo hi l + | otherwise = link x (between lo hi l) (between lo hi r) {-------------------------------------------------------------------- Utility functions that maintain the balance properties of the tree. From 2933680351e70415da779a2473d0b6e38983900f Mon Sep 17 00:00:00 2001 From: Reed Mullanix Date: Fri, 1 Oct 2021 20:42:34 -0700 Subject: [PATCH 4/6] Add property test for 'Set.between' --- containers-tests/tests/set-properties.hs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/containers-tests/tests/set-properties.hs b/containers-tests/tests/set-properties.hs index 2e8ed6abc..92958ea1a 100644 --- a/containers-tests/tests/set-properties.hs +++ b/containers-tests/tests/set-properties.hs @@ -100,6 +100,7 @@ main = defaultMain $ testGroup "set-properties" , testProperty "takeWhileAntitone" prop_takeWhileAntitone , testProperty "dropWhileAntitone" prop_dropWhileAntitone , testProperty "spanAntitone" prop_spanAntitone + , testProperty "between" prop_between , testProperty "take" prop_take , testProperty "drop" prop_drop , testProperty "splitAt" prop_splitAt @@ -672,6 +673,12 @@ prop_spanAntitone xs' = valid tw .&&. valid dw xs = fromList xs' (tw, dw) = spanAntitone isLeft xs +prop_between :: Int -> Int -> [Int] -> Property +prop_between lo hi xs' = valid tw .&&. tw === (filter (\x -> lo <= x && x <= hi) xs) + where + xs = fromList xs' + tw = between lo hi xs + prop_powerSet :: Set Int -> Property prop_powerSet xs = valid ps .&&. ps === ps' where From ca13f3235132e805bd8e29911428d9d20970a269 Mon Sep 17 00:00:00 2001 From: Reed Mullanix Date: Fri, 1 Oct 2021 20:42:47 -0700 Subject: [PATCH 5/6] Update the docs for 'Map.between' for some added clarity. --- containers/src/Data/Map/Internal.hs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index bbcb23146..e76a34dbb 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -2984,11 +2984,12 @@ spanAntitone p0 m = toPair (go p0 m) | p kx = let u :*: v = go p r in link kx x l u :*: v | otherwise = let u :*: v = go p l in u :*: link kx x v r --- | /O(log n)/. Filter a map such that its keys lie between two values (inclusive). +-- | /O(log n)/. Filter a map such that its keys lie between a lower and an upper bound (inclusive). -- -- > between 2 4 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == fromList [(3, "a"), (4, "c"), (2, "e")] -- > between 4 2 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == empty -- > between 2 2 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == singleton 2 "e" + between :: (Ord k) => k -> k -> Map k a -> Map k a between _ _ Tip = Tip between lo hi (Bin _ kx x l r) From fc9669b212cc9b2ad297dbd6f2ac0f7564bfe24c Mon Sep 17 00:00:00 2001 From: Reed Mullanix Date: Fri, 1 Oct 2021 21:06:32 -0700 Subject: [PATCH 6/6] Fix silly optimization oversight --- containers/src/Data/Map/Internal.hs | 2 +- containers/src/Data/Set/Internal.hs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs index e76a34dbb..dcf8ab3ba 100644 --- a/containers/src/Data/Map/Internal.hs +++ b/containers/src/Data/Map/Internal.hs @@ -2995,7 +2995,7 @@ between _ _ Tip = Tip between lo hi (Bin _ kx x l r) | kx < lo = between lo hi r | kx > hi = between lo hi l - | otherwise = link kx x (between lo hi l) (between lo hi r) + | otherwise = link kx x (dropWhileAntitone (< lo) l) (takeWhileAntitone (<= hi) r) -- | /O(n)/. Partition the map according to a predicate. The first -- map contains all elements that satisfy the predicate, the second all diff --git a/containers/src/Data/Set/Internal.hs b/containers/src/Data/Set/Internal.hs index 067b147e2..8cd34c2e3 100644 --- a/containers/src/Data/Set/Internal.hs +++ b/containers/src/Data/Set/Internal.hs @@ -1546,7 +1546,7 @@ between _ _ Tip = Tip between lo hi (Bin _ x l r) | x < lo = between lo hi r | x > hi = between lo hi l - | otherwise = link x (between lo hi l) (between lo hi r) + | otherwise = link x (dropWhileAntitone (< lo) l) (takeWhileAntitone (<= hi) r) {-------------------------------------------------------------------- Utility functions that maintain the balance properties of the tree.