Bitmap.hs

{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoFieldSelectors #-}

module Bench.Bitmap (benchmarks) where

import Bench.Values (Values (..), countWhere, dense, forEach, prefix, scattered, withValues)
import Control.Monad (forM, void)
import Data.Bitmap.Roaring.FFI.Bitmap (Bitmap (..))
import Data.Bitmap.Roaring.FFI.Bitmap qualified as Bitmap
import Data.Bitmap.Roaring.FFI.Bitmap.BulkContext (BulkContext (..))
import Data.Bitmap.Roaring.FFI.Bitmap.BulkContext qualified as BulkContext
import Data.Bitmap.Roaring.FFI.Bitmap.Iterator (Iterator (..))
import Data.Bitmap.Roaring.FFI.Bitmap.Iterator qualified as Iterator
import Data.List (sort)
import Data.Word (Word32, Word64)
import Foreign.C.ConstPtr (ConstPtr (..))
import Foreign.C.Types (CChar, CSize)
import Foreign.Marshal.Alloc (allocaBytesAligned, free, mallocBytes)
import Foreign.Marshal.Array (advancePtr, allocaArray, mallocArray, peekArray, pokeArray, withArrayLen)
import Foreign.Marshal.Utils (fillBytes)
import Foreign.Ptr (Ptr, alignPtr, freeHaskellFunPtr, nullPtr)
import Test.Tasty (withResource)
import Test.Tasty.Bench (Benchmark, bench, bgroup, whnfIO)

size :: Int
size = 1_000_000

benchmarks :: Benchmark
benchmarks =
  withValues (scattered 1 size) $ \getScattered ->
    withValues (scattered 2 size) $ \getOther ->
      withValues (dense size) $ \getDense ->
        withBitmap getScattered $ \getScatteredBitmap ->
          withBitmap getOther $ \getOtherBitmap ->
            withBitmap getDense $ \getDenseBitmap ->
              bgroup
                "Bitmap"
                [ bgroup
                    "construction"
                    [ bgroup "scattered" $ construction getScattered
                    , bgroup "dense" $ construction getDense <> [bench "fromRange" . whnfIO $ Bitmap.fromRange 0 (fromIntegral size) 1 >>= Bitmap.free]
                    ]
                , bgroup
                    "queries"
                    [ bgroup "hits" $ queries getScatteredBitmap getScattered
                    , bgroup "misses" $ queries getScatteredBitmap getOther
                    , rankMany getScatteredBitmap
                    ]
                , bgroup
                    "algebra"
                    [ bgroup "scattered" $ algebra getScatteredBitmap getOtherBitmap
                    , bgroup "mixed" $ algebra getScatteredBitmap getDenseBitmap
                    , many
                    ]
                , bgroup
                    "traversal"
                    [ bgroup "scattered" $ traversal getScatteredBitmap
                    , bgroup "dense" $ traversal getDenseBitmap
                    ]
                , bgroup
                    "serialization"
                    [ bgroup "scattered" $ serialization getScatteredBitmap
                    , bgroup "dense" $ serialization getDenseBitmap
                    ]
                ]

withBitmap :: IO (Values Word32) -> (IO Bitmap -> Benchmark) -> Benchmark
withBitmap getValues = withResource acquire Bitmap.free
  where
    acquire = do
      values <- getValues
      Bitmap.ofPtr (fromIntegral values.count) values.pointer

construction :: IO (Values Word32) -> [Benchmark]
construction getValues =
  [ bench "ofPtr" . whnfIO $ do
      values <- getValues
      Bitmap.ofPtr (fromIntegral values.count) values.pointer >>= Bitmap.free
  , filling "addMany" $ \bitmap values -> Bitmap.addMany bitmap (fromIntegral values.count) values.pointer
  , filling "add" $ \bitmap values -> forEach values $ Bitmap.add bitmap
  , filling "addChecked" $ \bitmap values -> forEach values $ void . Bitmap.addChecked bitmap
  , filling "addBulk" $ \bitmap values ->
      allocaBytesAligned BulkContext.size BulkContext.alignment $ \memory -> do
        fillBytes memory 0 BulkContext.size
        forEach values . Bitmap.addBulk bitmap $ BulkContext memory
  ]
  where
    filling :: String -> (Bitmap -> Values Word32 -> IO ()) -> Benchmark
    filling name fill =
      bench name . whnfIO $ do
        values <- getValues
        bitmap <- Bitmap.create
        fill bitmap values
        Bitmap.free bitmap

queries :: IO Bitmap -> IO (Values Word32) -> [Benchmark]
queries getBitmap getProbes =
  [ probing "contains" $ \bitmap probes -> countWhere probes $ Bitmap.contains bitmap
  , probing "containsBulk" $ \bitmap probes ->
      allocaBytesAligned BulkContext.size BulkContext.alignment $ \memory -> do
        fillBytes memory 0 BulkContext.size
        countWhere probes . Bitmap.containsBulk bitmap $ BulkContext memory
  , probing "rank of 1000" $ \bitmap probes -> countWhere (prefix 1000 probes) $ fmap (> 0) . Bitmap.rank bitmap
  , probing "getIndex of 1000" $ \bitmap probes -> countWhere (prefix 1000 probes) $ fmap (>= 0) . Bitmap.getIndex bitmap
  , bench "getCardinality" . whnfIO $ getBitmap >>= Bitmap.getCardinality
  ]
  where
    probing name query =
      bench name . whnfIO $ do
        bitmap <- getBitmap
        probes <- getProbes
        query bitmap probes

rankMany :: IO Bitmap -> Benchmark
rankMany getBitmap =
  withValues (sort $ scattered 1 size) $ \getSorted ->
    withResource (mallocArray size) free $ \getOutput ->
      bench "rankMany" . whnfIO $ do
        bitmap <- getBitmap
        sorted <- getSorted
        output <- getOutput
        Bitmap.rankMany bitmap sorted.pointer (advancePtr sorted.pointer sorted.count) output

algebra :: IO Bitmap -> IO Bitmap -> [Benchmark]
algebra getLeft getRight =
  [ fresh "and" Bitmap.and
  , fresh "or" Bitmap.or
  , fresh "xor" Bitmap.xor
  , fresh "andnot" Bitmap.andnot
  , measured "andCardinality" Bitmap.andCardinality
  , measured "orCardinality" Bitmap.orCardinality
  , measured "intersect" Bitmap.intersect
  , measured "equals" Bitmap.equals
  , bench "copy" . whnfIO $ getLeft >>= Bitmap.copy >>= Bitmap.free
  ]
  where
    fresh name operation = measured name $ \left right -> operation left right >>= Bitmap.free
    measured name operation =
      bench name . whnfIO $ do
        left <- getLeft
        right <- getRight
        operation left right

many :: Benchmark
many =
  withResource acquire release $ \getHandles ->
    bgroup
      "many"
      [ combined getHandles "orMany" $ Bitmap.orMany . fromIntegral
      , combined getHandles "orManyHeap" $ Bitmap.orManyHeap . fromIntegral
      , combined getHandles "xorMany" $ Bitmap.xorMany . fromIntegral
      ]
  where
    count = 16 :: Int
    combined getHandles name operation = bench name . whnfIO $ getHandles >>= operation count >>= Bitmap.free
    acquire = do
      bitmaps <- forM [1 .. fromIntegral count] $ \seed ->
        withArrayLen (scattered seed 100_000) $ \values ptr -> Bitmap.ofPtr (fromIntegral values) ptr
      handles <- mallocArray count
      pokeArray handles bitmaps
      pure handles
    release handles = do
      peekArray count handles >>= mapM_ Bitmap.free
      free handles

traversal :: IO Bitmap -> [Benchmark]
traversal getBitmap =
  [ withResource (mallocArray size) free $ \getOutput ->
      bench "toUint32Array" . whnfIO $ do
        bitmap <- getBitmap
        getOutput >>= Bitmap.toUint32Array bitmap
  , bench "Iterator.read" . whnfIO $ do
      bitmap <- getBitmap
      iterating bitmap $ \iterator ->
        allocaArray 256 $ \buffer -> do
          let loop !total = do
                count <- Iterator.read iterator buffer 256
                if count < 256 then pure $ total + count else loop $ total + count
          loop 0
  , bench "Iterator.advance" . whnfIO $ do
      bitmap <- getBitmap
      iterating bitmap $ \iterator -> do
        let loop !total = do
              present <- Iterator.hasValue iterator
              if present then do
                value <- Iterator.currentValue iterator
                _ <- Iterator.advance iterator
                loop $ total + fromIntegral value
              else
                pure (total :: Word64)
        loop 0
  , withResource (Bitmap.wrapCallback $ \_ _ -> pure True) freeHaskellFunPtr $ \getCallback ->
      bench "iterate" . whnfIO $ do
        bitmap <- getBitmap
        callback <- getCallback
        Bitmap.iterate bitmap callback nullPtr
  ]
  where
    iterating bitmap continue =
      allocaBytesAligned Iterator.size Iterator.alignment $ \memory -> do
        let iterator = Iterator memory
        Iterator.init bitmap iterator
        continue iterator

data Buffer = Buffer
  { allocation :: Ptr CChar
  , aligned :: Ptr CChar
  , size :: CSize
  }

withBuffer :: IO Bitmap -> (Bitmap -> IO CSize) -> (Bitmap -> Ptr CChar -> IO a) -> (IO Buffer -> Benchmark) -> Benchmark
withBuffer getBitmap measure write = withResource acquire release
  where
    acquire = do
      bitmap <- getBitmap
      measured <- measure bitmap
      allocation <- mallocBytes $ fromIntegral measured + 32
      let aligned = alignPtr allocation 32
      _ <- write bitmap aligned
      pure Buffer{allocation, aligned, size = measured}
    release buffer = free buffer.allocation

serialization :: IO Bitmap -> [Benchmark]
serialization getBitmap =
  [ bench "portableSizeInBytes" . whnfIO $ getBitmap >>= Bitmap.portableSizeInBytes
  , bench "internalValidate" . whnfIO $ getBitmap >>= (`Bitmap.internalValidate` nullPtr)
  , withBuffer getBitmap Bitmap.portableSizeInBytes Bitmap.portableSerialize $ \getBuffer ->
      bgroup
        "portable"
        [ bench "serialize" . whnfIO $ do
            bitmap <- getBitmap
            buffer <- getBuffer
            Bitmap.portableSerialize bitmap buffer.aligned
        , bench "deserializeBounded" . whnfIO $ do
            buffer <- getBuffer
            Bitmap.portableDeserializeBounded buffer.aligned buffer.size >>= Bitmap.free
        , bench "deserializeFrozen" . whnfIO $ do
            buffer <- getBuffer
            Bitmap.portableDeserializeFrozen buffer.aligned >>= Bitmap.free
        ]
  , withBuffer getBitmap Bitmap.frozenSizeInBytes Bitmap.frozenSerialize $ \getBuffer ->
      bgroup
        "frozen"
        [ bench "serialize" . whnfIO $ do
            bitmap <- getBitmap
            buffer <- getBuffer
            Bitmap.frozenSerialize bitmap buffer.aligned
        , bench "view" . whnfIO $ do
            buffer <- getBuffer
            Bitmap.frozenView buffer.aligned buffer.size >>= Bitmap.free . Bitmap . unConstPtr
        ]
  ]