Improve foldl and foldr for IntSet

containers-tests/benchmarks/IntSet.hs

@@ -15,6 +15,8 @@ main = do
     evaluate $ rnf [s, s_even, s_odd]
     defaultMain
         [ bench "member" $ whnf (member elems) s
+        , bench "foldr" $ whnf (sumUntilL (2^14 - 5)) fish
+        , bench "foldl" $ whnf (sumUntilR 5) fish
         , bench "insert" $ whnf (ins elems) S.empty
         , bench "map" $ whnf (S.map (+ 1)) s
         , bench "filter" $ whnf (S.filter ((== 0) . (`mod` 2))) s
@@ -38,10 +40,25 @@ main = do
         , bench "null.intersection:true" $ whnf (S.null. S.intersection s_odd) s_even
         ]
   where
+    fish = S.fromDistinctAscList [7..2^14]
     elems = [1..2^12]
     elems_even = [2,4..2^12]
     elems_odd = [1,3..2^12]
 
+sumUntilL :: Int -> S.IntSet -> Int
+sumUntilL n = \xs -> S.foldr go id xs 0
+  where
+    go x r !acc
+      | x >= n = acc
+      | otherwise = r (acc + x)
+
+sumUntilR :: Int -> S.IntSet -> Int
+sumUntilR n = \xs -> S.foldl go id xs 0
+  where
+    go r x !acc
+      | x <= n = acc
+      | otherwise = r (acc + x)
+
 member :: [Int] -> S.IntSet -> Int
 member xs s = foldl' (\n x -> if S.member x s then n + 1 else n) 0 xs
 

containers/src/Data/IntSet/Internal.hs

@@ -1436,26 +1436,24 @@ lowestBitSet x = indexOfTheOnlyBit (lowestBitMask x)
 
 highestBitSet x = indexOfTheOnlyBit (highestBitMask x)
 
-foldlBits prefix f z bitmap = go bitmap z
-  where go 0 acc = acc
-        go bm acc = go (bm `xor` bitmask) ((f acc) $! (prefix+bi))
-          where
-            !bitmask = lowestBitMask bm
-            !bi = indexOfTheOnlyBit bitmask
+foldlBits !prefix f z bitmap = go (revNat bitmap)
+  where go 0 = z
+        go bm = f (go (bm `xor` bitmask)) (prefix+(WORD_SIZE_IN_BITS-1)-bi)
+          where !bitmask = lowestBitMask bm
+                !bi = indexOfTheOnlyBit bitmask
 
 foldl'Bits prefix f z bitmap = go bitmap z
   where go 0 acc = acc
         go bm !acc = go (bm `xor` bitmask) ((f acc) $! (prefix+bi))
           where !bitmask = lowestBitMask bm
                 !bi = indexOfTheOnlyBit bitmask
 
-foldrBits prefix f z bitmap = go (revNat bitmap) z
-  where go 0 acc = acc
-        go bm acc = go (bm `xor` bitmask) ((f $! (prefix+(WORD_SIZE_IN_BITS-1)-bi)) acc)
+foldrBits !prefix f z bitmap = go bitmap
+  where go 0 = z
+        go bm = f (prefix + bi) (go (bm `xor` bitmask))
           where !bitmask = lowestBitMask bm
                 !bi = indexOfTheOnlyBit bitmask
 
-
 foldr'Bits prefix f z bitmap = go (revNat bitmap) z
   where go 0 acc = acc
         go bm !acc = go (bm `xor` bitmask) ((f $! (prefix+(WORD_SIZE_IN_BITS-1)-bi)) acc)
@@ -1489,10 +1487,11 @@ highestBitSet n0 =
     in b6
 
 foldlBits prefix f z bm = let lb = lowestBitSet bm
-                          in  go (prefix+lb) z (bm `shiftRL` lb)
-  where go !_ acc 0 = acc
-        go bi acc n | n `testBit` 0 = go (bi + 1) (f acc bi) (n `shiftRL` 1)
-                    | otherwise     = go (bi + 1)    acc     (n `shiftRL` 1)
+                          in  go (prefix+lb) (bm `shiftRL` lb)
+  where
+        go !_ 0 = z
+        go bi n | n `testBit` 0 = f (go (bi + 1) (n `shiftRL` 1)) bi
+                | otherwise     = go (bi + 1) (n `shiftRL` 1)
 
 foldl'Bits prefix f z bm = let lb = lowestBitSet bm
                            in  go (prefix+lb) z (bm `shiftRL` lb)