Bitset.hs

module Test.Bitset (tests) where

import Control.Exception (try)
import Control.Monad (forM, forM_)
import Data.Bitmap.Roaring.FFI.Bitset (Bitset (..))
import Data.Bitmap.Roaring.FFI.Bitset qualified as Bitset
import Data.Bitmap.Roaring.FFI.NullHandle (NullHandle (..))
import Data.IORef (readIORef)
import Data.Set qualified as Set
import Data.Word (Word64)
import Foreign.C.Types (CSize)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Marshal.Array (allocaArray, peekArray)
import Foreign.Ptr (nullPtr)
import Foreign.Storable (peek, poke)
import Hedgehog (Gen, Property, assert, evalIO, forAll, (===))
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Test.Bitmap (collect, visiting)
import Test.Property (Test, returns)
import Test.Property qualified as Property
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Hedgehog (testProperty)

tests :: TestTree
tests =
  testGroup
    "Bitset"
    [ testProperty "set bits and queries" queries
    , testProperty "set algebra" setAlgebra
    , testProperty "shifts" shifts
    , testProperty "buffer fields" fields
    , testProperty "forEach stops when the callback says so" forEachStop
    , testProperty "resize, trim and fill" sizing
    , testProperty "with frees the bitset and rejects NULL" scoped
    , testProperty "forEachWith visits, stops and rethrows" callbackHelpers
    , testProperty "in-place operations with the same bitset twice" selfAliasing
    ]

genBits :: CSize -> Gen [CSize]
genBits limit = Gen.list (Range.linear 0 200) . fmap fromIntegral . Gen.word64 . Range.constant 0 $ fromIntegral limit

created :: IO Bitset -> Test Bitset
created = Property.owned (Bitset nullPtr) Bitset.free

fromBits :: [CSize] -> Test Bitset
fromBits bits = do
  bitset <- created Bitset.create
  evalIO . forM_ bits $ Bitset.set bitset
  pure bitset

toList :: Bitset -> IO [CSize]
toList bitset =
  allocaArray 5 $ \buffer ->
    alloca $ \start -> do
      poke start 0
      let loop = do
            count <- fromIntegral <$> Bitset.nextSetBits bitset buffer 5 start
            found <- peekArray count buffer
            if count == 0 then
              pure []
            else do
              peek start >>= poke start . (+ 1)
              (found <>) <$> loop
      loop

queries :: Property
queries =
  Property.generated $ do
    bits <- forAll $ genBits 5000
    probes <- forAll $ genBits 6000
    bitset <- fromBits bits
    let
      model = Set.fromList bits
      next probe =
        alloca $ \index -> do
          poke index probe
          present <- Bitset.nextSetBit bitset index
          if present then Just <$> peek index else pure Nothing
    toList bitset `returns` Set.toAscList model
    Bitset.count bitset `returns` fromIntegral (Set.size model)
    Bitset.empty bitset `returns` Set.null model
    Bitset.minimum bitset `returns` maybe maxBound fst (Set.minView model)
    Bitset.maximum bitset `returns` maybe 0 fst (Set.maxView model)
    forM probes (Bitset.get bitset) `returns` map (`Set.member` model) probes
    forM probes next `returns` map (`Set.lookupGE` model) probes

setAlgebra :: Property
setAlgebra =
  Property.generated $ do
    leftBits <- forAll $ genBits 3000
    rightBits <- forAll $ genBits 3000
    left <- fromBits leftBits
    right <- fromBits rightBits
    let
      leftModel = Set.fromList leftBits
      rightModel = Set.fromList rightBits
      union = Set.union leftModel rightModel
      intersection = Set.intersection leftModel rightModel
      difference = Set.difference leftModel rightModel
      symmetric = union `Set.difference` intersection
      applied operation = do
        copy <- fromBits leftBits
        _ <- evalIO $ operation copy right
        evalIO $ toList copy
    Bitset.unionCount left right `returns` fromIntegral (Set.size union)
    Bitset.intersectionCount left right `returns` fromIntegral (Set.size intersection)
    Bitset.differenceCount left right `returns` fromIntegral (Set.size difference)
    Bitset.symmetricDifferenceCount left right `returns` fromIntegral (Set.size symmetric)
    Bitset.disjoint left right `returns` Set.disjoint leftModel rightModel
    Bitset.intersect left right `returns` not (Set.disjoint leftModel rightModel)
    Bitset.containsAll left right `returns` Set.isSubsetOf rightModel leftModel
    applied Bitset.inplaceUnion >>= (=== Set.toAscList union)
    applied Bitset.inplaceIntersection >>= (=== Set.toAscList intersection)
    applied Bitset.inplaceDifference >>= (=== Set.toAscList difference)
    applied Bitset.inplaceSymmetricDifference >>= (=== Set.toAscList symmetric)
    copied <- created $ Bitset.copy left
    toList copied `returns` Set.toAscList leftModel

shifts :: Property
shifts =
  Property.generated $ do
    bits <- forAll $ genBits 3000
    shift <- forAll . fmap fromIntegral . Gen.word64 $ Range.linear 0 200
    let ascending = Set.toAscList $ Set.fromList bits
    left <- fromBits bits
    evalIO $ Bitset.shiftLeft left shift
    toList left `returns` map (+ shift) ascending
    right <- fromBits bits
    evalIO $ Bitset.shiftRight right shift
    toList right `returns` [bit - shift | bit <- ascending, bit >= shift]

fields :: Property
fields =
  Property.example $ do
    bitset <- fromBits [0, 63, 64, 200]
    used <- evalIO $ Bitset.sizeInWords bitset
    Bitset.arraySize bitset `returns` used
    Bitset.sizeInBits bitset `returns` (used * 64)
    Bitset.sizeInBytes bitset `returns` (used * 8)
    capacity <- evalIO $ Bitset.capacity bitset
    assert $ capacity >= used
    take 4 <$> (Bitset.array bitset >>= peekArray (fromIntegral used)) `returns` [2 ^ (63 :: Int) + 1, 1, 0, 2 ^ (200 - 192 :: Int) :: Word64]

forEachStop :: Property
forEachStop =
  Property.example $ do
    bitset <- fromBits [1, 5, 9, 700]
    (seen, callback) <- collect Bitset.wrapCallback (< 9)
    Bitset.forEach bitset callback nullPtr `returns` False
    reverse <$> readIORef seen `returns` [1, 5, 9]

sizing :: Property
sizing =
  Property.example $ do
    bitset <- created $ Bitset.createWithCapacity 128
    Bitset.sizeInWords bitset `returns` 2
    Bitset.count bitset `returns` 0
    evalIO $ Bitset.fill bitset
    Bitset.count bitset `returns` 128
    evalIO $ Bitset.setToValue bitset 3 False
    Bitset.get bitset 3 `returns` False
    Bitset.resize bitset 4 True `returns` True
    Bitset.sizeInWords bitset `returns` 4
    Bitset.count bitset `returns` 127
    Bitset.grow bitset 8 `returns` True
    Bitset.sizeInWords bitset `returns` 8
    evalIO $ Bitset.clear bitset
    Bitset.empty bitset `returns` True
    Bitset.trim bitset `returns` True

scoped :: Property
scoped =
  Property.example $ do
    Bitset.with (Bitset.createWithCapacity 64) (\bitset -> Bitset.set bitset 5 >> Bitset.count bitset) `returns` 1
    outcome <- evalIO . try $ Bitset.with (pure $ Bitset nullPtr) Bitset.count
    outcome === Left (NullHandle "Data.Bitmap.Roaring.FFI.Bitset.with")

callbackHelpers :: Property
callbackHelpers =
  Property.example $ do
    bitset <- fromBits [1 .. 10]
    visiting (Bitset.forEachWith bitset) (const $ pure True) `returns` (Right True, [1 .. 10])
    visiting (Bitset.forEachWith bitset) (pure . (< 3)) `returns` (Right False, [1, 2, 3])
    visiting (Bitset.forEachWith bitset) stopAtThree `returns` (Left $ userError "stop", [1, 2, 3])
  where
    stopAtThree bit = if bit == 3 then ioError $ userError "stop" else pure True

selfAliasing :: Property
selfAliasing =
  Property.generated $ do
    bits <- forAll $ genBits 3000
    let
      expected = Set.toAscList $ Set.fromList bits
      applied operation = do
        bitset <- fromBits bits
        result <- evalIO $ operation bitset bitset
        (,) result <$> evalIO (toList bitset)
    applied Bitset.inplaceUnion >>= (=== (True, expected))
    applied Bitset.inplaceIntersection >>= (=== ((), expected))
    applied Bitset.inplaceDifference >>= (=== ((), []))
    applied Bitset.inplaceSymmetricDifference >>= (=== (True, []))