Bitmap64.hs

{-# LANGUAGE OverloadedRecordDot #-}

module Test.Bitmap64
  ( tests
  , genValues64
  , fromValues64
  ) where

import Control.Exception (try)
import Control.Monad (forM, forM_)
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.Statistics (Statistics (..))
import Data.Bitmap.Roaring.FFI.NullHandle (NullHandle (..))
import Data.IORef (readIORef)
import Data.Maybe (fromMaybe)
import Data.Set qualified as Set
import Data.Word (Word64)
import Foreign.C.Types (CChar, CSize)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Marshal.Array (allocaArray, peekArray, withArrayLen)
import Foreign.Marshal.Utils (fillBytes)
import Foreign.Ptr (Ptr, nullPtr, plusPtr)
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, fromValues, modelRank, symmetricDifference, 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
    "Bitmap64"
    [ testProperty "bool results in both polarities" boolResults
    , testGroup
        "model"
        [ testProperty "contents and queries" contentsAndQueries
        , testProperty "set algebra" setAlgebra
        , testProperty "in-place set algebra" inplaceSetAlgebra
        , testProperty "ranges" ranges
        , testProperty "offsets" offsets
        , testProperty "move from a 32-bit bitmap" moveFromRoaring32
        ]
    , testGroup
        "serialization"
        [ testProperty "portable round-trip" portableRoundTrip
        , testProperty "portable frozen round-trip" portableFrozenRoundTrip
        , testProperty "frozen view round-trip" frozenRoundTrip
        , testProperty "bounded portable load rejects a truncated buffer" truncatedPortable
        ]
    , testGroup
        "layout"
        [ testProperty "statistics of known containers" knownStatistics
        , testProperty "statistics poke and peek agree" statisticsStorable
        , testProperty "bulk context in caller memory" bulkContext
        ]
    , testProperty "iterate visits every value" iterateAll
    , testGroup
        "helpers"
        [ testProperty "with frees the bitmap and rejects NULL" scoped
        , testProperty "withPortableFrozen views valid bytes in place" portableFrozenScope
        , testProperty "withFrozenView views aligned bytes in place" frozenViewScope
        , testProperty "iterateWith visits, stops and rethrows" callbackHelpers
        , testProperty "in-place operations with the same bitmap twice" selfAliasing
        , testProperty "BulkContext.with" scopedBulkContext
        ]
    ]

genValues64 :: Gen [Word64]
genValues64 = concat <$> Gen.list (Range.linear 0 60) (Gen.choice [pure <$> genValue64, genRun])
  where
    genRun = do
      start <- genValue64
      len <- Gen.int $ Range.linear 0 3000
      pure . take len $ enumFrom start

genValue64 :: Gen Word64
genValue64 =
  Gen.choice
    [ Gen.word64 $ Range.constant 0 200000
    , Gen.word64 Range.constantBounded
    , (+) . (* 2 ^ (32 :: Int)) <$> Gen.word64 (Range.constant 1 3) <*> Gen.word64 (Range.constant 0 200000)
    , Gen.word64 $ Range.constant (maxBound - 200000) maxBound
    ]

adopt64 :: IO Bitmap64 -> Test Bitmap64
adopt64 = Property.adopt (Bitmap64 nullPtr) Bitmap64.free

owned64 :: IO Bitmap64 -> Test Bitmap64
owned64 = Property.owned (Bitmap64 nullPtr) Bitmap64.free

fromValues64 :: [Word64] -> Test Bitmap64
fromValues64 values = owned64 . withArrayLen values $ \count ptr -> Bitmap64.ofPtr (fromIntegral count) ptr

elements64 :: Bitmap64 -> IO [Word64]
elements64 bitmap = do
  count <- fromIntegral <$> Bitmap64.getCardinality bitmap
  allocaArray count $ \ptr -> do
    Bitmap64.toUint64Array bitmap ptr
    peekArray count ptr

toList64 :: Bitmap64 -> Test [Word64]
toList64 = evalIO . elements64

boolResults :: Property
boolResults =
  Property.example $ do
    bitmap <- owned64 Bitmap64.create
    Bitmap64.addChecked bitmap (2 ^ (40 :: Int)) `returns` True
    Bitmap64.addChecked bitmap (2 ^ (40 :: Int)) `returns` False
    Bitmap64.removeChecked bitmap (2 ^ (40 :: Int)) `returns` True
    Bitmap64.removeChecked bitmap (2 ^ (40 :: Int)) `returns` False
    Bitmap64.isEmpty bitmap `returns` True

contentsAndQueries :: Property
contentsAndQueries =
  Property.generated $ do
    values <- forAll genValues64
    probes <- forAll $ Gen.list (Range.linear 0 50) genValue64
    bitmap <- fromValues64 values
    let model = Set.fromList values
    toList64 bitmap >>= (=== Set.toAscList model)
    Bitmap64.getCardinality bitmap `returns` fromIntegral (Set.size model)
    Bitmap64.isEmpty bitmap `returns` Set.null model
    Bitmap64.minimum bitmap `returns` maybe maxBound fst (Set.minView model)
    Bitmap64.maximum bitmap `returns` maybe 0 fst (Set.maxView model)
    Bitmap64.internalValidate bitmap nullPtr `returns` True
    forM (values <> probes) (probe bitmap) `returns` map (modelProbe model) (values <> probes)
    forM [0 .. fromIntegral (Set.size model)] (select bitmap) `returns` (map Just (Set.toAscList model) <> [Nothing])
  where
    probe bitmap value = (,,) <$> Bitmap64.contains bitmap value <*> Bitmap64.rank bitmap value <*> index bitmap value
    index bitmap value =
      alloca $ \out -> do
        found <- Bitmap64.getIndex bitmap value out
        if found then Just <$> peek out else pure Nothing
    modelProbe model value =
      ( Set.member value model
      , modelRank model value
      , fromIntegral <$> Set.lookupIndex value model
      )
    select bitmap rank =
      alloca $ \out -> do
        found <- Bitmap64.select bitmap rank out
        if found then Just <$> peek out else pure Nothing

setAlgebra :: Property
setAlgebra =
  Property.generated $ do
    leftValues <- forAll genValues64
    rightValues <- forAll genValues64
    left <- fromValues64 leftValues
    right <- fromValues64 rightValues
    let
      leftModel = Set.fromList leftValues
      rightModel = Set.fromList rightValues
      intersection = Set.intersection leftModel rightModel
      union = Set.union leftModel rightModel
      difference = Set.difference leftModel rightModel
      symmetric = symmetricDifference leftModel rightModel
    owned64 (Bitmap64.and left right) >>= toList64 >>= (=== Set.toAscList intersection)
    owned64 (Bitmap64.or left right) >>= toList64 >>= (=== Set.toAscList union)
    owned64 (Bitmap64.andnot left right) >>= toList64 >>= (=== Set.toAscList difference)
    owned64 (Bitmap64.xor left right) >>= toList64 >>= (=== Set.toAscList symmetric)
    Bitmap64.andCardinality left right `returns` fromIntegral (Set.size intersection)
    Bitmap64.orCardinality left right `returns` fromIntegral (Set.size union)
    Bitmap64.andnotCardinality left right `returns` fromIntegral (Set.size difference)
    Bitmap64.xorCardinality left right `returns` fromIntegral (Set.size symmetric)
    Bitmap64.intersect left right `returns` not (Set.disjoint leftModel rightModel)
    Bitmap64.isSubset left right `returns` Set.isSubsetOf leftModel rightModel
    Bitmap64.isStrictSubset left right `returns` Set.isProperSubsetOf leftModel rightModel
    Bitmap64.equals left right `returns` (leftModel == rightModel)

inplaceSetAlgebra :: Property
inplaceSetAlgebra =
  Property.generated $ do
    leftValues <- forAll genValues64
    rightValues <- forAll genValues64
    right <- fromValues64 rightValues
    let
      leftModel = Set.fromList leftValues
      rightModel = Set.fromList rightValues
      applied :: (Bitmap64 -> Bitmap64 -> IO ()) -> Test [Word64]
      applied operation = do
        left <- fromValues64 leftValues
        evalIO $ operation left right
        toList64 left
    applied Bitmap64.andInplace >>= (=== Set.toAscList (Set.intersection leftModel rightModel))
    applied Bitmap64.orInplace >>= (=== Set.toAscList (Set.union leftModel rightModel))
    applied Bitmap64.andnotInplace >>= (=== Set.toAscList (Set.difference leftModel rightModel))
    applied Bitmap64.xorInplace >>= (=== Set.toAscList (symmetricDifference leftModel rightModel))
    applied Bitmap64.overwrite >>= (=== Set.toAscList rightModel)

ranges :: Property
ranges =
  Property.generated $ do
    values <- forAll genValues64
    start <- forAll genValue64
    len <- forAll . Gen.word64 $ Range.linear 0 100000
    bitmap <- fromValues64 values
    let
      model = Set.fromList values
      end = if maxBound - start < len then maxBound else start + len
      spanned = Set.fromList . take (fromIntegral $ end - start) $ enumFrom start
      inRange = Set.intersection model spanned
      modified :: (Bitmap64 -> Word64 -> Word64 -> IO ()) -> Test [Word64]
      modified operation = do
        copy <- fromValues64 values
        evalIO $ operation copy start end
        toList64 copy
    Bitmap64.rangeCardinality bitmap start end `returns` fromIntegral (Set.size inRange)
    Bitmap64.containsRange bitmap start end `returns` (inRange == spanned)
    Bitmap64.intersectWithRange bitmap start end `returns` not (Set.null inRange)
    owned64 (Bitmap64.flip bitmap start end) >>= toList64 >>= (=== Set.toAscList (symmetricDifference model spanned))
    modified Bitmap64.flipInplace >>= (=== Set.toAscList (symmetricDifference model spanned))
    modified Bitmap64.addRange >>= (=== Set.toAscList (Set.union model spanned))
    modified Bitmap64.removeRange >>= (=== Set.toAscList (Set.difference model spanned))

offsets :: Property
offsets =
  Property.generated $ do
    values <- forAll genValues64
    shift <- forAll genValue64
    bitmap <- fromValues64 values
    let
      ascending = Set.toAscList $ Set.fromList values
      down = [value - shift | value <- ascending, value >= shift]
    owned64 (Bitmap64.addOffset bitmap shift) >>= toList64 >>= (=== [value + shift | value <- ascending, value <= maxBound - shift])
    owned64 (Bitmap64.subOffset bitmap shift) >>= toList64 >>= (=== down)
    owned64 (Bitmap64.addOffsetSigned bitmap False shift) >>= toList64 >>= (=== down)

moveFromRoaring32 :: Property
moveFromRoaring32 =
  Property.generated $ do
    values <- forAll . Gen.list (Range.linear 0 200) $ Gen.word32 Range.constantBounded
    bitmap <- fromValues values
    owned64 (Bitmap64.moveFromRoaring32 bitmap) >>= toList64 >>= (=== map fromIntegral (Set.toAscList $ Set.fromList values))
    Bitmap.isEmpty bitmap `returns` True

serialized :: (Bitmap64 -> IO CSize) -> (Bitmap64 -> Ptr CChar -> IO CSize) -> Bitmap64 -> Test (Ptr CChar, CSize, CSize)
serialized measure write bitmap = do
  size <- evalIO $ measure bitmap
  buffer <- Property.alignedBytes (fromIntegral size) 64
  written <- evalIO $ write bitmap buffer
  pure (buffer, size, written)

portableRoundTrip :: Property
portableRoundTrip =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    (buffer, size, written) <- serialized Bitmap64.portableSizeInBytes Bitmap64.portableSerialize bitmap
    written === size
    Bitmap64.portableDeserializeSize buffer written `returns` written
    owned64 (Bitmap64.portableDeserializeBounded buffer written) >>= toList64 >>= (=== Set.toAscList (Set.fromList values))
    owned64 (fromMaybe (Bitmap64 nullPtr) <$> Bitmap64.portableDeserializeValidated buffer written) >>= toList64 >>= (=== Set.toAscList (Set.fromList values))

portableFrozenRoundTrip :: Property
portableFrozenRoundTrip =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    (buffer, _, written) <- serialized Bitmap64.portableSizeInBytes Bitmap64.portableSerialize bitmap
    owned64 (Bitmap64.portableDeserializeFrozen buffer written) >>= toList64 >>= (=== Set.toAscList (Set.fromList values))

frozenRoundTrip :: Property
frozenRoundTrip =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    _ <- evalIO $ Bitmap64.shrinkToFit bitmap
    (buffer, size, written) <- serialized Bitmap64.frozenSizeInBytes Bitmap64.frozenSerialize bitmap
    written === size
    owned64 (Bitmap64.frozenView buffer written) >>= toList64 >>= (=== Set.toAscList (Set.fromList values))
    adopt64 (Bitmap64.frozenView (buffer `plusPtr` 1) written) >>= (=== Bitmap64 nullPtr)

truncatedPortable :: Property
truncatedPortable =
  Property.example $ do
    bitmap <- fromValues64 [1, 2, 3, 2 ^ (40 :: Int)]
    (buffer, _, written) <- serialized Bitmap64.portableSizeInBytes Bitmap64.portableSerialize bitmap
    Bitmap64.portableDeserializeSize buffer (written - 1) `returns` 0
    adopt64 (Bitmap64.portableDeserializeBounded buffer $ written - 1) >>= (=== Bitmap64 nullPtr)
    Bitmap64.portableDeserializeValidated buffer (written - 1) `returns` Nothing

knownStatistics :: Property
knownStatistics =
  Property.example $ do
    bitmap <- owned64 $ Bitmap64.fromRange 131072 (131072 + 65536) 2
    evalIO $ do
      forM_ [0, 2 .. 18] $ Bitmap64.add bitmap
      Bitmap64.addRange bitmap 65536 $ 65536 + 10000
    Bitmap64.runOptimize bitmap `returns` True
    statistics <- evalIO . alloca $ \ptr -> Bitmap64.statistics bitmap ptr >> peek ptr
    Bitmap64.getStatistics bitmap `returns` statistics
    statistics
      === Statistics
        { containers = 3
        , arrayContainers = 1
        , runContainers = 1
        , bitsetContainers = 1
        , valuesArrayContainers = 10
        , valuesRunContainers = 10000
        , valuesBitsetContainers = 32768
        , bytesArrayContainers = statistics.bytesArrayContainers
        , bytesRunContainers = statistics.bytesRunContainers
        , bytesBitsetContainers = 8192
        , maxValue = 131072 + 65534
        , minValue = 0
        , cardinality = 10 + 10000 + 32768
        }
    assert $ statistics.bytesArrayContainers >= 20
    assert $ statistics.bytesRunContainers >= 4

statisticsStorable :: Property
statisticsStorable =
  Property.example $ do
    let statistics = Statistics 1 2 3 4 5 6 7 8 9 10 11 12 13
    ptr <- Property.array 1
    evalIO $ poke ptr statistics
    peek ptr `returns` statistics

bulkContext :: Property
bulkContext =
  Property.example $ do
    bitmap <- owned64 Bitmap64.create
    memory <- Property.alignedBytes BulkContext.size BulkContext.alignment
    let
      context = BulkContext memory
      values = [1, 2, 3, 2 ^ (40 :: Int), 2 ^ (40 :: Int) + 1, maxBound]
      reset = evalIO $ fillBytes memory 0 BulkContext.size
    reset
    evalIO . forM_ values $ Bitmap64.addBulk bitmap context
    reset
    forM (values <> [4, 2 ^ (40 :: Int) + 2]) (Bitmap64.containsBulk bitmap context) `returns` (map (const True) values <> [False, False])
    reset
    evalIO $ do
      Bitmap64.removeBulk bitmap context 2
      Bitmap64.removeBulk bitmap context 3
    toList64 bitmap >>= (=== [1, 2 ^ (40 :: Int), 2 ^ (40 :: Int) + 1, maxBound])

iterateAll :: Property
iterateAll =
  Property.example $ do
    bitmap <- fromValues64 [9, 1, 2 ^ (40 :: Int), 5]
    (seen, callback) <- collect Bitmap.wrapCallback64 $ const True
    Bitmap64.iterate bitmap callback nullPtr `returns` True
    reverse <$> readIORef seen `returns` [1, 5, 9, 2 ^ (40 :: Int)]

scoped :: Property
scoped =
  Property.example $ do
    Bitmap64.with (withArrayLen [1, 5, 2 ^ (40 :: Int)] $ \count values -> Bitmap64.ofPtr (fromIntegral count) values) elements64 `returns` [1, 5, 2 ^ (40 :: Int)]
    outcome <- evalIO . try $ Bitmap64.with (pure $ Bitmap64 nullPtr) elements64
    outcome === Left (NullHandle "Data.Bitmap.Roaring.FFI.Bitmap64.with")

portableFrozenScope :: Property
portableFrozenScope =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    (buffer, _, written) <- serialized Bitmap64.portableSizeInBytes Bitmap64.portableSerialize bitmap
    Bitmap64.withPortableFrozen buffer written elements64 `returns` Just (Set.toAscList $ Set.fromList values)
    Bitmap64.withPortableFrozen buffer (written - 1) elements64 `returns` Nothing

frozenViewScope :: Property
frozenViewScope =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    _ <- evalIO $ Bitmap64.shrinkToFit bitmap
    (buffer, _, written) <- serialized Bitmap64.frozenSizeInBytes Bitmap64.frozenSerialize bitmap
    Bitmap64.withFrozenView buffer written elements64 `returns` Just (Set.toAscList $ Set.fromList values)
    Bitmap64.withFrozenView (buffer `plusPtr` 1) written elements64 `returns` Nothing

callbackHelpers :: Property
callbackHelpers =
  Property.example $ do
    bitmap <- fromValues64 [1 .. 10]
    visiting (Bitmap64.iterateWith bitmap) (const $ pure True) `returns` (Right True, [1 .. 10])
    visiting (Bitmap64.iterateWith bitmap) (pure . (< 3)) `returns` (Right False, [1, 2, 3])
    visiting (Bitmap64.iterateWith bitmap) stopAtThree `returns` (Left $ userError "stop", [1, 2, 3])
  where
    stopAtThree value = if value == 3 then ioError $ userError "stop" else pure True

selfAliasing :: Property
selfAliasing =
  Property.generated $ do
    values <- forAll genValues64
    let
      expected = Set.toAscList $ Set.fromList values
      applied :: (Bitmap64 -> Bitmap64 -> IO ()) -> Test [Word64]
      applied operation = do
        bitmap <- fromValues64 values
        evalIO $ operation bitmap bitmap
        toList64 bitmap
    applied Bitmap64.orInplace >>= (=== expected)
    applied Bitmap64.xorInplace >>= (=== [])
    applied Bitmap64.andnotInplace >>= (=== [])
    applied Bitmap64.andInplace >>= (=== expected)
    applied Bitmap64.overwrite >>= (=== expected)

scopedBulkContext :: Property
scopedBulkContext =
  Property.example $ do
    bitmap <- owned64 Bitmap64.create
    evalIO . BulkContext.with $ \context -> forM_ values $ Bitmap64.addBulk bitmap context
    BulkContext.with (\context -> forM (values <> [4]) $ Bitmap64.containsBulk bitmap context) `returns` (map (const True) values <> [False])
    toList64 bitmap >>= (=== values)
  where
    values = [1, 2, 3, 2 ^ (40 :: Int), maxBound]