Update compareSize implementations (#1139)

The previous implementations went straight down the left spine,
which does not match the documented complexity. We must decrement and
potentially return at each step to meet O(min(n,c)).

Also add some benchmarks for size and compareSize.

Author: meooow25

containers-tests/benchmarks/IntMap.hs

@@ -79,6 +79,9 @@ main = do
             whnf (uncurry M.withoutKeys) (m_random, s_random2)
         , bench "restrictKeys:random" $ -- large keys, no overlap
             whnf (uncurry M.restrictKeys) (m_random, s_random2)
+        , bench "size" $ whnf M.size m
+        , bench "compareSize:2" $ whnf (flip M.compareSize 2) m
+        , bench "compareSize:n" $ whnf (flip M.compareSize bound) m
         , bgroup "folds" $ foldBenchmarks M.foldr M.foldl M.foldr' M.foldl' foldMap m
         , bgroup "folds with key" $
             foldWithKeyBenchmarks M.foldrWithKey M.foldlWithKey M.foldrWithKey' M.foldlWithKey' M.foldMapWithKey m

containers-tests/benchmarks/IntSet.hs

@@ -72,6 +72,12 @@ main = do
         , bench "eq" $ whnf (\s' -> s' == s') s -- worst case, compares everything
         , bench "compare:dense" $ whnf (\s' -> compare s' s') s -- worst case, compares everything
         , bench "compare:sparse" $ whnf (\s' -> compare s' s') s_sparse -- worst case, compares everything
+        , bench "size" $ whnf IS.size s
+        , bench "size:sparse" $ whnf IS.size s_sparse
+        , bench "compareSize:2" $ whnf (flip IS.compareSize 2) s
+        , bench "compareSize:sparse:2" $ whnf (flip IS.compareSize 2) s_sparse
+        , bench "compareSize:n" $ whnf (flip IS.compareSize bound) s
+        , bench "compareSize:sparse:n" $ whnf (flip IS.compareSize bound) s_sparse
         , bgroup "folds:dense" $ foldBenchmarks IS.foldr IS.foldl IS.foldr' IS.foldl' IS.foldMap s
         , bgroup "folds:sparse" $ foldBenchmarks IS.foldr IS.foldl IS.foldr' IS.foldl' IS.foldMap s_sparse
         ]

containers/src/Data/IntMap/Internal.hs

@@ -598,16 +598,17 @@ size = go 0
 --
 -- @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)
+compareSize _ c0 | c0 <= 0 = GT
+compareSize t c0 = compare 0 (go t (c0 - 1))
   where
+    go (Bin _ _ _) 0 = -1
     go (Bin _ l r) c
-      | c' <= 0 = -1
+      | c' < 0 = c'
       | otherwise = go r c'
       where
-        c' = go l c
-    go _ c = c - 1 -- Must be Tip (invariant)
+        c' = go l (c - 1)
+    go _ c = c
 
 -- | \(O(\min(n,W))\). Is the key a member of the map?
 --

containers/src/Data/IntSet/Internal.hs

@@ -377,16 +377,18 @@ size = go 0
 --
 -- @since FIXME
 compareSize :: IntSet -> Int -> Ordering
-compareSize !_ c0 | c0 < 0 = GT
-compareSize t c0 = compare 0 (go t c0)
+compareSize Nil c0 = compare 0 c0
+compareSize _ c0 | c0 <= 0 = GT
+compareSize t c0 = compare 0 (go t (c0 - 1))
   where
+    go (Bin _ _ _) 0 = -1
     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
+      | c' < 0 = c'
+      | otherwise = go r c'
+      where
+        c' = go l (c - 1)
+    go (Tip _ bm) c = c + 1 - popCount bm
+    go Nil !_ = error "compareSize.go: Nil"
 
 -- | \(O(\min(n,W))\). Is the value a member of the set?