Comparison.hs

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

module Bench.Comparison
  ( benchmarks
  , scaling
  ) where

import Bench.Baseline qualified as Baseline
import Bench.Values (dense, scattered)
import Control.DeepSeq (NFData)
import Control.Exception (bracket)
import Data.Bit (Bit (..), countBits, listBits, zipBits)
import Data.Bitmap.Roaring.FFI.Bitmap (Bitmap)
import Data.Bitmap.Roaring.FFI.Bitmap qualified as Bitmap
import Data.Bitmap.Roaring.FFI.Bitmap.Iterator (Iterator (..))
import Data.Bitmap.Roaring.FFI.Bitmap.Iterator qualified as Iterator
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.Bitmap.Roaring.FFI.Bitset (Bitset)
import Data.Bitmap.Roaring.FFI.Bitset qualified as Bitset
import Data.Bits ((.&.), (.|.))
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.Vector.Storable qualified as Storable
import Data.Vector.Unboxed qualified as Unboxed
import Data.Vector.Unboxed.Mutable qualified as Unboxed.Mutable
import Data.Word (Word32, Word64)
import Foreign.C.Types (CSize)
import Foreign.Marshal.Alloc (alloca, allocaBytesAligned)
import Foreign.Marshal.Array (allocaArray)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable, peek, peekElemOff, poke)
import GHC.Generics (Generic)
import Test.Tasty (withResource)
import Test.Tasty.Bench (Benchmark, Benchmarkable, bgroup, env, nf, whnf, whnfIO)

data Dataset = Dataset
  { name :: String
  , universe :: Maybe Int
  , left :: [Word32]
  , right :: [Word32]
  }

data Inputs = Inputs
  { left :: Storable.Vector Word32
  , right :: Storable.Vector Word32
  , probes :: Storable.Vector Word32
  , left64 :: Storable.Vector Word64
  , right64 :: Storable.Vector Word64
  , probes64 :: Storable.Vector Word64
  }
  deriving (Generic)

instance NFData Inputs

data Contender = Contender
  { name :: String
  , construction :: Benchmarkable
  , member :: Benchmarkable
  , cardinality :: Benchmarkable
  , union :: Benchmarkable
  , intersection :: Benchmarkable
  , iteration :: Benchmarkable
  }

size :: Int
size = 1_000_000

benchmarks :: Benchmark
benchmarks =
  bgroup "comparison" $
    map
      (dataset "comparison")
      [ Dataset
          { name = "scattered"
          , universe = Nothing
          , left = scattered 1 size
          , right = scattered 2 size
          }
      , Dataset
          { name = "bounded"
          , universe = Just 0x1000000
          , left = map (.&. 0xFFFFFF) $ scattered 1 size
          , right = map (.&. 0xFFFFFF) $ scattered 2 size
          }
      , Dataset
          { name = "dense"
          , universe = Just 0x200000
          , left = dense size
          , right = map (+ fromIntegral (size `div` 2)) $ dense size
          }
      ]

scaling :: Benchmark
scaling =
  bgroup "scaling" $ map (dataset "scaling" . sparse) [1_000, 3_000, 10_000, 30_000, 100_000, 300_000, 1_000_000]
  where
    sparse count =
      Dataset
        { name = show count
        , universe = Nothing
        , left = scattered 1 count
        , right = scattered 2 count
        }

dataset :: String -> Dataset -> Benchmark
dataset group spec =
  env (pure prepared) $ \inputs ->
    withIntSet inputs $ \intSet ->
      withBitmap inputs $ \bitmap ->
        withBitmap64 inputs $ \bitmap64 ->
          withBitVectors spec.universe inputs $ \bitVectors ->
            bgroup spec.name . contests group spec.name intSet $ bitmap : bitmap64 : bitVectors
  where
    left = Storable.fromList spec.left
    right = Storable.fromList spec.right
    probes = Storable.fromList . concat $ zipWith (\hit miss -> [hit, miss]) spec.left spec.right
    prepared =
      Inputs
        { left
        , right
        , probes
        , left64 = Storable.map fromIntegral left
        , right64 = Storable.map fromIntegral right
        , probes64 = Storable.map fromIntegral probes
        }

contests :: String -> String -> Contender -> [Contender] -> [Benchmark]
contests group datasetName baseline others =
  [ versus "construction" (.construction)
  , versus "member" (.member)
  , versus "cardinality" (.cardinality)
  , versus "union" (.union)
  , versus "intersection" (.intersection)
  , versus "iteration" (.iteration)
  ]
  where
    versus = Baseline.versus [datasetName, group] baseline others

withBitmap :: Inputs -> (Contender -> Benchmark) -> Benchmark
withBitmap inputs continue =
  withResource (bitmapOf inputs.left) Bitmap.free $ \getLeft ->
    withResource (bitmapOf inputs.right) Bitmap.free $ \getRight ->
      continue
        Contender
          { name = "Bitmap"
          , construction = whnfIO $ bitmapOf inputs.left >>= Bitmap.free
          , member = whnfIO $ do
              bitmap <- getLeft
              tallyIO (Bitmap.contains bitmap) inputs.probes
          , cardinality = whnfIO $ getLeft >>= Bitmap.getCardinality
          , union = whnfIO $ combined getLeft getRight Bitmap.or
          , intersection = whnfIO $ combined getLeft getRight Bitmap.and
          , iteration = whnfIO $ getLeft >>= sumBitmap
          }
  where
    combined getLeft getRight operation = do
      left <- getLeft
      right <- getRight
      operation left right >>= Bitmap.free

bitmapOf :: Storable.Vector Word32 -> IO Bitmap
bitmapOf values = Storable.unsafeWith values . Bitmap.ofPtr . fromIntegral $ Storable.length values

sumBitmap :: Bitmap -> IO Word64
sumBitmap bitmap =
  allocaBytesAligned Iterator.size Iterator.alignment $ \memory ->
    allocaArray chunk $ \buffer -> do
      let
        iterator = Iterator memory
        loop !accumulated = do
          count <- fromIntegral <$> Iterator.read iterator buffer (fromIntegral chunk)
          summed <- sumBuffer buffer count accumulated
          if count < chunk then pure summed else loop summed
      Iterator.init bitmap iterator
      loop 0
  where
    chunk = 256

withBitmap64 :: Inputs -> (Contender -> Benchmark) -> Benchmark
withBitmap64 inputs continue =
  withResource (bitmap64Of inputs.left64) Bitmap64.free $ \getLeft ->
    withResource (bitmap64Of inputs.right64) Bitmap64.free $ \getRight ->
      continue
        Contender
          { name = "Bitmap64"
          , construction = whnfIO $ bitmap64Of inputs.left64 >>= Bitmap64.free
          , member = whnfIO $ do
              bitmap <- getLeft
              tallyIO (Bitmap64.contains bitmap) inputs.probes64
          , cardinality = whnfIO $ getLeft >>= Bitmap64.getCardinality
          , union = whnfIO $ combined getLeft getRight Bitmap64.or
          , intersection = whnfIO $ combined getLeft getRight Bitmap64.and
          , iteration = whnfIO $ getLeft >>= sumBitmap64
          }
  where
    combined getLeft getRight operation = do
      left <- getLeft
      right <- getRight
      operation left right >>= Bitmap64.free

bitmap64Of :: Storable.Vector Word64 -> IO Bitmap64
bitmap64Of values = Storable.unsafeWith values . Bitmap64.ofPtr . fromIntegral $ Storable.length values

sumBitmap64 :: Bitmap64 -> IO Word64
sumBitmap64 bitmap =
  bracket (Iterator64.create bitmap) Iterator64.free $ \iterator ->
    allocaArray chunk $ \buffer -> do
      let loop !accumulated = do
            count <- fromIntegral <$> Iterator64.read iterator buffer (fromIntegral chunk)
            summed <- sumBuffer buffer count accumulated
            if count < chunk then pure summed else loop summed
      loop 0
  where
    chunk = 256

withIntSet :: Inputs -> (Contender -> Benchmark) -> Benchmark
withIntSet inputs continue =
  env (pure (intSetOf inputs.left, intSetOf inputs.right)) $ \ ~(left, right) ->
    continue
      Contender
        { name = "IntSet"
        , construction = nf intSetOf inputs.left
        , member = whnf (countMembers inputs.probes) left
        , cardinality = whnf IntSet.size left
        , union = nf (IntSet.union left) right
        , intersection = nf (IntSet.intersection left) right
        , iteration = whnf (IntSet.foldl' (+) 0) left
        }
  where
    countMembers probes set = tally isMember probes
      where
        isMember value = IntSet.member (fromIntegral value) set

intSetOf :: Storable.Vector Word32 -> IntSet
intSetOf = Storable.foldl' (\set value -> IntSet.insert (fromIntegral value) set) IntSet.empty

withBitVectors :: Maybe Int -> Inputs -> ([Contender] -> Benchmark) -> Benchmark
withBitVectors universe inputs continue =
  case universe of
    Nothing -> continue []
    Just bits ->
      withBitset bits inputs $ \bitset ->
        withBitvec bits inputs $ \bitvec ->
          continue [bitset, bitvec]

withBitset :: Int -> Inputs -> (Contender -> Benchmark) -> Benchmark
withBitset universe inputs continue =
  withResource (bitsetOf universe inputs.left) Bitset.free $ \getLeft ->
    withResource (bitsetOf universe inputs.right) Bitset.free $ \getRight ->
      continue
        Contender
          { name = "Bitset"
          , construction = whnfIO $ bitsetOf universe inputs.left >>= Bitset.free
          , member = whnfIO $ do
              bitset <- getLeft
              tallyIO (Bitset.get bitset . fromIntegral) inputs.probes
          , cardinality = whnfIO $ getLeft >>= Bitset.count
          , union = whnfIO $ combined getLeft getRight Bitset.inplaceUnion
          , intersection = whnfIO $ combined getLeft getRight Bitset.inplaceIntersection
          , iteration = whnfIO $ getLeft >>= sumBitset
          }
  where
    combined getLeft getRight operation = do
      left <- getLeft
      right <- getRight
      copy <- Bitset.copy left
      _ <- operation copy right
      Bitset.free copy

bitsetOf :: Int -> Storable.Vector Word32 -> IO Bitset
bitsetOf universe values = do
  bitset <- Bitset.createWithCapacity $ fromIntegral universe
  Storable.mapM_ (Bitset.set bitset . fromIntegral) values
  pure bitset

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

withBitvec :: Int -> Inputs -> (Contender -> Benchmark) -> Benchmark
withBitvec universe inputs continue =
  env (pure (bitvecOf universe inputs.left, bitvecOf universe inputs.right)) $ \ ~(left, right) ->
    continue
      Contender
        { name = "bitvec"
        , construction = whnf (bitvecOf universe) inputs.left
        , member = whnf (countMembers inputs.probes) left
        , cardinality = whnf countBits left
        , union = whnf (zipBits (.|.) left) right
        , intersection = whnf (zipBits (.&.) left) right
        , iteration = whnf (foldl' (+) 0 . listBits) left
        }
  where
    countMembers probes bits = tally isSet probes
      where
        isSet value = unBit . Unboxed.unsafeIndex bits $ fromIntegral value

bitvecOf :: Int -> Storable.Vector Word32 -> Unboxed.Vector Bit
bitvecOf universe values =
  Unboxed.create $ do
    bits <- Unboxed.Mutable.replicate universe $ Bit False
    Storable.forM_ values $ \value -> Unboxed.Mutable.unsafeWrite bits (fromIntegral value) $ Bit True
    pure bits

tallyIO :: (Storable a) => (a -> IO Bool) -> Storable.Vector a -> IO Int
tallyIO query values = go 0 0
  where
    go !hits index
      | index == Storable.length values = pure hits
      | otherwise = do
          hit <- query $ Storable.unsafeIndex values index
          go (counted hits hit) $ index + 1
{-# INLINE tallyIO #-}

tally :: (Storable a) => (a -> Bool) -> Storable.Vector a -> Int
tally query values = go 0 0
  where
    go !hits index
      | index == Storable.length values = hits
      | otherwise = go (counted hits . query $ Storable.unsafeIndex values index) $ index + 1
{-# INLINE tally #-}

counted :: Int -> Bool -> Int
counted hits hit = if hit then hits + 1 else hits
{-# INLINE counted #-}

sumBuffer :: (Storable a, Integral a, Num total) => Ptr a -> Int -> total -> IO total
sumBuffer buffer count = go 0
  where
    go index !accumulated
      | index == count = pure accumulated
      | otherwise = do
          value <- peekElemOff buffer index
          go (index + 1) $ accumulated + fromIntegral value
{-# INLINE sumBuffer #-}