Sizes.hs

module Main (main) where

import Bench.Values (dense, scattered)
import Control.Exception (bracket, evaluate)
import Control.Monad (forM_)
import Data.Binary qualified as Binary
import Data.Bit (Bit (..))
import Data.Bitmap.Roaring.FFI.Bitmap qualified as Bitmap
import Data.Bitmap.Roaring.FFI.Bitmap64 qualified as Bitmap64
import Data.Bitmap.Roaring.FFI.Bitset qualified as Bitset
import Data.Bits ((.&.))
import Data.ByteString.Lazy qualified as Lazy
import Data.IntSet qualified as IntSet
import Data.List (intercalate)
import Data.Vector.Unboxed qualified as Unboxed
import Data.Vector.Unboxed.Mutable qualified as Unboxed.Mutable
import Data.Word (Word32)
import Foreign.C.Types (CSize (..))
import Foreign.Marshal.Array (withArrayLen)
import GHC.Compact (compact, compactSize)
import System.Environment (getArgs)
import System.Exit (die)

foreign import ccall unsafe "croaring_ffi_count_allocations"
  countAllocations :: IO ()

foreign import ccall unsafe "croaring_ffi_allocated_bytes"
  allocatedBytes :: IO CSize

size :: Int
size = 1_000_000

main :: IO ()
main = do
  arguments <- getArgs
  case arguments of
    [] -> table >>= putStr
    [output] -> table >>= writeFile output
    _ -> die "usage: croaring-ffi-sizes [OUTPUT.csv]"

table :: IO String
table = do
  countAllocations
  rows <- concat <$> mapM measure datasets
  pure . unlines $ "dataset,structure,bytes,values" : map (intercalate ",") rows

datasets :: [(String, [Word32], Bool)]
datasets =
  [ ("sparse", scattered 1 size, False)
  , ("bounded", map (.&. 0xFFFFFF) $ scattered 1 size, True)
  , ("dense", dense size, True)
  ]

measure :: (String, [Word32], Bool) -> IO [[String]]
measure (name, values, bitArrays) = do
  intSet <- evaluate . IntSet.fromList $ map fromIntegral values
  let
    count = IntSet.size intSet
    row structure bytes = [name, structure, show bytes, show count]
  intSetMemory <- compact intSet >>= compactSize
  let intSetBinary = Lazy.length $ Binary.encode intSet
  (roaring32Memory, roaring32Portable) <- roaring32 values
  (roaring64Memory, roaring64Portable) <- roaring64 values
  bitArrayRows <-
    if bitArrays then do
      bitsetMemory <- bitset values
      bitvecMemory <- bitvec values
      pure [row "bitset-memory" bitsetMemory, row "bitvec-memory" bitvecMemory]
    else
      pure []
  pure $
    [ row "intset-memory" $ toInteger intSetMemory
    , row "intset-binary" $ toInteger intSetBinary
    , row "roaring32-memory" roaring32Memory
    , row "roaring32-portable" roaring32Portable
    , row "roaring64-memory" roaring64Memory
    , row "roaring64-portable" roaring64Portable
    ]
      <> bitArrayRows

allocatedBy :: IO a -> (a -> IO ()) -> (a -> IO b) -> IO (Integer, b)
allocatedBy create release inspect = do
  before <- allocatedBytes
  bracket create release $ \structure -> do
    after <- allocatedBytes
    result <- inspect structure
    pure (toInteger after - toInteger before, result)

roaring32 :: [Word32] -> IO (Integer, Integer)
roaring32 values = allocatedBy create Bitmap.free $ fmap toInteger . Bitmap.portableSizeInBytes
  where
    create = do
      bitmap <- withArrayLen values $ \count ptr -> Bitmap.ofPtr (fromIntegral count) ptr
      _ <- Bitmap.runOptimize bitmap
      _ <- Bitmap.shrinkToFit bitmap
      pure bitmap

roaring64 :: [Word32] -> IO (Integer, Integer)
roaring64 values = allocatedBy create Bitmap64.free $ fmap toInteger . Bitmap64.portableSizeInBytes
  where
    create = do
      bitmap <- withArrayLen (map fromIntegral values) $ \count ptr -> Bitmap64.ofPtr (fromIntegral count) ptr
      _ <- Bitmap64.runOptimize bitmap
      _ <- Bitmap64.shrinkToFit bitmap
      pure bitmap

bitset :: [Word32] -> IO Integer
bitset values = fst <$> allocatedBy create Bitset.free (const $ pure ())
  where
    create = do
      structure <- Bitset.create
      forM_ values $ Bitset.set structure . fromIntegral
      _ <- Bitset.trim structure
      pure structure

bitvec :: [Word32] -> IO Integer
bitvec values = do
  bits <- evaluate built
  pure . toInteger $ 8 * ((Unboxed.length bits + 63) `div` 64)
  where
    built =
      Unboxed.create $ do
        vector <- Unboxed.Mutable.replicate (fromIntegral (maximum values) + 1) $ Bit False
        forM_ values $ \value -> Unboxed.Mutable.unsafeWrite vector (fromIntegral value) $ Bit True
        pure vector