Bitset.hs

{-# LANGUAGE OverloadedRecordDot #-}

module Bench.Bitset (benchmarks) where

import Bench.Values (Values (..), countWhere, dense, forEach, scattered, withValues)
import Data.Bitmap.Roaring.FFI.Bitset (Bitset)
import Data.Bitmap.Roaring.FFI.Bitset qualified as Bitset
import Data.Bits ((.&.))
import Foreign.C.Types (CSize)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Marshal.Array (allocaArray)
import Foreign.Ptr (freeHaskellFunPtr, nullPtr)
import Foreign.Storable (peek, poke)
import Test.Tasty (withResource)
import Test.Tasty.Bench (Benchmark, bench, bgroup, whnfIO)

size :: Int
size = 1_000_000

benchmarks :: Benchmark
benchmarks =
  withValues (map (fromIntegral . (.&. 0xFFFFFF)) $ scattered 1 size) $ \getScattered ->
    withValues (map (fromIntegral . (.&. 0xFFFFFF)) $ scattered 2 size) $ \getOther ->
      withValues (dense size) $ \getDense ->
        withBitset getScattered $ \getScatteredBitset ->
          withBitset getOther $ \getOtherBitset ->
            bgroup
              "Bitset"
              [ bgroup
                  "construction"
                  [ bgroup "scattered" $ construction getScattered
                  , bgroup "dense" $ construction getDense
                  ]
              , bgroup
                  "queries"
                  [ probing getScatteredBitset getScattered "get hits"
                  , probing getScatteredBitset getOther "get misses"
                  , bench "count" . whnfIO $ getScatteredBitset >>= Bitset.count
                  ]
              , bgroup
                  "algebra"
                  [ measured getScatteredBitset getOtherBitset "unionCount" Bitset.unionCount
                  , measured getScatteredBitset getOtherBitset "intersectionCount" Bitset.intersectionCount
                  , bench "copy and inplaceUnion" . whnfIO $ do
                      left <- getScatteredBitset
                      right <- getOtherBitset
                      copy <- Bitset.copy left
                      _ <- Bitset.inplaceUnion copy right
                      Bitset.free copy
                  ]
              , bgroup
                  "traversal"
                  [ bench "nextSetBits" . whnfIO $ getScatteredBitset >>= nextSetBits
                  , withResource (Bitset.wrapCallback $ \_ _ -> pure True) freeHaskellFunPtr $ \getCallback ->
                      bench "forEach" . whnfIO $ do
                        bitset <- getScatteredBitset
                        callback <- getCallback
                        Bitset.forEach bitset callback nullPtr
                  ]
              ]

withBitset :: IO (Values CSize) -> (IO Bitset -> Benchmark) -> Benchmark
withBitset getValues = withResource acquire Bitset.free
  where
    acquire = do
      values <- getValues
      bitset <- Bitset.create
      forEach values $ Bitset.set bitset
      pure bitset

construction :: IO (Values CSize) -> [Benchmark]
construction getValues =
  [ bench "set" . whnfIO $ do
      values <- getValues
      bitset <- Bitset.create
      forEach values $ Bitset.set bitset
      Bitset.free bitset
  , bench "set with capacity" . whnfIO $ do
      values <- getValues
      bitset <- Bitset.createWithCapacity 0x1000000
      forEach values $ Bitset.set bitset
      Bitset.free bitset
  ]

probing :: IO Bitset -> IO (Values CSize) -> String -> Benchmark
probing getBitset getProbes name =
  bench name . whnfIO $ do
    bitset <- getBitset
    probes <- getProbes
    countWhere probes $ Bitset.get bitset

measured :: IO Bitset -> IO Bitset -> String -> (Bitset -> Bitset -> IO CSize) -> Benchmark
measured getLeft getRight name operation =
  bench name . whnfIO $ do
    left <- getLeft
    right <- getRight
    operation left right

nextSetBits :: Bitset -> IO CSize
nextSetBits bitset =
  allocaArray 256 $ \buffer ->
    alloca $ \start -> do
      poke start 0
      let loop !total = do
            count <- Bitset.nextSetBits bitset buffer 256 start
            if count == 0 then
              pure total
            else do
              peek start >>= poke start . (+ 1)
              loop $ total + count
      loop 0