{-# LANGUAGE LambdaCase #-}
module TaPL7 where

import           Data.Function                  ( on )
import           Data.List                      ( unfoldr )

data Term
    = TmVar Int
    | TmAbs Term
    | TmApp Term Term
    deriving Show


tShift :: Int -> Term -> Term
tShift d =
    let walk :: Int -> Term -> Term
        walk c = \case
            TmVar k | k < c     -> TmVar k
                    | otherwise -> TmVar (k + d)
            TmAbs t     -> TmAbs $ walk (succ c) t
            TmApp t1 t2 -> (TmApp `on` walk c) t1 t2
    in  walk 0


tSubst :: Int -> Term -> Term -> Term
tSubst j s = \case
    TmVar k | k == j    -> s
            | otherwise -> TmVar k
    TmAbs t     -> TmAbs $ tSubst (succ j) (tShift 1 s) t
    TmApp t1 t2 -> (TmApp `on` tSubst j s) t1 t2


tSubstTop :: Term -> Term -> Term
tSubstTop v t = tShift (-1) $ tSubst 0 (tShift 1 v) t


tEval :: Term -> Maybe Term
tEval = \case
    TmApp (TmAbs t1) t2    -> Just $ tSubstTop t2 t1
    TmApp t1 t2 | isval t1 -> TmApp t1 <$> tEval t2
    TmApp t1 t2            -> flip TmApp t2 <$> tEval t1
    _                      -> Nothing
  where
    isval :: Term -> Bool
    isval = \case
        TmAbs _ -> True
        _       -> False


tEvalAll :: Term -> [Term]
tEvalAll = unfoldr $ fmap delta . tEval where delta a = (a, a)


printEvals :: Term -> IO ()
printEvals = mapM_ printTerm . tEvalAll


printTerm :: Term -> IO ()
printTerm = putStrLn . showTerm


showTerm :: Term -> String
showTerm = \case
    TmVar n     -> show n
    TmAbs t     -> "(λ." ++ showTerm t ++ ")"
    TmApp t1 t2 -> showTerm t1 ++ " " ++ showTerm t2