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