Skip to content

Commit dcfb408

Browse files
meooow25meooow25
authored
Add fromDescList for IntSet and IntMap (#1194)
1 parent d3cedea commit dcfb408

File tree

9 files changed

+161
-4
lines changed

9 files changed

+161
-4
lines changed

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,7 @@ main = defaultMain $ testGroup "intmap-properties"
266266
, testProperty "fromAscListWith" prop_fromAscListWith
267267
, testProperty "fromAscListWithKey" prop_fromAscListWithKey
268268
, testProperty "fromDistinctAscList" prop_fromDistinctAscList
269+
, testProperty "fromDescList" prop_fromDescList
269270
, testProperty "fromListWith" prop_fromListWith
270271
, testProperty "fromListWithKey" prop_fromListWithKey
271272
, testProperty "compareSize" prop_compareSize
@@ -2182,6 +2183,14 @@ prop_fromDistinctAscList kxs =
21822183
List.sortBy (comparing fst) kxs
21832184
t = fromDistinctAscList nubSortedKxs
21842185

2186+
prop_fromDescList :: [(Int, A)] -> Property
2187+
prop_fromDescList kxs =
2188+
valid t .&&.
2189+
t === fromList sortedKxs
2190+
where
2191+
sortedKxs = List.sortBy (comparing (Down . fst)) kxs
2192+
t = fromDescList sortedKxs
2193+
21852194
prop_fromListWith :: Fun (A, A) A -> [(Int, A)] -> Property
21862195
prop_fromListWith f kxs =
21872196
valid m .&&.

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import Data.Functor.Identity (Identity(..))
1414
import qualified Data.List as List
1515
import qualified Data.List.NonEmpty as NE
1616
import Data.Maybe (catMaybes, mapMaybe)
17-
import Data.Ord (comparing)
17+
import Data.Ord (Down(..), comparing)
1818
import Test.ChasingBottoms.IsBottom
1919
import Test.Tasty (TestTree, defaultMain, testGroup)
2020
import Test.Tasty.QuickCheck (testProperty)
@@ -193,6 +193,17 @@ prop_lazyFromDistinctAscList kvs = isNotBottomProp (L.fromDistinctAscList kvs')
193193
where
194194
kvs' = uniqOn fst $ List.sortBy (comparing fst) (coerce kvs) :: [(Key, A)]
195195

196+
prop_strictFromDescList :: [(Key, Bot A)] -> Property
197+
prop_strictFromDescList kvs =
198+
isBottom (M.fromDescList kvs') === isBottom (M.fromList kvs')
199+
where
200+
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(Key, A)]
201+
202+
prop_lazyFromDescList :: [(Key, Bot A)] -> Property
203+
prop_lazyFromDescList kvs = isNotBottomProp (L.fromDescList kvs')
204+
where
205+
kvs' = List.sortBy (comparing (Down . fst)) (coerce kvs) :: [(Key, A)]
206+
196207
prop_strictInsert :: Key -> Bot A -> IntMap A -> Property
197208
prop_strictInsert k (Bot x) m = isBottom (M.insert k x m) === isBottom x
198209

@@ -1049,6 +1060,7 @@ tests =
10491060
, testPropStrictLazy "fromAscListWith" prop_strictFromAscListWith prop_lazyFromAscListWith
10501061
, testPropStrictLazy "fromAscListWithKey" prop_strictFromAscListWithKey prop_lazyFromAscListWithKey
10511062
, testPropStrictLazy "fromDistinctAscList" prop_strictFromDistinctAscList prop_lazyFromDistinctAscList
1063+
, testPropStrictLazy "fromDescList" prop_strictFromDescList prop_lazyFromDescList
10521064
, testPropStrictLazy "insert" prop_strictInsert prop_lazyInsert
10531065
, testPropStrictLazy "insertWith" prop_strictInsertWith prop_lazyInsertWith
10541066
, testPropStrictLazy "insertWithKey" prop_strictInsertWithKey prop_lazyInsertWithKey

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ main = defaultMain $ testGroup "intset-properties"
100100
, testProperty "deleteMax" prop_deleteMax
101101
, testProperty "fromAscList" prop_fromAscList
102102
, testProperty "fromDistinctAscList" prop_fromDistinctAscList
103+
, testProperty "fromDescList" prop_fromDescList
103104
, testProperty "compareSize" prop_compareSize
104105
, testLaws $ Laws.eqLaws (Proxy :: Proxy IntSet)
105106
, testLaws $ Laws.ordLaws (Proxy :: Proxy IntSet)
@@ -583,5 +584,14 @@ prop_fromDistinctAscList xs =
583584
nubSortedXs = List.map NE.head $ NE.group $ sort xs
584585
t = fromDistinctAscList nubSortedXs
585586

587+
prop_fromDescList :: [Int] -> Property
588+
prop_fromDescList xs =
589+
valid t .&&.
590+
toList t === nubSortedXs
591+
where
592+
sortedXs = sort xs
593+
nubSortedXs = List.map NE.head $ NE.group sortedXs
594+
t = fromDescList (reverse sortedXs)
595+
586596
prop_compareSize :: IntSet -> Int -> Property
587597
prop_compareSize t c = compareSize t c === compare (size t) c

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,7 @@ module Data.IntMap.Internal (
237237
, fromAscListWith
238238
, fromAscListWithKey
239239
, fromDistinctAscList
240+
, fromDescList
240241

241242
-- * Filter
242243
, filter
@@ -296,6 +297,9 @@ module Data.IntMap.Internal (
296297
, Stack(..)
297298
, ascLinkTop
298299
, ascLinkAll
300+
, descInsert
301+
, descLinkTop
302+
, descLinkAll
299303
, IntMapBuilder(..)
300304
, BStack(..)
301305
, emptyB
@@ -3605,6 +3609,22 @@ fromDistinctAscList :: [(Key,a)] -> IntMap a
36053609
fromDistinctAscList = fromAscList
36063610
{-# INLINE fromDistinctAscList #-} -- Inline for list fusion
36073611

3612+
-- | \(O(n)\). Build a map from a list of key\/value pairs where
3613+
-- the keys are in descending order.
3614+
--
3615+
-- __Warning__: This function should be used only if the keys are in
3616+
-- non-increasing order. This precondition is not checked. Use 'fromList' if the
3617+
-- precondition may not hold.
3618+
--
3619+
-- > fromDescList [(5,"a"), (3,"b")] == fromList [(3,"b"), (5,"a")]
3620+
-- > fromDescList [(5,"a"), (5,"b"), (3,"b")] == fromList [(3,"b"), (5,"b")]
3621+
--
3622+
-- @since FIXME
3623+
fromDescList :: [(Key,a)] -> IntMap a
3624+
fromDescList xs =
3625+
descLinkAll (Foldable.foldl' (\s (ky, y) -> descInsert ky y s) MSNada xs)
3626+
{-# INLINE fromDescList #-} -- Inline for list fusion
3627+
36083628
data Stack a
36093629
= Nada
36103630
| Push {-# UNPACK #-} !Int !(IntMap a) !(Stack a)
@@ -3647,6 +3667,40 @@ ascLinkStack stk !rk r = case stk of
36473667
where
36483668
p = mask rk m
36493669

3670+
-- Insert an entry. The key must be <= the last inserted key. If it is equal
3671+
-- to the previous key, the previous value is replaced.
3672+
descInsert :: Int -> a -> MonoState a -> MonoState a
3673+
descInsert !ky y s = case s of
3674+
MSNada -> MSPush ky y Nada
3675+
MSPush kx x stk
3676+
| kx == ky -> MSPush ky y stk
3677+
| otherwise -> let m = branchMask kx ky
3678+
in MSPush ky y (descLinkTop kx (Tip kx x) m stk)
3679+
{-# INLINE descInsert #-}
3680+
3681+
descLinkTop :: Int -> IntMap a -> Int -> Stack a -> Stack a
3682+
descLinkTop !lk l !lm stk = case stk of
3683+
Nada -> Push lm l stk
3684+
Push m r stk'
3685+
| i2w m < i2w lm -> let p = mask lk m
3686+
in descLinkTop lk (Bin p l r) lm stk'
3687+
| otherwise -> Push lm l stk
3688+
3689+
descLinkAll :: MonoState a -> IntMap a
3690+
descLinkAll s = case s of
3691+
MSNada -> Nil
3692+
MSPush kx x stk -> descLinkStack kx (Tip kx x) stk
3693+
{-# INLINABLE descLinkAll #-}
3694+
3695+
descLinkStack :: Int -> IntMap a -> Stack a -> IntMap a
3696+
descLinkStack !lk l stk = case stk of
3697+
Nada -> l
3698+
Push m r stk'
3699+
| signBranch p -> Bin p r l
3700+
| otherwise -> descLinkStack lk (Bin p l r) stk'
3701+
where
3702+
p = mask lk m
3703+
36503704
{--------------------------------------------------------------------
36513705
IntMapBuilder
36523706
--------------------------------------------------------------------}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,11 +110,12 @@ module Data.IntMap.Lazy (
110110
, fromListWith
111111
, fromListWithKey
112112

113-
-- ** From Ascending Lists
113+
-- ** From Ordered Lists
114114
, fromAscList
115115
, fromAscListWith
116116
, fromAscListWithKey
117117
, fromDistinctAscList
118+
, fromDescList
118119

119120
-- * Insertion
120121
, insert

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,11 +128,12 @@ module Data.IntMap.Strict (
128128
, fromListWith
129129
, fromListWithKey
130130

131-
-- ** From Ascending Lists
131+
-- ** From Ordered Lists
132132
, fromAscList
133133
, fromAscListWith
134134
, fromAscListWithKey
135135
, fromDistinctAscList
136+
, fromDescList
136137

137138
-- * Insertion
138139
, insert

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

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,12 @@ module Data.IntMap.Strict.Internal (
6969
, fromListWith
7070
, fromListWithKey
7171

72-
-- ** From Ascending Lists
72+
-- ** From Ordered Lists
7373
, fromAscList
7474
, fromAscListWith
7575
, fromAscListWithKey
7676
, fromDistinctAscList
77+
, fromDescList
7778

7879
-- * Insertion
7980
, insert
@@ -246,6 +247,8 @@ import Data.IntMap.Internal
246247
, Stack(..)
247248
, ascLinkTop
248249
, ascLinkAll
250+
, descInsert
251+
, descLinkAll
249252
, IntMapBuilder(..)
250253
, BStack(..)
251254
, emptyB
@@ -1269,6 +1272,22 @@ fromDistinctAscList :: [(Key,a)] -> IntMap a
12691272
fromDistinctAscList = fromAscList
12701273
{-# INLINE fromDistinctAscList #-} -- Inline for list fusion
12711274

1275+
-- | \(O(n)\). Build a map from a list of key\/value pairs where
1276+
-- the keys are in descending order.
1277+
--
1278+
-- __Warning__: This function should be used only if the keys are in
1279+
-- non-increasing order. This precondition is not checked. Use 'fromList' if the
1280+
-- precondition may not hold.
1281+
--
1282+
-- > fromDescList [(5,"a"), (3,"b")] == fromList [(3,"b"), (5,"a")]
1283+
-- > fromDescList [(5,"a"), (5,"b"), (3,"b")] == fromList [(3,"b"), (5,"b")]
1284+
--
1285+
-- @since FIXME
1286+
fromDescList :: [(Key,a)] -> IntMap a
1287+
fromDescList xs =
1288+
descLinkAll (Foldable.foldl' (\s (!ky, !y) -> descInsert ky y s) MSNada xs)
1289+
{-# INLINE fromDescList #-} -- Inline for list fusion
1290+
12721291
{--------------------------------------------------------------------
12731292
IntMapBuilder
12741293
--------------------------------------------------------------------}

‎containers/src/Data/IntSet.hs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ module Data.IntSet (
103103
, fromRange
104104
, fromAscList
105105
, fromDistinctAscList
106+
, fromDescList
106107

107108
-- * Insertion
108109
, insert

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ module Data.IntSet.Internal (
179179
, toDescList
180180
, fromAscList
181181
, fromDistinctAscList
182+
, fromDescList
182183

183184
-- * Debugging
184185
, showTree
@@ -1530,6 +1531,19 @@ fromDistinctAscList :: [Key] -> IntSet
15301531
fromDistinctAscList = fromAscList
15311532
{-# INLINE fromDistinctAscList #-} -- Inline for list fusion
15321533

1534+
-- | \(O(n)\). Build a set from an descending list of elements.
1535+
--
1536+
-- __Warning__: This function should be used only if the elements are in
1537+
-- non-increasing order. This precondition is not checked. Use 'fromList' if the
1538+
-- precondition may not hold.
1539+
--
1540+
-- @since FIXME
1541+
1542+
-- See Note [fromAscList implementation] in Data.IntMap.Internal.
1543+
fromDescList :: [Key] -> IntSet
1544+
fromDescList xs = descLinkAll (Foldable.foldl' descInsert MSNada xs)
1545+
{-# INLINE fromDescList #-} -- Inline for list fusion
1546+
15331547
data Stack
15341548
= Nada
15351549
| Push {-# UNPACK #-} !Int !IntSet !Stack
@@ -1574,6 +1588,42 @@ ascLinkStack stk !rk r = case stk of
15741588
where
15751589
p = mask rk m
15761590

1591+
-- Insert an element. The element must be <= the last inserted element.
1592+
descInsert :: MonoState -> Int -> MonoState
1593+
descInsert s !ky = case s of
1594+
MSNada -> MSPush py bmy Nada
1595+
MSPush px bmx stk
1596+
| px == py -> MSPush py (bmx .|. bmy) stk
1597+
| otherwise -> let m = branchMask px py
1598+
in MSPush py bmy (descLinkTop px (Tip px bmx) m stk)
1599+
where
1600+
py = prefixOf ky
1601+
bmy = bitmapOf ky
1602+
{-# INLINE descInsert #-}
1603+
1604+
descLinkTop :: Int -> IntSet -> Int -> Stack -> Stack
1605+
descLinkTop !lk l !lm stk = case stk of
1606+
Nada -> Push lm l stk
1607+
Push m r stk'
1608+
| i2w m < i2w lm -> let p = mask lk m
1609+
in descLinkTop lk (Bin p l r) lm stk'
1610+
| otherwise -> Push lm l stk
1611+
1612+
descLinkAll :: MonoState -> IntSet
1613+
descLinkAll s = case s of
1614+
MSNada -> Nil
1615+
MSPush px bmx stk -> descLinkStack px (Tip px bmx) stk
1616+
{-# INLINABLE descLinkAll #-}
1617+
1618+
descLinkStack :: Int -> IntSet -> Stack -> IntSet
1619+
descLinkStack !lk l stk = case stk of
1620+
Nada -> l
1621+
Push m r stk'
1622+
| signBranch p -> Bin p r l
1623+
| otherwise -> descLinkStack lk (Bin p l r) stk'
1624+
where
1625+
p = mask lk m
1626+
15771627
{--------------------------------------------------------------------
15781628
IntSetBuilder
15791629
--------------------------------------------------------------------}

0 commit comments

Comments
 (0)