Iterator.hs

{-# LANGUAGE OverloadedRecordDot #-}

module Test.Bitmap64.Iterator (tests) where

import Control.Monad (forM)
import Data.Bitmap.Roaring.FFI.Bitmap64.Iterator (Iterator (..), Range (..))
import Data.Bitmap.Roaring.FFI.Bitmap64.Iterator qualified as Iterator
import Data.Set qualified as Set
import Data.Word (Word64)
import Foreign.Marshal.Array (allocaArray, peekArray)
import Foreign.Ptr (nullPtr)
import Foreign.Storable (peek, poke)
import Hedgehog (Property, evalIO, forAll)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Test.Bitmap.Iterator (modelRanges, readAll)
import Test.Bitmap64 (fromValues64, genValues64)
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.Iterator"
    [ testProperty "walks forward and backward" walks
    , testProperty "reads values in chunks" chunks
    , testProperty "reads maximal ranges" ranges
    , testProperty "moves to the first value at or above a probe" moves
    , testProperty "copies and reinitialises" copies
    , testProperty "with and withLast scope an iterator" scoped
    , testProperty "range poke and peek agree" rangeStorable
    ]

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

walk :: (Iterator -> IO Bool) -> Iterator -> IO [Word64]
walk step iterator = do
  present <- Iterator.hasValue iterator
  if present then do
    value <- Iterator.value iterator
    _ <- step iterator
    (value :) <$> walk step iterator
  else
    pure []

walks :: Property
walks =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    let expected = Set.toAscList $ Set.fromList values
    forward <- created $ Iterator.create bitmap
    walk Iterator.advance forward `returns` expected
    backward <- created $ Iterator.createLast bitmap
    walk Iterator.previous backward `returns` reverse expected

chunks :: Property
chunks =
  Property.generated $ do
    values <- forAll genValues64
    chunk <- forAll . Gen.int $ Range.linear 1 50
    bitmap <- fromValues64 values
    let expected = Set.toAscList $ Set.fromList values
    forward <- created $ Iterator.create bitmap
    readAll Iterator.read chunk forward `returns` expected
    backward <- created $ Iterator.createLast bitmap
    readAll Iterator.readBackward chunk backward `returns` reverse expected

ranges :: Property
ranges =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    let expected = modelRanges . Set.toAscList $ Set.fromList values
    forward <- created $ Iterator.create bitmap
    readRangesWith Iterator.readRanges forward `returns` expected
    backward <- created $ Iterator.createLast bitmap
    readRangesWith Iterator.readPrevRanges backward `returns` reverse expected
  where
    readRangesWith readNext iterator =
      allocaArray 3 $ \buffer -> do
        let loop = do
              count <- fromIntegral <$> readNext iterator buffer 3
              found <- map (\range -> (range.min, range.max)) <$> peekArray count buffer
              if count < 3 then pure found else (found <>) <$> loop
        loop

moves :: Property
moves =
  Property.generated $ do
    values <- forAll genValues64
    probes <- forAll . Gen.list (Range.linear 0 50) $ Gen.word64 Range.constantBounded
    bitmap <- fromValues64 values
    iterator <- created $ Iterator.create bitmap
    let
      model = Set.fromList values
      probe value = do
        Iterator.reinit bitmap iterator
        present <- Iterator.moveEqualOrLarger iterator value
        if present then Just <$> Iterator.value iterator else pure Nothing
    forM probes probe `returns` map (`Set.lookupGE` model) probes

copies :: Property
copies =
  Property.example $ do
    bitmap <- fromValues64 [3, 2 ^ (40 :: Int), 9]
    original <- created $ Iterator.create bitmap
    Iterator.advance original `returns` True
    copied <- created $ Iterator.copy original
    Iterator.advance original `returns` True
    Iterator.value original `returns` 2 ^ (40 :: Int)
    walk Iterator.advance copied `returns` [9, 2 ^ (40 :: Int)]
    evalIO $ Iterator.reinit bitmap copied
    Iterator.value copied `returns` 3
    evalIO $ Iterator.reinitLast bitmap copied
    walk Iterator.previous copied `returns` [2 ^ (40 :: Int), 9, 3]

rangeStorable :: Property
rangeStorable =
  Property.example $ do
    let range = Range 5 (2 ^ (40 :: Int))
    ptr <- Property.array 1
    evalIO $ poke ptr range
    peek ptr `returns` range

scoped :: Property
scoped =
  Property.generated $ do
    values <- forAll genValues64
    bitmap <- fromValues64 values
    let expected = Set.toAscList $ Set.fromList values
    Iterator.with bitmap (walk Iterator.advance) `returns` expected
    Iterator.withLast bitmap (walk Iterator.previous) `returns` reverse expected