{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}

module Memoize
  ( memoizeBenchmark
  ) where

import Data.Bits
import Data.Chimera
import Data.Foldable
import Data.Function
import qualified Data.Vector.Unboxed as U
import Test.Tasty.Bench

memoizeBenchmark :: Benchmark
memoizeBenchmark = bgroup "memoize"
  [ bgroup "memoizeFix" $ memoizeFixBenchmark memoizeFix
  , bgroup "memoizeFix unboxed" $ memoizeFixBenchmark (index . tabulateFix @U.Vector)
  , bgroup "fix memoize" $ memoizeFixBenchmark (fix . (memoize .))
  ]

memoizeFixBenchmark :: (forall a. U.Unbox a => ((Word -> a) -> Word -> a) -> Word -> a) -> [Benchmark]
memoizeFixBenchmark fixer =
  [ bench "isOdd" $ nf (\f -> let isOdd = fixer f in
      foldl' (\acc n -> xor acc (isOdd n)) False [0..10000]) isOddF
  , bench "isPrime" $ nf (\f -> let isPrime = fixer f in
      foldl' (\acc n -> xor acc (isPrime n)) False [0..10000]) isPrimeF
  , bench "fibo" $ nf (\f -> let fibo = fixer f in
      foldl' (\acc n -> acc + fibo n) 0 [0..10000]) fiboF
  , bench "collatz" $ nf (\f -> let collatz = fixer f in
      foldl' (\acc n -> acc + collatz n) 0 [0..1000]) collatzF
  ]

isOddF :: (Word -> Bool) -> Word -> Bool
isOddF f n = n /= 0 && not (f (n - 1))

isPrimeF :: (Word -> Bool) -> Word -> Bool
isPrimeF f n = n > 1 && and [ n `rem` d /= 0 | d <- [2 .. floor (sqrt (fromIntegral n :: Double))], f d]

fiboF :: (Word -> Word) -> Word -> Word
fiboF f n = if n < 2 then fromIntegral n else f (n - 1) + f (n - 2)

collatzF :: (Word -> Word) -> Word -> Word
collatzF f n = if n <= 1 then 0 else 1 + f (if even n then n `quot` 2 else 3 * n + 1)