import Control.Applicative
import Control.Monad.Reader
import Data.Traversable
import Data.Foldable
import Data.Monoid
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map (Map)
import qualified Data.Map as Map

-- We split our datatype into the top-level structure (Expr') and the recursive
-- knot-tying (Expr)
type Name = String
data Expr' e = Var Name | App e e | Lam Name e
newtype Expr = Expr { unExpr :: Expr' Expr }

-- Expr' has all kinds of nice properties.

instance Functor Expr' where
  fmap f (Var x)     = Var x
  fmap f (App e1 e2) = App (f e1) (f e2)
  fmap f (Lam x e)   = Lam x (f e)

instance Foldable Expr' where
  foldMap f (Var _)     = mempty
  foldMap f (App e1 e2) = f e1 `mappend` f e2
  foldMap f (Lam _ e)   = f e

instance Traversable Expr' where
  traverse f (Var x)     = pure (Var x)
  traverse f (App e1 e2) = App <$> f e1 <*> f e2
  traverse f (Lam x e)   = Lam x <$> f e

-- Once we get the instances out of the way we can define freeVars quite
-- elegantly.

freeVars :: Expr -> Set Name
freeVars (Expr e) = case e of
  Var x   -> Set.singleton x
  Lam x e -> Set.delete x (freeVars e)
  _       -> foldMap freeVars e

-- Why isn't this instance in the libraries??
instance Applicative (Reader e) where
  pure  = return
  (<*>) = ap

-- Another more interesting example.

alphaRename :: Expr -> Expr
alphaRename e = runReader (alpha e) (names, Map.empty)
  where
    names = [ s ++ [c] | s <- "":names, c <- ['a'..'z'] ]

    alpha (Expr e) = Expr <$> case e of
      Var x   -> Var <$> rename x
      Lam x e -> fresh x $ \y -> Lam y <$> alpha e
      _       -> traverse alpha e

    rename x = do
      (_, ren) <- ask
      return $ maybe x id $ Map.lookup x ren

    fresh x f = do
      (y:names, ren) <- ask
      local (const (names, Map.insert x y ren)) (f y)

-- Examples

lam x e   = Expr (Lam x e)
app e1 e2 = Expr (App e1 e2)
var x     = Expr (Var x)

e1 = lam "x" $ var "x"
e2 = lam "unused" $ var "C"
e3 = lam "x" $ lam "y" $ lam "z" $ app (var "x") (var "z") `app` app (var "y") (var "z")

instance Show Expr where
  showsPrec p (Expr e) = showsPrec p e

instance Show e => Show (Expr' e) where
  showsPrec p (Var x) = showString x
  showsPrec p (App e1 e2) = showParen (p>1) $
    showsPrec 1 e1 . showString " " . showsPrec 2 e2
  showsPrec p (Lam x e) = showParen (p>0) $
    showString ("\\" ++ x ++ " -> ") . shows e