Iterator.hs

{-# LANGUAGE OverloadedRecordDot #-}

module Test.Bitmap.Iterator
  ( tests
  , readAll
  , modelRanges
  ) where

import Control.Monad (forM)
import Data.Bitmap.Roaring.FFI.Bitmap (Bitmap)
import Data.Bitmap.Roaring.FFI.Bitmap.Iterator (Iterator (..), Range (..))
import Data.Bitmap.Roaring.FFI.Bitmap.Iterator qualified as Iterator
import Data.Set qualified as Set
import Data.Word (Word32)
import Foreign.Marshal.Array (allocaArray, peekArray)
import Foreign.Ptr (Ptr, nullPtr)
import Foreign.Storable (Storable, peek, poke)
import Hedgehog (Property, evalIO, forAll)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Test.Bitmap (fromValues, genValues)
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
    "Bitmap.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 "skips forward and backward" skips
    , testProperty "heap iterators copy independently" heap
    , testProperty "with and withLast scope an iterator" scoped
    , testProperty "range poke and peek agree" rangeStorable
    ]

placed :: (Bitmap -> Iterator -> IO ()) -> Bitmap -> Test Iterator
placed position bitmap = do
  iterator <- Iterator <$> Property.alignedBytes Iterator.size Iterator.alignment
  evalIO $ position bitmap iterator
  pure iterator

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

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

current :: Iterator -> IO (Maybe Word32)
current iterator = do
  present <- Iterator.hasValue iterator
  if present then Just <$> Iterator.currentValue iterator else pure Nothing

walks :: Property
walks =
  Property.generated $ do
    values <- forAll genValues
    bitmap <- fromValues values
    let expected = Set.toAscList $ Set.fromList values
    forward <- placed Iterator.init bitmap
    walk Iterator.advance forward `returns` expected
    backward <- placed Iterator.initLast bitmap
    walk Iterator.previous backward `returns` reverse expected

chunks :: Property
chunks =
  Property.generated $ do
    values <- forAll genValues
    chunk <- forAll . Gen.int $ Range.linear 1 50
    bitmap <- fromValues values
    let expected = Set.toAscList $ Set.fromList values
    forward <- placed Iterator.init bitmap
    readAll Iterator.read chunk forward `returns` expected
    backward <- placed Iterator.initLast bitmap
    readAll Iterator.readBackward chunk backward `returns` reverse expected

readAll :: (Storable value, Integral count) => (iterator -> Ptr value -> count -> IO count) -> Int -> iterator -> IO [value]
readAll readNext chunk iterator =
  allocaArray chunk $ \buffer -> do
    let loop = do
          count <- fromIntegral <$> readNext iterator buffer (fromIntegral chunk)
          found <- peekArray count buffer
          if count < chunk then pure found else (found <>) <$> loop
    loop

ranges :: Property
ranges =
  Property.generated $ do
    values <- forAll genValues
    bitmap <- fromValues values
    let expected = modelRanges . Set.toAscList $ Set.fromList values
    forward <- placed Iterator.init bitmap
    readRangesWith Iterator.readRanges forward `returns` expected
    backward <- placed Iterator.initLast 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

modelRanges :: (Eq a, Num a) => [a] -> [(a, a)]
modelRanges = foldr extend []
  where
    extend value ((low, high) : rest)
      | value + 1 == low = (value, high) : rest
    extend value grouped = (value, value) : grouped

moves :: Property
moves =
  Property.generated $ do
    values <- forAll genValues
    probes <- forAll . Gen.list (Range.linear 0 50) $ Gen.word32 Range.constantBounded
    bitmap <- fromValues values
    iterator <- Iterator <$> Property.alignedBytes Iterator.size Iterator.alignment
    let
      model = Set.fromList values
      probe value = do
        Iterator.init bitmap iterator
        _ <- Iterator.moveEqualOrLarger iterator value
        current iterator
    forM probes probe `returns` map (`Set.lookupGE` model) probes

skips :: Property
skips =
  Property.generated $ do
    values <- forAll genValues
    count <- forAll . Gen.int $ Range.linear 0 5000
    bitmap <- fromValues values
    let
      model = Set.fromList values
      total = Set.size model
      skippable = fromIntegral $ min count total
    forward <- placed Iterator.init bitmap
    Iterator.skip forward (fromIntegral count) `returns` skippable
    current forward `returns` if count < total then Just (Set.elemAt count model) else Nothing
    backward <- placed Iterator.initLast bitmap
    Iterator.skipBackward backward (fromIntegral count) `returns` skippable
    current backward `returns` if count < total then Just (Set.elemAt (total - 1 - count) model) else Nothing

heap :: Property
heap =
  Property.example $ do
    bitmap <- fromValues [3, 70000, 9]
    original <- created $ Iterator.create bitmap
    Iterator.advance original `returns` True
    copied <- created $ Iterator.copy original
    Iterator.advance original `returns` True
    Iterator.currentValue original `returns` 70000
    walk Iterator.advance copied `returns` [9, 70000]

rangeStorable :: Property
rangeStorable =
  Property.example $ do
    let range = Range 5 99
    ptr <- Property.array 1
    evalIO $ poke ptr range
    peek ptr `returns` range

scoped :: Property
scoped =
  Property.generated $ do
    values <- forAll genValues
    bitmap <- fromValues 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