Add compareSize for IntSet and IntMap (#1135)

This function short-circuits, which can be more efficient than
calculating the size.
Other structures for which `length` takes O(n) provide similar
`compareLength`s: list, Text, lazy Text, lazy ByteString.

Author: meooow25

containers-tests/tests/intmap-properties.hs

@@ -138,6 +138,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testCase "maxViewWithKey" test_maxViewWithKey
              , testCase "minimum" test_minimum
              , testCase "maximum" test_maximum
+             , testCase "compareSize" test_compareSize
              , testProperty "valid"                prop_valid
              , testProperty "hasPrefix"            prop_hasPrefix
              , testProperty "empty valid"          prop_emptyValid
@@ -254,6 +255,7 @@ main = defaultMain $ testGroup "intmap-properties"
              , testProperty "fromDistinctAscList"  prop_fromDistinctAscList
              , testProperty "fromListWith"         prop_fromListWith
              , testProperty "fromListWithKey"      prop_fromListWithKey
+             , testProperty "compareSize"          prop_compareSize
              ]
 
 {--------------------------------------------------------------------
@@ -1196,6 +1198,12 @@ test_maximum = do
     maximum (elems testOrdMap) @?= maximum testOrdMap
   where getOW (OrdWith s _) = s
 
+-- Check cases where there is a risk of overflow
+test_compareSize :: Assertion
+test_compareSize = do
+  compareSize (fromList [(1,'a')]) minBound @?= GT
+  compareSize (fromList [(1,'a')]) maxBound @?= LT
+
 testOrdMap :: IntMap (OrdWith Int)
 testOrdMap = fromList [(1,OrdWith "max" 1),(-1,OrdWith "min" 1)]
 
@@ -2033,3 +2041,6 @@ prop_fromListWithKey f kxs =
   m === List.foldl' (\m' (k,x) -> insertWith (applyFun3 f k) k x m') empty kxs
   where
     m = fromListWithKey (applyFun3 f) kxs
+
+prop_compareSize :: IntMap A -> Int -> Property
+prop_compareSize t c = compareSize t c === compare (size t) c

containers-tests/tests/intset-properties.hs

@@ -24,6 +24,7 @@ main = defaultMain $ testGroup "intset-properties"
                    , testCase "lookupGE" test_lookupGE
                    , testCase "split" test_split
                    , testCase "isProperSubsetOf" test_isProperSubsetOf
+                   , testCase "compareSize" test_compareSize
                    , testProperty "prop_Valid" prop_Valid
                    , testProperty "prop_EmptyValid" prop_EmptyValid
                    , testProperty "prop_SingletonValid" prop_SingletonValid
@@ -91,6 +92,7 @@ main = defaultMain $ testGroup "intset-properties"
                    , testProperty "deleteMax" prop_deleteMax
                    , testProperty "fromAscList" prop_fromAscList
                    , testProperty "fromDistinctAscList" prop_fromDistinctAscList
+                   , testProperty "compareSize" prop_compareSize
                    ]
 
 ----------------------------------------------------------------
@@ -136,6 +138,12 @@ test_isProperSubsetOf = do
     -- See Github #1007
     isProperSubsetOf (fromList [-65,-1]) (fromList [-65,-1,0]) @?= True
 
+-- Check cases where there is a risk of overflow
+test_compareSize :: Assertion
+test_compareSize = do
+  compareSize (fromList [1]) minBound @?= GT
+  compareSize (fromList [1]) maxBound @?= LT
+
 {--------------------------------------------------------------------
   Arbitrary, reasonably balanced trees
 --------------------------------------------------------------------}
@@ -535,3 +543,6 @@ prop_fromDistinctAscList xs =
   where
     nubSortedXs = List.map NE.head $ NE.group $ sort xs
     t = fromDistinctAscList nubSortedXs
+
+prop_compareSize :: IntSet -> Int -> Property
+prop_compareSize t c = compareSize t c === compare (size t) c

containers/changelog.md

@@ -2,6 +2,11 @@
 
 ## Next release
 
+### Additions
+
+* Add `compareSize` for `IntSet` and `IntMap`. (Soumik Sarkar)
+  ([#1135](https://github.com/haskell/containers/pull/1135))
+
 ### Performance improvements
 
 * Improved performance for `Data.IntMap.restrictKeys` and

containers/src/Data/IntMap/Internal.hs

@@ -95,6 +95,7 @@ module Data.IntMap.Internal (
     -- * Query
     , null
     , size
+    , compareSize
     , member
     , notMember
     , lookup
@@ -579,13 +580,33 @@ null _   = False
 -- > size empty                                   == 0
 -- > size (singleton 1 'a')                       == 1
 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3
+--
+-- See also: 'compareSize'
 size :: IntMap a -> Int
 size = go 0
   where
     go !acc (Bin _ l r) = go (go acc l) r
     go acc (Tip _ _) = 1 + acc
     go acc Nil = acc
 
+-- | \(O(\min(n,c))\). Compare the number of entries in the map to an @Int@.
+--
+-- @compareSize m c@ returns the same result as @compare ('size' m) c@ but is
+-- more efficient when @c@ is smaller than the size of the map.
+--
+-- @since FIXME
+compareSize :: IntMap a -> Int -> Ordering
+compareSize !_ c0 | c0 < 0 = GT
+compareSize Nil c0 = compare 0 c0
+compareSize t c0 = compare 0 (go t c0)
+  where
+    go (Bin _ l r) c
+      | c' <= 0 = -1
+      | otherwise = go r c'
+      where
+        c' = go l c
+    go _ c = c - 1 -- Must be Tip (invariant)
+
 -- | \(O(\min(n,W))\). Is the key a member of the map?
 --
 -- > member 5 (fromList [(5,'a'), (3,'b')]) == True

containers/src/Data/IntMap/Lazy.hs

@@ -147,6 +147,7 @@ module Data.IntMap.Lazy (
     -- ** Size
     , IM.null
     , size
+    , compareSize
 
     -- * Combine
 

containers/src/Data/IntMap/Strict.hs

@@ -165,6 +165,7 @@ module Data.IntMap.Strict (
     -- ** Size
     , null
     , size
+    , compareSize
 
     -- * Combine
 

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

@@ -108,6 +108,7 @@ module Data.IntMap.Strict.Internal (
     -- ** Size
     , null
     , size
+    , compareSize
 
     -- * Combine
 
@@ -305,6 +306,7 @@ import Data.IntMap.Internal
   , spanAntitone
   , restrictKeys
   , size
+  , compareSize
   , split
   , splitLookup
   , splitRoot

containers/src/Data/IntSet.hs

@@ -122,6 +122,7 @@ module Data.IntSet (
             , lookupGE
             , IS.null
             , size
+            , compareSize
             , isSubsetOf
             , isProperSubsetOf
             , disjoint

containers/src/Data/IntSet/Internal.hs

@@ -98,6 +98,7 @@ module Data.IntSet.Internal (
     -- * Query
     , null
     , size
+    , compareSize
     , member
     , notMember
     , lookupLT
@@ -358,13 +359,33 @@ null _   = False
 {-# INLINE null #-}
 
 -- | \(O(n)\). Cardinality of the set.
+--
+-- See also: 'compareSize'
 size :: IntSet -> Int
 size = go 0
   where
     go !acc (Bin _ l r) = go (go acc l) r
     go acc (Tip _ bm) = acc + popCount bm
     go acc Nil = acc
 
+-- | \(O(\min(n,c))\). Compare the number of elements in the set to an @Int@.
+--
+-- @compareSize m c@ returns the same result as @compare ('size' m) c@ but is
+-- more efficient when @c@ is smaller than the size of the set.
+--
+-- @since FIXME
+compareSize :: IntSet -> Int -> Ordering
+compareSize !_ c0 | c0 < 0 = GT
+compareSize t c0 = compare 0 (go t c0)
+  where
+    go (Bin _ l r) c
+        | c' <= 0 = -1
+        | otherwise = go r c'
+        where
+          c' = go l c
+    go (Tip _ bm) c = c - popCount bm
+    go Nil c = c
+
 -- | \(O(\min(n,W))\). Is the value a member of the set?
 
 -- See Note: Local 'go' functions and capturing.