-- | Untyped lambda calc for benchmarking
{-# LANGUAGE DeriveGeneric, DeriveDataTypeable, CPP #-}
module BenchLam where

import Control.Applicative
import Control.Monad (replicateM)
import Data.List (foldl')

import GHC.Generics (Generic)
import Data.Typeable (Typeable)

import Control.DeepSeq (NFData(..), deepseq)
#if MIN_VERSION_deepseq (1,4,0)
#else
import Control.DeepSeq.Generics (genericRnf)
#endif
import Criterion (Benchmark, env, bench, nf)

import Unbound.Generics.LocallyNameless

type Var = Name Term

data Term =
  V !Var
  | App !Term !Term
  | Lam !(Bind Var Term)
    deriving (Show, Generic, Typeable)

instance Alpha Term
#if MIN_VERSION_deepseq (1,4,0)
instance NFData Term
#else
instance NFData Term where rnf = genericRnf
#endif


-- | lambda abstract over all the given vars
lams :: [Var] -> Term -> Term
lams [] = id
lams (v:vs) = Lam . bind v . lams vs

-- | apply the given term to the given terms
apps :: Term -> [Term] -> Term
apps = foldl' App

-- eta-expand a term a given number of times
etaN :: Fresh m => Term -> Int -> m Term
etaN m n = do
  vs <- replicateM n (fresh $ s2n "v")
  let ms = map V vs
  return (lams vs $ apps m ms)

-- | While the head is a lambda, descend under it and then do something;
-- then close all the lambdas back up.
workHeadUnderLams :: Fresh m => (Term -> m Term) -> (Term -> m Term)
workHeadUnderLams comp = go
  where
    go m =
      case m of
      Lam bnd -> do
        (x, m') <- unbind bnd
        m'' <- go m'
        return $ Lam $ bind x m''
      _ -> comp m

freshNeutralTermHead :: (Applicative m, Fresh m) => Term -> m Term
freshNeutralTermHead (App m n) = App <$> freshNeutralTermHead m <*> pure n
freshNeutralTermHead (V v) = V <$> fresh v
freshNeutralTermHead lam@(Lam {}) = return lam

-- | A benchmark that creates an eta expansion of the term "x" of a given size
-- and then freshens the "x" by traversing down below all the lambdas.
--
-- Every time we go under a lambda, we freshen the body, so to go
-- under N lambdas, we do O(N²) work.
freshenEtaTermBench :: Int -> Benchmark
freshenEtaTermBench n =
  let name = "freshen eta term of size " ++ show n
  in name `deepseq` env setup $ \m ->
  bench name $ nf (runFreshM . workHeadUnderLams freshNeutralTermHead) m
  where
    setup = return $ runFreshM $ etaN (V $ s2n "x") n