Unfair.hs

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

module Bench.Unfair (benchmarks) where

import Bench.Baseline qualified as Baseline
import Bench.Values (dense, scattered)
import Control.Exception (bracket)
import Data.Binary qualified as Binary
import Data.Bitmap.Roaring.FFI.Bitmap (Bitmap)
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.Iterator qualified as Iterator64
import Data.ByteString.Lazy qualified as Lazy
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.Word (Word32, Word64)
import Foreign.C.Types (CChar, CSize)
import Foreign.Marshal.Alloc (alloca, allocaBytes, free, mallocBytes)
import Foreign.Marshal.Array (allocaArray, peekArray, withArrayLen)
import Foreign.Ptr (Ptr)
import Foreign.Storable (peek)
import Test.Tasty (withResource)
import Test.Tasty.Bench (Benchmark, Benchmarkable, bgroup, env, nf, nfIO, whnf, whnfIO)

data Dataset = Dataset
  { name :: String
  , left :: [Word32]
  , right :: [Word32]
  , low :: Word32
  , high :: Word32
  }

data Contender = Contender
  { name :: String
  , difference :: Benchmarkable
  , intersectionSize :: Benchmarkable
  , jaccard :: Benchmarkable
  , rangeCount :: Benchmarkable
  , rank :: Benchmarkable
  , page :: Benchmarkable
  , shift :: Benchmarkable
  , serialize :: Benchmarkable
  , deserialize :: Benchmarkable
  , view :: Benchmarkable
  }

size :: Int
size = 1_000_000

pageOffset :: Int
pageOffset = size `div` 2

pageSize :: Int
pageSize = 100

offset :: Int
offset = 1_000_000

benchmarks :: Benchmark
benchmarks =
  bgroup "unfair" $
    map
      dataset
      [ Dataset
          { name = "sparse"
          , left = scattered 1 size
          , right = scattered 2 size
          , low = 0x40000000
          , high = 0xC0000000
          }
      , Dataset
          { name = "dense"
          , left = dense size
          , right = map (+ fromIntegral (size `div` 2)) $ dense size
          , low = 250_000
          , high = 750_000
          }
      ]

dataset :: Dataset -> Benchmark
dataset spec =
  withIntSet spec $ \intSet ->
    withBitmap spec $ \bitmap ->
      withBitmap64 spec $ \bitmap64 ->
        bgroup spec.name $ contests spec.name intSet [bitmap, bitmap64]

contests :: String -> Contender -> [Contender] -> [Benchmark]
contests datasetName baseline others =
  [ versus "difference" (.difference)
  , versus "intersection-size" (.intersectionSize)
  , versus "jaccard" (.jaccard)
  , versus "range-count" (.rangeCount)
  , versus "rank" (.rank)
  , versus "page" (.page)
  , versus "shift" (.shift)
  , versus "serialize" (.serialize)
  , versus "deserialize" (.deserialize)
  , versus "view" (.view)
  ]
  where
    versus = Baseline.versus [datasetName, "unfair"] baseline others

middle :: Dataset -> Word32
middle spec = spec.low + (spec.high - spec.low) `div` 2

withIntSet :: Dataset -> (Contender -> Benchmark) -> Benchmark
withIntSet spec continue =
  env (pure (left, right, Binary.encode left)) $ \ ~(intSet, other, encoded) ->
    continue
      Contender
        { name = "IntSet"
        , difference = nf (IntSet.difference intSet) other
        , intersectionSize = whnf (IntSet.size . IntSet.intersection intSet) other
        , jaccard = whnf (jaccardOf intSet) other
        , rangeCount = whnf rangeCountOf intSet
        , rank = whnf rankOf intSet
        , page = nf pageOf intSet
        , shift = nf (IntSet.mapMonotonic (+ offset)) intSet
        , serialize = nf Binary.encode intSet
        , deserialize = nf decode encoded
        , view = nf decode encoded
        }
  where
    left = IntSet.fromList $ map fromIntegral spec.left
    right = IntSet.fromList $ map fromIntegral spec.right
    decode :: Lazy.ByteString -> IntSet
    decode = Binary.decode
    jaccardOf one two = fromIntegral common / fromIntegral (IntSet.size one + IntSet.size two - common) :: Double
      where
        common = IntSet.size $ IntSet.intersection one two
    rangeCountOf set = IntSet.size inside
      where
        (_, above) = IntSet.split (fromIntegral spec.low - 1) set
        (inside, _) = IntSet.split (fromIntegral spec.high) above
    rankOf set = IntSet.size below + fromEnum found
      where
        (below, found, _) = IntSet.splitMember (fromIntegral $ middle spec) set
    pageOf = take pageSize . drop pageOffset . IntSet.toAscList

withBitmap :: Dataset -> (Contender -> Benchmark) -> Benchmark
withBitmap spec continue =
  withResource (bitmapOf spec.left) Bitmap.free $ \getLeft ->
    withResource (bitmapOf spec.right) Bitmap.free $ \getRight ->
      withResource (getLeft >>= serialized) (free . fst) $ \getBytes ->
        let both operation = do
              left <- getLeft
              right <- getRight
              operation left right
        in continue
             Contender
               { name = "Bitmap"
               , difference = whnfIO $ both Bitmap.andnot >>= Bitmap.free
               , intersectionSize = whnfIO $ both Bitmap.andCardinality
               , jaccard = whnfIO $ both Bitmap.jaccardIndex
               , rangeCount = whnfIO $ getLeft >>= \bitmap -> Bitmap.rangeCardinality bitmap (fromIntegral spec.low) $ fromIntegral spec.high
               , rank = whnfIO $ getLeft >>= \bitmap -> Bitmap.rank bitmap $ middle spec
               , page = nfIO $ getLeft >>= pageOf
               , shift = whnfIO $ getLeft >>= \bitmap -> Bitmap.addOffset bitmap (fromIntegral offset) >>= Bitmap.free
               , serialize = whnfIO $ getLeft >>= serializeInto
               , deserialize = whnfIO $ getBytes >>= \(bytes, count) -> Bitmap.portableDeserializeBounded bytes count >>= Bitmap.free
               , view = whnfIO $ getBytes >>= \(bytes, _) -> Bitmap.portableDeserializeFrozen bytes >>= Bitmap.free
               }
  where
    serialized bitmap = do
      count <- Bitmap.portableSizeInBytes bitmap
      bytes <- mallocBytes $ fromIntegral count
      written <- Bitmap.portableSerialize bitmap bytes
      pure (bytes, written)
    serializeInto bitmap = do
      count <- Bitmap.portableSizeInBytes bitmap
      allocaBytes (fromIntegral count) $ Bitmap.portableSerialize bitmap
    pageOf bitmap =
      allocaArray pageSize $ \buffer -> do
        _ <- Bitmap.rangeUint32Array bitmap (fromIntegral pageOffset) (fromIntegral pageSize) buffer
        peekArray pageSize buffer

bitmapOf :: [Word32] -> IO Bitmap
bitmapOf values = withArrayLen values $ \count ptr -> Bitmap.ofPtr (fromIntegral count) ptr

withBitmap64 :: Dataset -> (Contender -> Benchmark) -> Benchmark
withBitmap64 spec continue =
  withResource (bitmap64Of spec.left) Bitmap64.free $ \getLeft ->
    withResource (bitmap64Of spec.right) Bitmap64.free $ \getRight ->
      withResource (getLeft >>= serialized) (free . fst) $ \getBytes ->
        let both operation = do
              left <- getLeft
              right <- getRight
              operation left right
        in continue
             Contender
               { name = "Bitmap64"
               , difference = whnfIO $ both Bitmap64.andnot >>= Bitmap64.free
               , intersectionSize = whnfIO $ both Bitmap64.andCardinality
               , jaccard = whnfIO $ both Bitmap64.jaccardIndex
               , rangeCount = whnfIO $ getLeft >>= \bitmap -> Bitmap64.rangeCardinality bitmap (fromIntegral spec.low) $ fromIntegral spec.high
               , rank = whnfIO $ getLeft >>= \bitmap -> Bitmap64.rank bitmap . fromIntegral $ middle spec
               , page = nfIO $ getLeft >>= pageOf
               , shift = whnfIO $ getLeft >>= \bitmap -> Bitmap64.addOffset bitmap (fromIntegral offset) >>= Bitmap64.free
               , serialize = whnfIO $ getLeft >>= serializeInto
               , deserialize = whnfIO $ getBytes >>= \(bytes, count) -> Bitmap64.portableDeserializeBounded bytes count >>= Bitmap64.free
               , view = whnfIO $ getBytes >>= \(bytes, count) -> Bitmap64.portableDeserializeFrozen bytes count >>= Bitmap64.free
               }
  where
    serialized :: Bitmap64 -> IO (Ptr CChar, CSize)
    serialized bitmap = do
      count <- Bitmap64.portableSizeInBytes bitmap
      bytes <- mallocBytes $ fromIntegral count
      written <- Bitmap64.portableSerialize bitmap bytes
      pure (bytes, written)
    serializeInto bitmap = do
      count <- Bitmap64.portableSizeInBytes bitmap
      allocaBytes (fromIntegral count) $ Bitmap64.portableSerialize bitmap
    pageOf :: Bitmap64 -> IO [Word64]
    pageOf bitmap =
      alloca $ \start -> do
        _ <- Bitmap64.select bitmap (fromIntegral pageOffset) start
        first <- peek start
        bracket (Iterator64.create bitmap) Iterator64.free $ \iterator ->
          allocaArray pageSize $ \buffer -> do
            _ <- Iterator64.moveEqualOrLarger iterator first
            count <- Iterator64.read iterator buffer $ fromIntegral pageSize
            peekArray (fromIntegral count) buffer

bitmap64Of :: [Word32] -> IO Bitmap64
bitmap64Of values = withArrayLen (map fromIntegral values) $ \count ptr -> Bitmap64.ofPtr (fromIntegral count) ptr