Bitmap64.hs

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

module Bench.Bitmap64 (benchmarks) where

import Bench.Values (Values (..), countWhere, dense, forEach, prefix, scattered64, withValues)
import Control.Monad (void)
import Data.Bitmap.Roaring.FFI.Bitmap qualified as Bitmap
import Data.Bitmap.Roaring.FFI.Bitmap64 (Bitmap64)
import Data.Bitmap.Roaring.FFI.Bitmap64 qualified as Bitmap64
import Data.Bitmap.Roaring.FFI.Bitmap64.BulkContext (BulkContext (..))
import Data.Bitmap.Roaring.FFI.Bitmap64.BulkContext qualified as BulkContext
import Data.Bitmap.Roaring.FFI.Bitmap64.Iterator qualified as Iterator
import Data.Word (Word64)
import Foreign.C.Types (CChar, CSize)
import Foreign.Marshal.Alloc (allocaBytesAligned, free, mallocBytes)
import Foreign.Marshal.Array (allocaArray, mallocArray)
import Foreign.Marshal.Utils (fillBytes)
import Foreign.Ptr (Ptr, freeHaskellFunPtr, nullPtr)
import Test.Tasty (withResource)
import Test.Tasty.Bench (Benchmark, bench, bgroup, whnfIO)

size :: Int
size = 1_000_000

benchmarks :: Benchmark
benchmarks =
  withValues (scattered64 1 size) $ \getScattered ->
    withValues (scattered64 2 size) $ \getOther ->
      withValues (dense size) $ \getDense ->
        withBitmap getScattered $ \getScatteredBitmap ->
          withBitmap getOther $ \getOtherBitmap ->
            withBitmap getDense $ \getDenseBitmap ->
              bgroup
                "Bitmap64"
                [ bgroup
                    "construction"
                    [ bgroup "scattered" $ construction getScattered
                    , bgroup "dense" $ construction getDense
                    ]
                , bgroup
                    "queries"
                    [ bgroup "hits" $ queries getScatteredBitmap getScattered
                    , bgroup "misses" $ queries getScatteredBitmap getOther
                    ]
                , bgroup
                    "algebra"
                    [ bgroup "scattered" $ algebra getScatteredBitmap getOtherBitmap
                    , bgroup "mixed" $ algebra getScatteredBitmap getDenseBitmap
                    ]
                , bgroup
                    "traversal"
                    [ bgroup "scattered" $ traversal getScatteredBitmap
                    , bgroup "dense" $ traversal getDenseBitmap
                    ]
                , bgroup
                    "serialization"
                    [ bgroup "scattered" $ serialization getScatteredBitmap
                    , bgroup "dense" $ serialization getDenseBitmap
                    ]
                ]

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

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

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

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

traversal :: IO Bitmap64 -> [Benchmark]
traversal getBitmap =
  [ withResource (mallocArray size) free $ \getOutput ->
      bench "toUint64Array" . whnfIO $ do
        bitmap <- getBitmap
        getOutput >>= Bitmap64.toUint64Array 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.value iterator
                _ <- Iterator.advance iterator
                loop $ total + value
              else
                pure total
        loop 0
  , withResource (Bitmap.wrapCallback64 $ \_ _ -> pure True) freeHaskellFunPtr $ \getCallback ->
      bench "iterate" . whnfIO $ do
        bitmap <- getBitmap
        callback <- getCallback
        Bitmap64.iterate bitmap callback nullPtr
  ]
  where
    iterating bitmap continue = do
      iterator <- Iterator.create bitmap
      result <- continue iterator
      Iterator.free iterator
      pure result

data Buffer = Buffer
  { pointer :: Ptr CChar
  , size :: CSize
  }

withPortable :: IO Bitmap64 -> (IO Buffer -> Benchmark) -> Benchmark
withPortable getBitmap = withResource acquire release
  where
    acquire = do
      bitmap <- getBitmap
      measured <- Bitmap64.portableSizeInBytes bitmap
      pointer <- mallocBytes $ fromIntegral measured
      _ <- Bitmap64.portableSerialize bitmap pointer
      pure Buffer{pointer, size = measured}
    release buffer = free buffer.pointer

serialization :: IO Bitmap64 -> [Benchmark]
serialization getBitmap =
  [ bench "portableSizeInBytes" . whnfIO $ getBitmap >>= Bitmap64.portableSizeInBytes
  , bench "internalValidate" . whnfIO $ getBitmap >>= (`Bitmap64.internalValidate` nullPtr)
  , withPortable getBitmap $ \getBuffer ->
      bgroup
        "portable"
        [ bench "serialize" . whnfIO $ do
            bitmap <- getBitmap
            buffer <- getBuffer
            Bitmap64.portableSerialize bitmap buffer.pointer
        , bench "deserializeBounded" . whnfIO $ do
            buffer <- getBuffer
            Bitmap64.portableDeserializeBounded buffer.pointer buffer.size >>= Bitmap64.free
        , bench "deserializeFrozen" . whnfIO $ do
            buffer <- getBuffer
            Bitmap64.portableDeserializeFrozen buffer.pointer buffer.size >>= Bitmap64.free
        ]
  ]