{-# LANGUAGE DeriveAnyClass, FlexibleInstances, StandaloneDeriving #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Main where

import Alignment
import Criterion.Main
import Data.Bifunctor.Join
import Data.String
import Data.These
import Prologue
import Test.QuickCheck

main :: IO ()
main = do
  benchmarks <- sequenceA [ generativeBenchmark "numberedRows" 10 length (nf (numberedRows :: [Join These ()] -> [Join These (Int, ())])) ]
  defaultMain benchmarks

-- | Defines a named group of n benchmarks over inputs generated by an `Arbitrary` instance.
-- |
-- | The inputs’ sizes at each iteration are measured by a metric function, which gives the name of the benchmark. This makes it convenient to correlate a benchmark of some function over lists with e.g. input `length`.
generativeBenchmark :: (Arbitrary a, Show m) => String -> Int -> (a -> m) -> (a -> Benchmarkable) -> IO Benchmark
generativeBenchmark name n metric benchmark = do
  benchmarks <- traverse measure (replicate n defaultSize)
  pure $! bgroup name benchmarks
  where measure n = do
          input <- generate (resize n arbitrary)
          let measurement = metric input
          pure $! bench (show measurement) (benchmark input)
        defaultSize = 100


-- Instances

deriving instance (NFData a, NFData b) => NFData (These a b)
deriving instance NFData a => NFData (Join These a)

instance (Arbitrary a, Arbitrary b) => Arbitrary (These a b) where
  arbitrary = oneof [ This <$> arbitrary
                    , That <$> arbitrary
                    , These <$> arbitrary <*> arbitrary ]
  shrink = these (fmap This . shrink) (fmap That . shrink) (\ a b -> (This <$> shrink a) ++ (That <$> shrink b) ++ (These <$> shrink a <*> shrink b))

instance Arbitrary a => Arbitrary (Join These a) where
  arbitrary = Join <$> arbitrary
  shrink (Join a) = Join <$> shrink a