Set.powerSet experiments

Main.hs

{-# LANGUAGE ScopedTypeVariables #-}
import Control.DeepSeq
import Control.Exception
import Data.Bits
import Data.Functor
import Test.Tasty
import Test.Tasty.QuickCheck hiding ((.&.))
import qualified Test.Tasty.Bench as B
import qualified Data.Foldable as F
import qualified Data.List as L

import qualified Data.Array.Unboxed as UA
import qualified Data.Array.ST as STA

import Data.Set.Internal

main :: IO ()
main = do
  let sets = [("set" ++ show i, fromList [1..i]) | i <- [5,10,15,20]]
      funs :: [(String, PowFun)]
      funs =
        [ ("powerSet", powerSet)
        , ("powerSet2", powerSet2)
        , ("powerSet3", powerSet3)
        ]
  evaluate $ rnf sets `seq` rnf funs
  B.defaultMain
    [ testGroup "test" $
        funs <&> \(pwname, pw) ->
          testProperty pwname (prop_powerSet pw)
    , B.bgroup "bench" $
        sets <&> \(name, s) ->
          B.bgroup ("prop_powerSet_" ++ name) $
            funs <&> \(pwname, pw) ->
              B.bench pwname $ B.whnf pw s
    ]

type PowFun = Set Int -> Set (Set Int)

-- O(2^n)
powerSet2 :: forall a. Set a -> Set (Set a)
powerSet2 xs0 = fromDistinctAscList (L.map (sets UA.!) sortedMasks)
  where
    n = size xs0
    xs = UA.listArray (0, n - 1) (toList xs0) :: UA.Array Int a

    middleBit :: UA.Array Int Int
    middleBit = UA.listArray (1, bit n - 1) (L.map f [1 .. bit n - 1])
      where
        f 1 = 0
        f msk
          | even msk            = (middleBit UA.! (msk `quot` 2)) + 1
          | even (popCount msk) = middleBit UA.! (msk - 1)
          | otherwise           = let x = middleBit UA.! (msk - 1)
                                      -- return the next lowest set bit after x
                                      msk' = msk .&. (bit x - 1)
                                  in finiteBitSize msk' - 1 - countLeadingZeros msk'

    sets :: UA.Array Int (Set a)
    sets = UA.listArray (0, bit n - 1) (L.map f [0 .. bit n - 1])
      where
        f 0 = Tip
        f msk = bin x (sets UA.! mskl) (sets UA.! mskr)
          where
            m = middleBit UA.! msk
            x = xs UA.! m
            mskl = msk .&. (bit m - 1)
            mskr = msk `xor` bit m `xor` mskl

    sortedMasks :: [Int]
    sortedMasks = 0 : L.foldl' step [] [n-1, n-2 .. 0]
      where
        step msks i = bit i : L.map (bit i +) msks ++ msks

-- Same as powerSet2 but be strict and avoid thunks. O(2^n)
powerSet3 :: forall a. Set a -> Set (Set a)
powerSet3 xs0 = result
  where
    n = size xs0
    xs = UA.listArray (0, n - 1) (toList xs0) :: UA.Array Int a

    middleBit :: UA.UArray Int Int
    middleBit = STA.runSTUArray $ do
      a <- STA.newArray_ (1, bit n - 1)
      STA.writeArray a 1 0
      F.for_ [2 .. bit n - 1] $ \i -> do
        y <-
            if even i
            then (1+) <$> STA.readArray a (i `quot` 2)
            else if even (popCount i)
            then STA.readArray a (i - 1)
            else do
                x <- STA.readArray a (i - 1)
                -- find the next lowest set bit after x
                let i' = i .&. (bit x - 1)
                pure $! finiteBitSize i' - 1 - countLeadingZeros i'
        STA.writeArray a i y
      pure a

    sets :: UA.Array Int (Set a)
    sets = STA.runSTArray $ do
      a <- STA.newArray (0, bit n - 1) Tip
      F.for_ [1 .. bit n - 1] $ \i -> do
        let m = middleBit UA.! i
            x = xs UA.! m
            li = i .&. (bit m - 1)
            ri = i `xor` bit m `xor` li
        l <- STA.readArray a li
        r <- STA.readArray a ri
        STA.writeArray a i $! bin x l r
      pure a

    sortedMasks :: UA.UArray Int Int
    sortedMasks = STA.runSTUArray $ do
      a <- STA.newArray_ (0, bit n - 1)
      STA.writeArray a 0 0
      F.for_ [n-1, n-2 .. 0] $ \i -> do
        let d = bit (n - 1 - i) - 1
            start = bit n - 2 * d
        F.for_ [start .. start + d - 1] $ \j -> do
          STA.readArray a (j + d) >>=
            STA.writeArray a j . (bit i +)
        STA.writeArray a (start - 1) (bit i)
      pure a

    buildSet :: Int -> Int -> Set (Set a)
    buildSet l u | l > u = Tip
    buildSet l u = bin (sets UA.! i) (buildSet l (m-1)) (buildSet (m+1) u)
      where
        m = l + (u - l) `quot` 2
        i = sortedMasks UA.! m

    result = buildSet 0 (bit n - 1)

prop_powerSet :: PowFun -> Property
prop_powerSet pw = forAll (resize 10 arbitrary :: Gen (Set Int)) $ \xs ->
   let ps = pw xs
   in valid ps .&&. all valid ps .&&. size ps === 2^size xs

out.txt

All
  test
    powerSet:    OK
      +++ OK, passed 100 tests.
    powerSet2:   OK
      +++ OK, passed 100 tests.
    powerSet3:   OK
      +++ OK, passed 100 tests.
  bench
    prop_powerSet_set5
      powerSet:  OK (0.13s)
        1.06 μs ±  83 ns, 7.1 KB allocated,   4 B  copied, 6.0 MB peak memory
      powerSet2: OK (0.13s)
        4.03 μs ± 382 ns,  17 KB allocated,   1 B  copied, 6.0 MB peak memory
      powerSet3: OK (0.16s)
        1.17 μs ±  99 ns, 4.9 KB allocated,   1 B  copied, 6.0 MB peak memory
    prop_powerSet_set10
      powerSet:  OK (0.17s)
        41.5 μs ± 2.7 μs, 254 KB allocated, 5.6 KB copied, 6.0 MB peak memory
      powerSet2: OK (0.28s)
        137  μs ± 5.7 μs, 506 KB allocated, 7.9 KB copied, 6.0 MB peak memory
      powerSet3: OK (0.27s)
        32.1 μs ± 1.4 μs, 122 KB allocated, 854 B  copied, 8.0 MB peak memory
    prop_powerSet_set15
      powerSet:  OK (0.19s)
        5.99 ms ± 474 μs, 8.6 MB allocated, 7.9 MB copied,  16 MB peak memory
      powerSet2: OK (0.34s)
        10.9 ms ± 463 μs,  16 MB allocated, 8.7 MB copied,  17 MB peak memory
      powerSet3: OK (0.39s)
        1.53 ms ±  47 μs, 3.7 MB allocated, 903 KB copied,  17 MB peak memory
    prop_powerSet_set20
      powerSet:  OK (8.72s)
        567  ms ±  11 ms, 306 MB allocated, 707 MB copied, 634 MB peak memory
      powerSet2: OK (3.10s)
        1.027 s ±  15 ms, 507 MB allocated, 456 MB copied, 634 MB peak memory
      powerSet3: OK (1.75s)
        246  ms ± 5.8 ms, 120 MB allocated, 218 MB copied, 634 MB peak memory

Can you use the exact benchmarks that I submitted recently? haskell/containers#897

cf. also my numbers haskell/containers#904 (comment)

perhaps you can use next_pattern https://github.com/haskell/containers/blob/daf6133f73c5d21ea01342405fa10ef33f056b36/containers/src/Data/Set/Internal.hs#L1886