Fix.hs

{-# LANGUAGE DeriveFunctor, LambdaCase #-}
import Control.Monad.State
import Control.Applicative
import Text.Printf

data TermF a = Var String 
             | App a a 
             | Lam String a 
             deriving(Functor)

newtype ITermF a = ITermF (Int,TermF a) 
                   deriving(Functor)

newtype Fix f = In (f (Fix f))

cata :: Functor f => (f a -> a) -> Fix f -> a
cata f (In t) = f (fmap (cata f) t)

newtype Term = Term (Fix TermF)
newtype ITerm = ITerm (Fix ITermF)

incr :: State Int Int
incr = do
    i <- get
    put $! i + 1
    return i

eval :: Term -> Maybe Term
eval (Term _t) = Term <$> go _t where
    go (In _t) = In <$> case _t of
        Var _ -> empty
        App (In (Lam x t)) t2 -> return $ subst x t2 t
        App t1 t2 -> (App <$> go t1 <*> pure t2) `mplus`
                     (App <$> pure t1 <*> go t2)
        Lam x t -> Lam x <$> go t
    subst x (In tx) (In _t) = case _t of
        Var y | x == y -> tx
        App t1 t2 -> App (In $ subst x (In tx) t1) 
                         (In $ subst x (In tx) t2)
        Lam y t | x /= y -> Lam y (In $ subst x (In tx) t)
        _ -> _t

assignId :: Term -> State Int ITerm
assignId (Term _t) = ITerm <$> cata f _t where
    f (Var s) = c $ pure $ Var s
    f (App t1 t2) = c $ App <$> t1 <*> t2
    f (Lam x t) = c $ Lam x <$> t
    c t = In . ITermF <$> ((,) <$> incr <*> t)

instance Show Term where
    showsPrec i (Term (In _t)) = case _t of
        Var s -> showString s
        App t1 t2 -> 
            let s1 = showsPrec 1 (Term t1)
                s2 = showsPrec 1 (Term t2)
                s  = showString "App" . showChar ' ' 
                                      . s1
                                      . showChar ' '
                                      . s2 in
            showParen (i > 0) $ s
        Lam x t ->
            let s = showString "Lam" . showChar ' '
                                     . showString x
                                     . showChar ' '
                                     . showsPrec 1 (Term t) in
            showParen (i > 0) $ s

instance Show ITerm where
    showsPrec i (ITerm (In (ITermF (n,_t)))) = case _t of
        Var s -> showParen True $ showString s . showString ", "
                                . showString (show n)
        App t1 t2 -> 
            let s1 = showsPrec 1 (ITerm t1)
                s2 = showsPrec 1 (ITerm t2)
                s  = showString "App" . showChar ' ' 
                                      . showString (show n)
                                      . showChar ' '
                                      . s1
                                      . showChar ' '
                                      . s2 in
            showParen (i > 0) $ s
        Lam x t ->
            let s = showString "Lam" . showChar ' '
                                     . showString (show n)
                                     . showChar ' '
                                     . showString x
                                     . showChar ' '
                                     . showsPrec 1 (ITerm t) in
            showParen (i > 0) $ s

var :: String -> Term
var x = Term $ In $ Var x

app :: Term -> Term -> Term
app (Term t1) (Term t2) = Term $ In $ App t1 t2

lam :: String -> Term -> Term
lam x (Term t) = Term $ In $ Lam x t

main :: IO ()
main = do
    let t1 = lam "x" (var "x")
    let t = app (app t1 (var "z")) (app t1 (var "y"))
    printf "t = %s\n" $ show t
    printf "assignId t 0 = %s\n" $ show $ evalState (assignId t) 0
    printf "eval t = %s\n" $ show $ eval t