diff --git a/containers-tests/benchmarks/IntSet.hs b/containers-tests/benchmarks/IntSet.hs index 6927d3033..22a2d8f34 100644 --- a/containers-tests/benchmarks/IntSet.hs +++ b/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 diff --git a/containers/src/Data/IntSet/Internal.hs b/containers/src/Data/IntSet/Internal.hs index 3bc157ba1..d83e1c9e6 100644 --- a/containers/src/Data/IntSet/Internal.hs +++ b/containers/src/Data/IntSet/Internal.hs @@ -1436,12 +1436,11 @@ 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 @@ -1449,13 +1448,12 @@ foldl'Bits prefix f z bitmap = go bitmap z 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)