Simply Typed Lambda Calculus with Type Inference.hs

-- | A simply typed lambda calculus with type reconstruction.
-- extended with integers and let bindings.
-- No alpha reduction.
--

module Lambda where 

import Evaluator
import Control.Monad.Writer
import Control.Monad.Reader
import qualified Data.Map as M
import Data.Char
import Control.Monad ((>=>))


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)

instance Num Expr where
  fromInteger = Value . fromInteger
  (+) = Add
  (-) = Sub
  (*) = Mul
  abs = undefined
  signum = undefined

data Type = TInt 
          | TVar String    -- Type variables
          | TFun Type Type 
          deriving Eq

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

a = TVar "a"
b = TVar "b"
c = TVar "c"
d = TVar "d"
x = Variable "x"
y = Variable "y"
z = Variable "z"
int = TInt
t1 ~> t2 = TFun t1 t2

infixr 9 ~>


type Substitution = (Type, Type)   -- [t1 -> t2]
type Substitutions = [Substitution]
type Constraint = (Type, Type)     -- t1 = t2
type Constraints = [Constraint]

-- * Constraints
--
-- Hindley-Milner type-inference works as follows:
--
--   1. Take untyped lambda calculus and introduce
--      a new fresh variable for every lambda abstraction.
--   2. Walk the AST as if we were type checking, but 
--      instead of type checking, we generate constraints
--      (and still return types as well)
--      For example: when we have the application:
--       
--        t1:X t2:Y
--
--      We still return the codomain of t1 as the type 
--      of the whole term, but instead of applying the 
--      the type Y to X, we generate a new free variable 
--      Z and add a type constraint {X = Y -> Z}.
--
--      We create a constraint for every function application:
--      t1 t2 generates a constraint 
--        typoOf(t1) = typeOf (t2) -> X
--      where X is a free variable.
--
--      For binary functions `t1 + t2`, `t1 - t2`, etc. which takes two terms
--      of type int, we generate two constraints:
--        int = typeOf(t1)
--        int = typeOf(t2)
--
--   3. Use unification to convert the constraints 
--      to a number of type substitions.
--
--   4. Use the substitutions to reconstruct the type 
--      returned in step 2.
--
-- Most of the complexity in @constraints@ comes from 
-- the need for fresh variable generation. Ideally 
-- this would be hidden in the Evaluator monad, 
-- which is left for future work at the moment, 
-- but should be feasible.
--
constraints :: Expr -> Eval (Int, Env Type) [String] (Type, Int, Constraints)
constraints (Value _) = ask >>= \(ix,_) -> succeeds (int, ix, [])
constraints (Variable v) = do 
  ty <- envLookupWith (\v e -> M.lookup v (snd e)) v
  ask >>= \(ix,_) -> succeeds (ty, ix, [])
constraints (Lambda v t e) = do 
  (t', ix, c) <- local (\(f, e) -> (f, M.insert v t e)) (constraints e)
  succeeds (t ~> t', ix, c) 
constraints (App e1 e2) = do 
  (t1, ix1, c1) <- local (\(f,e) -> (f + 1, e)) $ constraints e1
  (t2, ix2, c2) <- local (\(_,e) -> (ix1  , e)) $ constraints e2
  let v = TVar (toVar ix2)
  succeeds (v, ix2, c1 ++ c2 ++ [(t2 ~> v, t1)]) 
constraints (Let s e1 e2) = do 
  (t1, ix1, c1) <- constraints e1
  (t2, ix2, c2) <- local (\(f,e)-> (ix1, M.insert s t1 e)) (constraints e2)
  succeeds (t2, ix2, c1 ++ c2)
constraints (Add e1 e2) = binOp e1 e2 
constraints (Sub e1 e2) = binOp e1 e2
constraints (Mul e1 e2) = binOp e1 e2
constraints (Div e1 e2) = binOp e1 e2

binOp e1 e2 = do 
  (t1, ix1, c1) <- constraints e1
  (t2, ix2, c2) <- local (\(f,e) -> (ix1, e)) $ constraints e2
  succeeds (int, ix2, c1 ++ c2 ++ [(int, t1), (int, t2)])

reconstruct :: (Type, Int, Constraints) -> Eval (Int, Env Type) [String] Type
reconstruct (ty, _, cs) = unify cs >>= succeeds . foldr substitute ty

reconstruct' = eval' (constraints >=> reconstruct)

-- * Unification
-- Convert constraints into substitutions

-- | Tries to unify the constraints; producing 
-- a set of substitutions.
--
unify :: Constraints -> Eval env [String] Substitutions
unify [] = succeeds []
unify ((c1, c2):cs) = if c1 == c2 then unify cs else unify' c1 c2 cs

unify' :: Type -> Type -> Constraints -> Eval env [String] Substitutions
unify' v@(TVar _) c cs = occursCheck v c cs
unify' c v@(TVar _) cs = occursCheck v c cs
unify' (TFun d1 r1) (TFun d2 r2) cs = unify (cs ++ [(d1, d2), (r1, r2)])
unify' v1 v2 _ = failsWith $ "Failed to unify constraints. " ++ show v1 ++ " = " ++ show v2


-- | Occurs check: 
-- In the constraint X = t2, check whether X is 
-- not in the free variables of t2.
--
occursCheck :: Type -> Type -> Constraints -> Eval env [String] Substitutions
occursCheck v@(TVar v') c cs = 
  if elem v' (freeVar c)
    then failsWith ("Failed occurs check: " ++ v' ++ " = " ++ show c)
    else unify (substituteC (v, c) cs) >>= \ss -> succeeds ((v, c):ss)

freeVar :: Type -> [String]
freeVar TInt = []
freeVar (TFun t1 t2) = freeVar t1 ++ freeVar t2
freeVar (TVar a) = [a]

-- * Type Substitutions
-- Substitute type variables.
--
substituteC :: (Type, Type) -> Constraints -> Constraints 
substituteC _ [] = []
substituteC s ((c1, c2):cs) = (substitute s c1, substitute s c2) : substituteC s cs

substitute :: (Type, Type) -> Type -> Type
substitute s TInt = TInt
substitute s (TFun d r) = substitute s d ~> substitute s r
substitute (v, t) v' = if v == v' then t else v'

toVar n | n < 20 = [chr (n + ord 'e')]
        | otherwise = [chr (mod n 20 + ord 'e')] ++ replicate (n `div` 20) '\''


eval' e = flip evalEval (0, (M.fromList [])) . e
eval expr = case reconstruct' expr of
              (Nothing, msg) -> (Nothing, msg)
              (Just t, msg) -> evalEval (tell msg >> evalExpr expr) (M.fromList [])


-- * Evaluation
-- Again, nothing changes here.
--
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)