Simply Typed Lambda Calculus.hs

module Lambda where 

import Evaluator
import Control.Monad.Writer
import Control.Monad.Reader
import qualified Data.Map as M

-- | A simply typed lambda calculus, 
-- extended with integers and let bindings.
-- No alpha reduction.
--
data Expr = 
    Value Int
  | Variable String 
  | Lambda String Type Expr -- (only) lambdas need to have a type annotation
  | App Expr Expr
  | Let String Expr Expr -- let s = e1 in e2
  | Add Expr Expr
  | Sub Expr Expr
  | Mul Expr Expr
  | Div Expr Expr
  deriving (Show, Eq)

data Type = TInt | TFun Type Type deriving Eq

instance Show Type where
  show TInt = "int"
  show (TFun t1 t2) = show t1 ++ " -> " ++ show t2

int = TInt
t1 ~> t2 = TFun t1 t2

-- * Type checking
--
typeOf :: Expr -> Eval (Env Type) [String] Type
typeOf (Value i) = succeeds TInt
typeOf (Variable s) = envLookup s
typeOf (Lambda v t e) = local (M.insert v t) (typeOf e) >>= succeeds . TFun t
typeOf (Let s e1 e2) = typeOf e1 >>= \ty -> local (M.insert s ty) (typeOf e2)
typeOf (Add e1 e2) = intType e1 e2
typeOf (Sub e1 e2) = intType e1 e2
typeOf (Mul e1 e2) = intType e1 e2
typeOf (Div e1 e2) = intType e1 e2
typeOf (App e1 e2) = typeOf e1 >>= \t1 -> typeOf e2 >>= flip checkFunc t1

checkFunc ty (TFun dom ran) | dom == ty = succeeds ty
                            | otherwise = failsWith ("Expecting " ++ show (dom ~> ran) 
                                                 ++ ", but got: " ++ show (ty ~> ran))
checkFunc _ ty = failsWith ("Expecting function, got: " ++ show ty)

intType e1 e2 = do t1 <- typeOf e1
                   t2 <- typeOf e2
                   if t1 == int && t2 == int
                     then succeeds TInt
                     else failsWith ("Expecting int -> int -> int, but got: " ++ 
                                     show (t1 ~> t2 ~> int))
                   

-- * Evaluating
-- (Pretty much the same as expr.hs
--
eval :: Expr -> (Maybe Expr, [String])
eval expr = case evalEval (typeOf expr) (M.fromList []) of
              (Nothing, msg) -> (Nothing, msg)
              (Just t, msg) -> evalEval (tell msg >> evalExpr expr) (M.fromList [])
  
instance Num Expr where
  fromInteger = Value . fromInteger
  (+) = Add
  (-) = Sub
  (*) = Mul
  abs = undefined
  signum = undefined

evalExpr :: Expr -> Eval (Env Expr) [String] Expr
evalExpr (Value i) = succeeds (Value i)
evalExpr (Variable s) = envLookup s
evalExpr l@(Lambda v _ e) = succeeds l
evalExpr (App e1 e2) = do func <- evalExpr e1
                          case func of 
                            (Lambda var _ e) -> updatedEnv var e2 e
evalExpr (Let v e1 e2) = updatedEnv v e1 e2
evalExpr (Add e1 e2) = binOp e1 e2 (+)
evalExpr (Sub e1 e2) = binOp e1 e2 (-)
evalExpr (Mul e1 e2) = binOp e1 e2 (*)
evalExpr (Div e1 e2) = do 
    v1 <- getValue $ evalExpr e1
    v2 <- getValue $ evalExpr e2
    case v2 of 
      0 -> failsWith "Error: division by zero"
      v -> succeeds (Value $ v1 `div` v2)

binOp e1 e2 op = do
  v1 <- getValue $ evalExpr e1
  v2 <- getValue $ evalExpr e2
  succeeds (Value $ v1 `op` v2)

getValue e = do v <- e
                case v of 
                  (Value i) -> return i

updatedEnv var e1 e2 = do val <- evalExpr e1
                          local (M.insert var val) (evalExpr e2)