Values.hs

{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE NoFieldSelectors #-}

module Bench.Values
  ( Values (..)
  , dense
  , scattered
  , scattered64
  , withValues
  , prefix
  , forEach
  , countWhere
  ) where

import Data.Bits (shiftL, shiftR)
import Data.Word (Word32, Word64)
import Foreign.Marshal.Alloc (free)
import Foreign.Marshal.Array (mallocArray, pokeArray)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable, peekElemOff)
import Test.Tasty (withResource)
import Test.Tasty.Bench (Benchmark)

data Values a = Values
  { pointer :: Ptr a
  , count :: Int
  }

dense :: (Num a, Enum a) => Int -> [a]
dense size = take size $ enumFrom 0

scattered :: Word32 -> Int -> [Word32]
scattered seed size = take size . drop 1 $ iterate step seed
  where
    step value = value * 1664525 + 1013904223

scattered64 :: Word32 -> Int -> [Word64]
scattered64 seed = zipWith place (cycle [0 .. 3]) . scattered seed
  where
    place segment value = segment `shiftL` 40 + fromIntegral value `shiftR` 2

withValues :: (Storable a) => [a] -> (IO (Values a) -> Benchmark) -> Benchmark
withValues values = withResource acquire release
  where
    acquire = do
      let count = length values
      pointer <- mallocArray count
      pokeArray pointer values
      pure Values{pointer, count}
    release prepared = free prepared.pointer

prefix :: Int -> Values a -> Values a
prefix limit values = Values{pointer = values.pointer, count = min limit values.count}

forEach :: (Storable a) => Values a -> (a -> IO ()) -> IO ()
forEach values action = go 0
  where
    go index
      | index == values.count = pure ()
      | otherwise = do
          peekElemOff values.pointer index >>= action
          go $ index + 1
{-# INLINE forEach #-}

countWhere :: (Storable a) => Values a -> (a -> IO Bool) -> IO Int
countWhere values query = go 0 0
  where
    go !hits index
      | index == values.count = pure hits
      | otherwise = do
          hit <- peekElemOff values.pointer index >>= query
          go (if hit then hits + 1 else hits) $ index + 1
{-# INLINE countWhere #-}