{-# LANGUAGE FlexibleContexts, GADTs, GeneralizedNewtypeDeriving, PatternGuards, ScopedTypeVariables #-}

-- | Lambda expressions using first-order syntax

module Lambda where



import Control.Monad.Reader
import Data.Array
import Data.Dynamic
import Data.Tree hiding (Node)
import qualified Data.Tree as Tree



--------------------------------------------------------------------------------
-- * Types
--------------------------------------------------------------------------------

-- | Lambda expressions parameterized by an arbitrary expression type. These
-- external expressions can be turned into lambda expressions using the 'Inject'
-- constructor.
data Lambda expr a
  where
    Variable :: Typeable a => VarId -> Lambda expr a

    Lambda :: (Typeable a, Typeable b) =>
        VarId -> Lambda expr b -> Lambda expr (a -> b)

    (:$:) :: Typeable a =>
        Lambda expr (a -> b) -> Lambda expr a -> Lambda expr b

    Inject :: expr a -> Lambda expr a



-- | Variable identifier
newtype VarId = VarId { varInteger :: Integer }
  deriving (Eq, Ord, Num, Enum, Ix)

instance Show VarId
  where
    show (VarId i) = show i



--------------------------------------------------------------------------------
-- * General operations on expressions
--------------------------------------------------------------------------------

-- | Equality for expressions. The difference between 'Eq' and 'ExprEq' is that
-- 'ExprEq' allows comparison of expressions with different value types. It is
-- assumed that when the types differ, the expressions also differ. The reason
-- for allowing comparison of different types is that this is convenient when
-- the types are existentially quantified.
class ExprEq expr
  where
    exprEq :: expr a -> expr b -> Bool

-- | Evaluation of expressions
class Eval expr
  where
    eval :: expr a -> a

class ExprShow expr
  where
    exprShow :: expr a -> String

printExpr :: ExprShow expr => expr a -> IO ()
printExpr = putStrLn . exprShow



--------------------------------------------------------------------------------
-- * Alpha-equivalence on 'Lambda'
--------------------------------------------------------------------------------

-- | Alpha-equivalence on 'Lambda' expressions. Free variables are taken to be
-- equvalent if they have the same name.
eqLambda :: ExprEq expr =>
    Lambda expr a -> Lambda expr b -> Reader [(VarId,VarId)] Bool
eqLambda (Variable i1) (Variable i2) = do
    env <- ask
    case lookup i1 env of
      Nothing  -> return (i1==i2)   -- Free variables
      Just i2' -> return (i2==i2')
eqLambda (Lambda i1 a1) (Lambda i2 a2) = local ((i1,i2):) $ eqLambda a1 a2
eqLambda (f1 :$: a1) (f2 :$: a2) = do
    e <- eqLambda f1 f2
    if e then eqLambda a1 a2 else return False
eqLambda (Inject a) (Inject b) = return (exprEq a b)
eqLambda _ _                   = return False

-- | Alpha-equivalence on 'Lambda' expressions. Free variables are taken to be
-- equvalent if they have the same name.
instance ExprEq expr => ExprEq (Lambda expr)
  where
    exprEq expr1 expr2 = flip runReader [] $ eqLambda expr1 expr2

-- | Alpha-equivalence on 'Lambda' expressions. Free variables are taken to be
-- equvalent if they have the same name.
instance ExprEq expr => Eq (Lambda expr a)
  where
    (==) = exprEq



--------------------------------------------------------------------------------
-- * Showing 'Lambda'
--------------------------------------------------------------------------------

-- | Parser for infix operators of the form @"(op)"@
viewInfix :: String -> Maybe String
viewInfix ('(':op)
    | (')':op') <- reverse op = Just op'
viewInfix _ = Nothing

showVar :: VarId -> String
showVar v = "var" ++ show v

-- | Shows a partially applied expression
exprShowApp :: ExprShow expr
    => [String]       -- ^ Missing arguments
    -> Lambda expr a  -- ^ Partially applied expression
    -> String
exprShowApp args (f :$: a) = exprShowApp (exprShowLam a : args) f
exprShowApp args f = case (viewInfix fStr, args) of
    (Just op, [a,b]) -> "(" ++ unwords [a,op,b] ++ ")"
    _                -> "(" ++ unwords (fStr : args) ++ ")"
  where
    fStr = exprShowLam f

exprShowLam :: ExprShow expr => Lambda expr a -> String
exprShowLam (Variable v) = showVar v
exprShowLam (Lambda v a) = "(\\" ++ showVar v ++ " -> "  ++ exprShowLam a ++ ")"
exprShowLam (f :$: a)    = exprShowApp [exprShowLam a] f
exprShowLam (Inject a)   = exprShow a

instance ExprShow expr => ExprShow (Lambda expr)
  where
    exprShow = exprShowLam

instance ExprShow (Lambda expr) => Show (Lambda expr a)
  where
    show = exprShow

-- | Converts a partially applied expression to a tree
lamToTreeApp :: ExprShow expr
    => Forest String  -- ^ Missing arguments
    -> Lambda expr a  -- ^ Partially applied expression
    -> Tree String
lamToTreeApp args (f :$: a) = lamToTreeApp (lamToTree a : args) f
lamToTreeApp args f         = Tree.Node "Apply" (lamToTree f : args)

-- | Converts a lambda expression to a tree
lamToTree :: ExprShow expr => Lambda expr a -> Tree String
lamToTree (Variable i) = Tree.Node ("Variable " ++ show i) []
lamToTree (Lambda i a) = Tree.Node ("Lambda " ++ show i) [lamToTree a]
lamToTree (f :$: a)    = lamToTreeApp [lamToTree a] f
lamToTree (Inject a)   = Tree.Node (exprShow a) []

-- | Show a lambda expression as a tree
showLamTree :: ExprShow expr => Lambda expr a -> String
showLamTree = drawTree . lamToTree

-- | Print a lambda expression as a tree
drawLambda :: ExprShow expr => Lambda expr a -> IO ()
drawLambda = putStrLn . showLamTree



--------------------------------------------------------------------------------
-- * Evaluating 'Lambda'
--------------------------------------------------------------------------------

evalLambda
    :: (Eval expr, MonadReader [(VarId,Dynamic)] m)
    => Lambda expr a -> m a
evalLambda (Variable v) = do
    env <- ask
    case lookup v env of
      Nothing -> return $ error "eval: evaluating free variable"
      Just a  -> case fromDynamic a of
        Just a -> return a
        _      -> return $ error "eval: internal type error"
evalLambda (Lambda v f) = do
    env <- ask
    return $ \a -> flip runReader ((v,toDyn a):env) $ evalLambda f
evalLambda (f :$: a)  = liftM2 ($) (evalLambda f) (evalLambda a)
evalLambda (Inject a) = return (eval a)

instance Eval expr => Eval (Lambda expr)
  where
    eval a = flip runReader [] $ evalLambda a