{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TemplateHaskell #-}

module BoundRecursionSchemes where

import Bound
import Data.Functor.Foldable
import Data.Functor.Foldable.TH
import Control.Monad (ap)
import Data.Functor.Classes
import Text.Show.Deriving (deriveShow1)
import Data.Eq.Deriving (deriveEq1)

data Exp a
  = V a
  | Lam (Scope () Exp a)
  | Exp a :@ Exp a

  | Int Int
  | Plus (Exp a) (Exp a)
  deriving (Functor, Foldable, Traversable)

makeBaseFunctor ''Exp

instance Eq1   Exp where
  liftEq eq (V a)        (V b)      = eq a b
  liftEq eq (a :@ a')    (b :@ b')  = liftEq eq a b && liftEq eq a' b'
  liftEq eq (Lam e1)  (Lam e2)      = liftEq eq e1 e2 -- n == n' && liftEq eq p p' && liftEq eq e e'
  liftEq _  _            _          = False

deriveShow1 ''Exp
deriving instance Show a => Show (Exp a) -- note: must be after the TH, so uses standalone deriving

instance Applicative Exp where
  pure = V
  (<*>) = ap

instance Monad Exp where
  return = pure
  V x >>= s = s x
  Lam x >>= s = Lam (x >>>= s)
  g :@ x >>= s = (g >>= s) :@ (x >>= s)

  Int x >>= _ = Int x
  Plus e1 e2 >>= s = Plus (e1 >>= s) (e2 >>= s)

nf :: Exp a -> Exp a
nf e@V{}   = e
nf (Lam b) = Lam $ toScope $ nf $ fromScope b
nf (f :@ a) = case whnf f of
  Lam b -> nf (instantiate1 a b)
  f' -> nf f' :@ nf a
nf x@(Int _) = x
nf (Plus e1 e2) = case (nf e1, nf e2) of
  (Int x, Int y) -> Int $ x + y
  (x, y) -> Plus x y

whnf :: Exp a -> Exp a
whnf e@V{}   = e
whnf e@Lam{} = e
whnf (f :@ a) = case whnf f of
  Lam b -> whnf (instantiate1 a b)
  f'    -> f' :@ a
whnf x@(Int _) = x
whnf (Plus e1 e2) = case (whnf e1, whnf e2) of
  (Int x, Int y) -> Int $ x + y
  (x, y) -> Plus x y


--------------------------------------------------
-- * Test

lam :: Eq a => a -> Exp a -> Exp a
lam x f = Lam $ abstract1 x f

iso :: Base (Exp a) (Exp a) -> Exp a
iso (VF x) = V x
iso (LamF f) = Lam f
iso (f :@$ x) = f :@ x
iso (IntF x) = Int x
iso (PlusF e1 e2) = Plus e1 e2

-- | Evaluates to 42
e0 :: Exp ()
e0 = lam () (Plus (V ()) (Int 3))
     :@ Int 39

-- | Use cata to replace `Int 39`
e1 :: Exp ()
e1 = cata go e0 where
  go :: Base (Exp a) (Exp a) -> Exp a
  go (IntF 39) = Int 97
  go x = iso x

fourtyTwo = nf e0
oneHundred = nf e1