1{-# LANGUAGE BangPatterns #-}
2module Data.IntPSQ.Benchmark
3    ( benchmark
4    ) where
5
6import           Data.List (foldl')
7import qualified Data.IntPSQ as IntPSQ
8import           Criterion.Main
9import           Prelude hiding (lookup)
10import           BenchmarkTypes
11
12benchmark :: String -> [BElem] -> BenchmarkSet
13benchmark name elems = BenchmarkSet
14    { bGroupName        = name
15    , bMinView          = whnf bench_minView              initialPSQ
16    , bLookup           = whnf (bench_lookup keys)        initialPSQ
17    , bInsertEmpty      = nf   (bench_insert firstElems)  IntPSQ.empty
18    , bInsertNew        = nf   (bench_insert secondElems) initialPSQ
19    , bInsertDuplicates = nf   (bench_insert firstElems)  initialPSQ
20    , bDelete           = nf   (bench_delete firstKeys)   initialPSQ
21    }
22  where
23    (firstElems, secondElems) = splitAt (numElems `div` 2) elems
24    numElems  = length elems
25    keys      = map (\(x, _, _) -> x) elems
26    firstKeys = map (\(x, _, _) -> x) firstElems
27
28    initialPSQ = IntPSQ.fromList firstElems :: IntPSQ.IntPSQ Int ()
29
30-- Get the sum of all priorities by getting all elements using 'lookup'
31bench_lookup :: [Int] -> IntPSQ.IntPSQ Int () -> Int
32bench_lookup xs m = foldl' (\n k -> maybe n fst (IntPSQ.lookup k m)) 0 xs
33
34-- Insert a list of elements one-by-one into a PSQ
35bench_insert
36    :: [(Int, Int, ())] -> IntPSQ.IntPSQ Int () -> IntPSQ.IntPSQ Int ()
37bench_insert xs m0 = foldl' (\m (k, p, v) -> IntPSQ.insert k p v m) m0 xs
38
39-- Get the sum of all priorities by sequentially popping all elements using
40-- 'minView'
41bench_minView :: IntPSQ.IntPSQ Int () -> Int
42bench_minView = go 0
43  where
44    go !n t = case IntPSQ.minView t of
45      Nothing            -> n
46      Just (k, p, _, t') -> go (n + k + p) t'
47
48-- Empty a queue by sequentially removing all elements
49bench_delete :: [Int] -> IntPSQ.IntPSQ Int () -> IntPSQ.IntPSQ Int ()
50bench_delete keys t0 = foldl' (\t k -> IntPSQ.delete k t) t0 keys
51